import os
import gradio as gr
import openai
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# --- 1. INITIALIZATION & SETUP ---
# Retrieve API key from environment variables (Hugging Face Secrets)
api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
    raise ValueError("OPENAI_API_KEY environment variable is not set.")
openai.api_key = api_key

# Define Ayesha's Persona
system_prompt_str = (
    "You are Cirilla, the Decoding Data Science (DDS) Enterprise HR Chatbot. "
    "Your objective is to interact politely and professionally with employees, answering only HR-related questions. "
    "Use only information directly from the connected HR documents to provide your answers. "
    "Always provide an explicit citation indicating the document source for every answer. "
    "Do not offer information or suggestions beyond what is present in these documents. "
    "If the requested information cannot be found, politely instruct the user to email connect@decodingdatascience.com. "
    "For questions outside of HR, inform the user that you can only answer HR-related questions."
)

# Configure LLM and Embeddings
Settings.llm = OpenAI(model="gpt-5.4-nano", temperature=0.2, system_prompt=system_prompt_str)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.chunk_size = 600
Settings.chunk_overlap = 200

# Load data and create index 
print("Initializing Knowledge Base...")
try:
    documents = SimpleDirectoryReader(
        input_dir="data",
        required_exts=[".pdf"]
    ).load_data()
    print(f"Successfully loaded {len(documents)} documents.")
except Exception as e:
    print(f"Error loading documents: {e}. Please ensure the 'data' folder exists and contains PDFs.")
    documents = []

# Build the Vector Index
index = VectorStoreIndex.from_documents(documents=documents)

# Upgrade to a Chat Engine so it remembers conversation history
chat_engine = index.as_chat_engine(
    chat_mode="condense_question",
    verbose=True
)

# --- 2. BACKEND CHAT FUNCTION ---
def chat_with_ayesha(user_message, history):
    if not user_message.strip():
        return "", history

    # Append user message to history immediately for UI responsiveness
    history.append((user_message, None))
    yield "", history

    # Query the LlamaIndex chat engine
    response = chat_engine.chat(user_message)

    # Update the history with the AI's response
    history[-1] = (user_message, response.response)
    yield "", history

def clear_chat():
    chat_engine.reset()
    return []

# --- 3. FRONTEND UI ARCHITECTURE ---
# REMOVED theme from gr.Blocks to comply with Gradio 6.0+
with gr.Blocks(theme=gr.themes.Soft(), fill_width=True) as app:
    with gr.Row():
        gr.Markdown(
            """
            # 🏢 Decoding Data Science HR - Assist
            ### *Ask Cirilla: Your Enterprise HR Assistant*
            """
        )

    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            # FIXED: Removed show_download_button=False to fix the crash
            gr.Image(value="https://cdn-icons-png.flaticon.com/512/4712/4712010.png", width=150, show_label=False)
            gr.Markdown("### 🤖 Agent Status: **Online**")
            gr.Markdown("**Role:** Enterprise HR Specialist")
            gr.Markdown("**Knowledge Base:** Connected to DDS internal HR PDFs.")

            with gr.Accordion("How to use this tool", open=False):
                gr.Markdown(
                    """
                    - Ask about leave policies, office hours, or benefits.
                    - Cirilla will cite the specific document she pulled the answer from.
                    - For non-HR issues, please contact IT.
                    """
                )

            clear_btn = gr.Button("🗑️ Clear Conversation", variant="secondary")

        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="Conversation with Cirilla",
                height=500,
                #layout="bubble",
                bubble_full_width=False,
                #type="tuples", 
                avatar_images=(None, "https://cdn-icons-png.flaticon.com/512/4712/4712139.png") 
            )

            with gr.Row():
                msg_input = gr.Textbox(
                    show_label=False,
                    placeholder="Type your HR question here and press Enter...",
                    container=False,
                    scale=4
                )
                submit_btn = gr.Button("Send ✉️", variant="primary", scale=1)

    # --- 4. EVENT BINDING ---
    msg_input.submit(
        fn=chat_with_ayesha,
        inputs=[msg_input, chatbot],
        outputs=[msg_input, chatbot]
    )

    submit_btn.click(
        fn=chat_with_ayesha,
        inputs=[msg_input, chatbot],
        outputs=[msg_input, chatbot]
    )

    clear_btn.click(
        fn=clear_chat,
        inputs=[],
        outputs=[chatbot]
    )

# MOVED theme configuration here to comply with Gradio 6.0+ requirements
#app.launch(theme=gr.themes.Soft())
app.launch(debug=True)