# AYesha — DDS Enterprise HR Chatbot (Hugging Face Space / Render)
# !pip install -U "gradio>=6.28" spaces pinecone llama-index llama-index-vector-stores-pinecone llama-index-readers-file pypdf llama-index-llms-openai llama-index-embeddings-openai

# --- Imports ---
import logging
import os
import sys

import gradio as gr
import spaces

from pinecone import Pinecone, ServerlessSpec
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    Settings,
    PromptTemplate,
)
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.readers.file import PDFReader
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# --- Logging ---
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Keys (set these in Space Settings -> Variables and secrets) ---
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")

if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY is not set. Add it under Settings -> Variables and secrets.")
if not PINECONE_API_KEY:
    raise RuntimeError("PINECONE_API_KEY is not set. Add it under Settings -> Variables and secrets.")

# --- LlamaIndex global settings ---
Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0.2)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")
Settings.chunk_size = 600
Settings.chunk_overlap = 200

SYSTEM_PROMPT = """You are AYesha, the Decoding Data Science (DDS) Enterprise HR Chatbot. Answer questions exclusively using the attached DDS HR Handbook. Base all responses on the most up-to-date information available in the handbook. Only respond to queries directly related to DDS HR policies as outlined in the handbook.
Important instructions:
- Only answer questions directly supported by the latest DDS HR handbook.
- Decline politely and redirect to the provided email address for any questions outside scope or for confidential information.
- Always reason before concluding. Only present the answer after checking scope and source.
Remember: As AYesha, the DDS HR Enterprise Chatbot, you must never provide information outside authorized HR handbook content and always respond respectfully according to these constraints.
"""

# LlamaIndex's query engine doesn't accept a `system_prompt=` kwarg — that param only
# exists on as_chat_engine(). The equivalent for a query engine is a custom QA prompt
# template, which is what actually enforces AYesha's scope/persona.
QA_TEMPLATE = PromptTemplate(
    SYSTEM_PROMPT
    + "\n\nContext information is below.\n---------------------\n{context_str}\n---------------------\n"
    "Given the context information and not prior knowledge, answer the query.\n"
    "Query: {query_str}\nAnswer: "
)

# --- Initialize Pinecone ---
pc = Pinecone(api_key=PINECONE_API_KEY)
INDEX_NAME = "quickstart"
DIMENSION = 1536

existing_indexes = [idx["name"] for idx in pc.list_indexes()]

if INDEX_NAME not in existing_indexes:
    # Only create the index the first time. Deleting + recreating on every
    # restart re-embeds every PDF (real API cost) and risks a race if two
    # cold starts overlap — so we only build the index when it doesn't exist yet.
    pc.create_index(
        name=INDEX_NAME,
        dimension=DIMENSION,
        metric="euclidean",
        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
    )

pinecone_index = pc.Index(INDEX_NAME)
vector_store = PineconeVectorStore(pinecone_index=pinecone_index)

# If the index is already populated (from a prior run), just attach to it.
# Otherwise load the PDFs and build it now.
stats = pinecone_index.describe_index_stats()
if stats.get("total_vector_count", 0) > 0:
    index = VectorStoreIndex.from_vector_store(vector_store)
else:
    documents = SimpleDirectoryReader(
        input_dir="Data",  # folder name is case-sensitive — must match exactly what you upload
        required_exts=[".pdf"],
        file_extractor={".pdf": PDFReader()},
    ).load_data()

    if not documents:
        raise ValueError(
            "No PDF documents were loaded from the 'Data' folder. "
            "Make sure a folder named exactly 'Data' with your PDF(s) is uploaded alongside app.py."
        )

    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

query_engine = index.as_query_engine(text_qa_template=QA_TEMPLATE)


# --- Query function (logic unchanged) ---
@spaces.GPU
def query_doc(prompt):
    try:
        response = query_engine.query(prompt)
        return str(response)
    except Exception as e:
        logger.exception("Query failed")
        return f"Error: {str(e)}"


# =====================================================================
# UI LAYER — everything below is presentation only.
# Every answer still comes from query_doc(prompt) exactly as before.
# =====================================================================

DDS_LOGO_URL = (
    "https://raw.githubusercontent.com/Decoding-Data-Science/airesidency/main/"
    "DDS%20HIGH%20RES%20Logo%20%20copywhite.png"
)

# FAQ groups shown in the sidebar. Edit these to mirror the sections of your handbook.
FAQ_GROUPS = {
    "Leave & time off": [
        "How many days of annual leave am I entitled to?",
        "What is the process for applying for sick leave?",
        "What public holidays does DDS observe?",
    ],
    "Working at DDS": [
        "What are the standard working hours?",
        "What is the policy on remote or hybrid work?",
        "How long is the probation period for new employees?",
    ],
    "Policies & process": [
        "What is the notice period for resignation?",
        "How do I raise a grievance or workplace concern?",
        "What does the code of conduct cover?",
    ],
}


def add_user_message(message, history):
    """Show the user's question in the chat immediately and stash it for the answer step."""
    message = (message or "").strip()
    history = list(history or [])
    if not message:
        return "", history, ""
    history.append({"role": "user", "content": message})
    return "", history, message


def add_bot_answer(pending_question, history):
    """Run the unchanged query_doc() and append AYesha's answer."""
    history = list(history or [])
    if not pending_question:
        return history
    answer = query_doc(pending_question)
    history.append({"role": "assistant", "content": answer})
    return history


def clear_chat():
    return [], "", ""


# --- Theme: DDS Futuristic palette ---
dds_theme = gr.themes.Base(
    primary_hue=gr.themes.colors.cyan,
    secondary_hue=gr.themes.colors.violet,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
    radius_size=gr.themes.sizes.radius_md,
).set(
    body_background_fill="#070b14",
    body_background_fill_dark="#070b14",
    body_text_color="#e2e8f0",
    body_text_color_dark="#e2e8f0",
    body_text_color_subdued="#94a3b8",
    body_text_color_subdued_dark="#94a3b8",
    background_fill_primary="#0c1322",
    background_fill_primary_dark="#0c1322",
    background_fill_secondary="#0a101d",
    background_fill_secondary_dark="#0a101d",
    block_background_fill="#0c1322",
    block_background_fill_dark="#0c1322",
    block_border_color="#1e293b",
    block_border_color_dark="#1e293b",
    block_border_width="1px",
    block_label_text_color="#94a3b8",
    block_label_text_color_dark="#94a3b8",
    border_color_primary="#1e293b",
    border_color_primary_dark="#1e293b",
    input_background_fill="#0a101d",
    input_background_fill_dark="#0a101d",
    input_border_color="#243049",
    input_border_color_dark="#243049",
    input_border_color_focus="#22d3ee",
    input_border_color_focus_dark="#22d3ee",
    button_primary_background_fill="#22d3ee",
    button_primary_background_fill_dark="#22d3ee",
    button_primary_background_fill_hover="#67e8f9",
    button_primary_background_fill_hover_dark="#67e8f9",
    button_primary_text_color="#04111a",
    button_primary_text_color_dark="#04111a",
    button_secondary_background_fill="#111a2c",
    button_secondary_background_fill_dark="#111a2c",
    button_secondary_background_fill_hover="#16213a",
    button_secondary_background_fill_hover_dark="#16213a",
    button_secondary_text_color="#cbd5e1",
    button_secondary_text_color_dark="#cbd5e1",
    button_secondary_border_color="#243049",
    button_secondary_border_color_dark="#243049",
    color_accent_soft="#132238",
    color_accent_soft_dark="#132238",
)

CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@500;600;700&display=swap');

.gradio-container {
    max-width: 1240px !important;
    margin: 0 auto !important;
    padding: 24px 20px 12px !important;
}

/* ---------- Header ---------- */
#dds-header {
    display: flex;
    align-items: center;
    justify-content: space-between;
    gap: 24px;
    flex-wrap: wrap;
    padding: 22px 28px;
    border: 1px solid #1e293b;
    border-top: 3px solid #22d3ee;
    border-radius: 14px;
    background:
        radial-gradient(900px 220px at 100% 0%, rgba(139, 92, 246, 0.16), transparent 60%),
        #0a101d;
}
#dds-header .brand {
    display: flex;
    align-items: center;
    gap: 22px;
}
#dds-header .brand img {
    height: 76px;
    width: auto;
    display: block;
}
#dds-header .divider {
    width: 1px;
    height: 48px;
    background: #243049;
}
#dds-header h1 {
    font-family: 'Space Grotesk', 'Inter', sans-serif;
    font-size: 30px;
    font-weight: 700;
    letter-spacing: -0.02em;
    margin: 0;
    color: #f8fafc;
    line-height: 1.1;
}
#dds-header p {
    margin: 4px 0 0;
    font-size: 14px;
    color: #94a3b8;
    max-width: 560px;
}
#dds-header .status {
    display: flex;
    align-items: center;
    gap: 8px;
    font-size: 13px;
    color: #cbd5e1;
    padding: 8px 14px;
    border: 1px solid #243049;
    border-radius: 999px;
    background: #0c1322;
}
#dds-header .dot {
    width: 8px;
    height: 8px;
    border-radius: 50%;
    background: #22d3ee;
    box-shadow: 0 0 0 4px rgba(34, 211, 238, 0.15);
}

/* ---------- Panels ---------- */
.dds-panel {
    border: 1px solid #1e293b !important;
    border-radius: 14px !important;
    background: #0a101d !important;
    padding: 18px !important;
}
.dds-panel-title h3 {
    font-family: 'Space Grotesk', 'Inter', sans-serif;
    font-size: 16px;
    font-weight: 600;
    color: #f1f5f9;
    margin: 0 0 2px;
}
.dds-panel-title p {
    font-size: 13px;
    color: #94a3b8;
    margin: 0;
}
.faq-group-label p {
    font-size: 12px;
    font-weight: 600;
    color: #f5a623;
    margin: 14px 0 6px !important;
}

/* FAQ buttons read like list items, not CTAs */
.faq-btn {
    justify-content: flex-start !important;
    text-align: left !important;
    white-space: normal !important;
    font-weight: 400 !important;
    font-size: 13.5px !important;
    line-height: 1.4 !important;
    padding: 10px 12px !important;
    border-left: 2px solid #243049 !important;
}
.faq-btn:hover {
    border-left-color: #22d3ee !important;
    color: #f8fafc !important;
}

.scope-note {
    margin-top: 16px;
    padding: 12px 14px;
    border: 1px solid #2a2140;
    border-left: 3px solid #8b5cf6;
    border-radius: 10px;
    background: #0f0c1d;
    font-size: 13px;
    color: #c4b5fd;
    line-height: 1.5;
}

/* ---------- Chat ---------- */
#dds-chat {
    border: 1px solid #1e293b !important;
    border-radius: 12px !important;
    background: #070b14 !important;
}
#dds-input textarea {
    font-size: 15px !important;
}
#send-btn {
    font-weight: 600 !important;
    min-height: 44px;
}

/* ---------- Footer ---------- */
#dds-footer {
    text-align: center;
    font-size: 12.5px;
    color: #64748b;
    padding: 18px 0 4px;
    border-top: 1px solid #1e293b;
    margin-top: 8px;
}
#dds-footer a {
    color: #22d3ee;
    text-decoration: none;
}

footer { display: none !important; }

@media (max-width: 768px) {
    #dds-header { padding: 18px; }
    #dds-header .divider { display: none; }
    #dds-header .brand { flex-direction: column; align-items: flex-start; gap: 12px; }
    #dds-header h1 { font-size: 24px; }
}
"""

HEADER_HTML = f"""
<div id="dds-header">
  <div class="brand">
    <img src="{DDS_LOGO_URL}" alt="Decoding Data Science">
    <div class="divider"></div>
    <div>
      <h1>AYesha</h1>
      <p>The DDS Enterprise HR assistant. Answers come only from the current DDS HR Handbook.</p>
    </div>
  </div>
  <div class="status"><span class="dot"></span>Handbook connected</div>
</div>
"""

FOOTER_HTML = """
<div id="dds-footer">
  Built by <a href="https://decodingdatascience.com" target="_blank" rel="noopener">Decoding Data Science</a>
  &nbsp;|&nbsp; Answers are generated from the DDS HR Handbook. Always confirm critical decisions with HR.
</div>
"""

CHAT_PLACEHOLDER = """
<div style="text-align:center; padding: 40px 16px;">
  <h3 style="font-family:'Space Grotesk',Inter,sans-serif; font-size:20px; color:#f1f5f9; margin:0 0 8px;">
    Ask AYesha about DDS HR policy
  </h3>
  <p style="color:#94a3b8; font-size:14px; margin:0;">
    Type a question below, or pick one from the common questions on the left.
  </p>
</div>
"""


with gr.Blocks(title="AYesha | DDS Enterprise HR Chatbot", fill_height=False) as demo:
    pending_question = gr.State("")

    gr.HTML(HEADER_HTML)

    with gr.Row(equal_height=False):
        # ---------- Sidebar: FAQ ----------
        with gr.Column(scale=4, min_width=280, elem_classes="dds-panel"):
            gr.Markdown(
                "### Common questions\nClick any question to ask it instantly.",
                elem_classes="dds-panel-title",
            )
            faq_buttons = []
            for group_name, questions in FAQ_GROUPS.items():
                gr.Markdown(group_name, elem_classes="faq-group-label")
                for q in questions:
                    faq_buttons.append(
                        gr.Button(q, variant="secondary", size="sm", elem_classes="faq-btn")
                    )
            gr.HTML(
                '<div class="scope-note">AYesha only answers from the DDS HR Handbook. '
                "For confidential matters or anything the handbook does not cover, "
                "please contact the DDS HR team directly.</div>"
            )

        # ---------- Main: Chat ----------
        with gr.Column(scale=8, min_width=360, elem_classes="dds-panel"):
            chatbot = gr.Chatbot(
                elem_id="dds-chat",
                label="Conversation",
                show_label=False,
                height=520,
                placeholder=CHAT_PLACEHOLDER,
                buttons=["copy"],
                feedback_options=None,
                layout="bubble",
            )
            with gr.Row(equal_height=True):
                msg = gr.Textbox(
                    elem_id="dds-input",
                    placeholder="e.g. How many days of annual leave do I get?",
                    show_label=False,
                    lines=1,
                    max_lines=4,
                    scale=8,
                    autofocus=True,
                )
                send_btn = gr.Button("Ask AYesha", variant="primary", scale=2, elem_id="send-btn")
            with gr.Row():
                gr.HTML("")
                clear_btn = gr.Button("Clear conversation", variant="secondary", size="sm", scale=0, min_width=180)

    gr.HTML(FOOTER_HTML)

    # ---------- Events ----------
    # Step 1 shows the question instantly; step 2 calls query_doc() for the answer.
    for trigger in (msg.submit, send_btn.click):
        trigger(
            add_user_message,
            inputs=[msg, chatbot],
            outputs=[msg, chatbot, pending_question],
            queue=False,
        ).then(
            add_bot_answer,
            inputs=[pending_question, chatbot],
            outputs=[chatbot],
        )

    # A button passed as an input sends its own label, i.e. the FAQ question text.
    for btn in faq_buttons:
        btn.click(
            add_user_message,
            inputs=[btn, chatbot],
            outputs=[msg, chatbot, pending_question],
            queue=False,
        ).then(
            add_bot_answer,
            inputs=[pending_question, chatbot],
            outputs=[chatbot],
        )

    clear_btn.click(clear_chat, outputs=[chatbot, msg, pending_question], queue=False)


if __name__ == "__main__":
    # Gradio 6: theme and css are passed to launch(), not to gr.Blocks().
    # PORT is read so the same file runs on Hugging Face Spaces (7860) and Render ($PORT).
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        theme=dds_theme,
        css=CUSTOM_CSS,
        favicon_path=None,
    )