import os
import time
import openai
import gradio as gr

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# OpenAI API Key
openai.api_key = os.getenv("OPENAI_API_KEY")

# Load documents
documents = SimpleDirectoryReader("data").load_data()

# Build index
index = VectorStoreIndex.from_documents(documents=documents)

# Query engine
query_engine = index.as_query_engine()

# Query Function
def query_document(query, history):

    if history is None:
        history = []

    if not query.strip():
        return history, ""

    start_time = time.time()

    response = query_engine.query(query)

    end_time = time.time()

    execution_time = f"{end_time - start_time:.2f}"

    bot_response = f"""
{response}

⏱️ Response generated in {execution_time} sec
"""

    # Add user message
    history.append({
        "role": "user",
        "content": query
    })

    # Add assistant message
    history.append({
        "role": "assistant",
        "content": bot_response
    })

    return history, ""



# Custom CSS
custom_css = """
.gradio-container {
    background: linear-gradient(135deg, #0f172a, #111827);
    font-family: 'Segoe UI', sans-serif;
}

#chatbot {
    height: 520px;
    border-radius: 18px;
    border: 1px solid #374151;
    background: #1e293b;
    box-shadow: 0 8px 30px rgba(0,0,0,0.35);
}

textarea {
    border-radius: 14px !important;
    background: #111827 !important;
    color: white !important;
    border: 1px solid #374151 !important;
    padding: 12px !important;
    font-size: 15px !important;
}

button {
    border-radius: 12px !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}

button:hover {
    transform: scale(1.03);
}

.footer-text {
    text-align: center;
    color: #9ca3af;
    margin-top: 12px;
    font-size: 13px;
}
"""

# Theme
theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="slate",
    neutral_hue="gray",
    radius_size="lg",
)

# UI
with gr.Blocks(
    theme=theme,
    css=custom_css,
    title="DDS RAG Application"
) as demo:

    gr.Markdown(
        """
        # 🧠 DDS RAG Application Using LlamaIndex

        ### Intelligent Document Question Answering System

        Ask questions from uploaded documents [Paul_Graham] using AI-powered Retrieval Augmented Generation (RAG).
        """
    )

    chatbot = gr.Chatbot(
        label="AI Assistant",
        elem_id="chatbot",
        #bubble_full_width=False
    )

    query_box = gr.Textbox(
        placeholder="Ask something about your documents...",
        label="Enter Your Question",
        lines=2
    )

    with gr.Row():

        submit_btn = gr.Button(
            "🚀 Ask AI",
            variant="primary"
        )

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

    gr.Markdown(
        """
        <div class="footer-text">
        Powered by LlamaIndex • OpenAI • Gradio
        </div>
        """
    )

    submit_btn.click(
        fn=query_document,
        inputs=[query_box, chatbot],
        outputs=[chatbot, query_box]
    )

    query_box.submit(
        fn=query_document,
        inputs=[query_box, chatbot],
        outputs=[chatbot, query_box]
    )

    clear_btn.click(
        lambda: [],
        outputs=chatbot
    )

# Launch
if __name__ == "__main__":
    demo.launch()