import os
import gradio as gr
from pathlib import Path
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex

# Ensure data directory exists
Path("data").mkdir(exist_ok=True)

# Set OpenAI API key
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

# Global state for engines
state = {"insurance_engine": None, "paul_engine": None, "custom_engine": None}

# Build index from file
def build_index_from_file(file_path):
    documents = SimpleDirectoryReader(input_files=[file_path]).load_data()
    index = VectorStoreIndex.from_documents(documents)
    return index.as_query_engine()

# Preload fixed documents
def preload_indexes():
    try:
        state["insurance_engine"] = build_index_from_file("data/insurance_FirstRAG.pdf")
    except:
        print("❌ Failed to load insurance.pdf")

    try:
        state["paul_engine"] = build_index_from_file("data/paul_graham_FirstRAG.txt")
    except:
        print("❌ Failed to load paul_graham.txt")

# Upload custom file
def upload_and_refresh(file_path):
    if file_path is None:
        return "⚠️ No file uploaded."
    try:
        state["custom_engine"] = build_index_from_file(file_path)
        return f"✅ Uploaded & Indexed: {os.path.basename(file_path)}"
    except Exception as e:
        return f"❌ Failed: {str(e)}"

# Query helpers
def ask_insurance(query): return _query_engine(state["insurance_engine"], query)
def ask_paul(query): return _query_engine(state["paul_engine"], query)
def ask_custom(query): return _query_engine(state["custom_engine"], query)

def _query_engine(engine, query):
    if not query:
        return "❗ Please enter a question."
    if engine is None:
        return "📂 Document not loaded yet."
    try:
        return str(engine.query(query))
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Summarize
def summarize_insurance(): return _query_engine(state["insurance_engine"], "Summarize the insurance document.")
def summarize_paul(): return _query_engine(state["paul_engine"], "Summarize Paul Graham's main ideas.")
def summarize_custom(): return _query_engine(state["custom_engine"], "Summarize the uploaded documents.")

# Clear
def clear_fields(): return "", ""

# Launch app
def launch():
    preload_indexes()

    with gr.Blocks(
        title="RAG App with LlamaIndex",
        css="""
        body {
            background-color: #f5f5dc;
            font-family: 'Georgia', 'Merriweather', serif;
        }
        h1 {
            font-size: 2.5em;
            font-weight: bold;
            color: #4e342e;
            margin-bottom: 0.3em;
        }
        .subtitle {
            font-size: 1.2em;
            color: #6d4c41;
            margin-bottom: 1.5em;
        }
        .gr-box, .gr-column, .gr-group {
            border-radius: 15px;
            padding: 20px;
            background-color: #fffaf0;
            box-shadow: 2px 4px 14px rgba(0, 0, 0, 0.1);
            margin-top: 10px;
        }
        textarea, input[type="text"], input[type="file"] {
            background-color: #fffaf0;
            border: 1px solid #d2b48c;
            color: #4e342e;
            border-radius: 8px;
        }
        button {
            background-color: #a1887f;
            color: white;
            font-weight: bold;
            border-radius: 8px;
            transition: background-color 0.3s ease;
        }
        button:hover {
            background-color: #8d6e63;
        }
        .gr-button {
            border-radius: 8px !important;
        }
        .tabitem {
            border-radius: 15px !important;
        }
        """
    ) as app:
        
        with gr.Column():
            gr.Markdown("""
                <div style='text-align: center;'>
                    <h1>RAG Application with LlamaIndex</h1>
                    <div class='subtitle'>📄 Ask questions from documents using Retrieval-Augmented Generation (RAG)</div>
                </div>
            """)

        # 📘 Insurance Tab
        with gr.Tab("📘 Insurance Summary"):
            with gr.Column():
                gr.Markdown("### 🛡️ Ask about Insurance Document")
                with gr.Group():
                    insurance_q = gr.Textbox(label="Your Question", placeholder="e.g., What does this cover?")
                    with gr.Row():
                        insurance_ask = gr.Button("Submit")
                        insurance_clear = gr.Button("Clear")
                    insurance_ans = gr.Textbox(label="Response", lines=6)
                    insurance_ask.click(fn=ask_insurance, inputs=insurance_q, outputs=insurance_ans)
                    insurance_clear.click(fn=clear_fields, outputs=[insurance_q, insurance_ans])

                with gr.Group():
                    summarize_btn = gr.Button("Summarize Insurance Document")
                    summarize_output = gr.Textbox(label="Summary", lines=6)
                    summarize_btn.click(fn=summarize_insurance, outputs=summarize_output)

        # 🧠 Paul Graham Tab
        with gr.Tab("🧠 Paul Graham"):
            with gr.Column():
                gr.Markdown("### 🧠 Ask about Paul Graham's Writings")
                with gr.Group():
                    paul_q = gr.Textbox(label="Your Question", placeholder="e.g., What does Paul say about startups?")
                    with gr.Row():
                        paul_ask = gr.Button("Ask")
                        paul_clear = gr.Button("Clear")
                    paul_ans = gr.Textbox(label="Response", lines=6)
                    paul_ask.click(fn=ask_paul, inputs=paul_q, outputs=paul_ans)
                    paul_clear.click(fn=clear_fields, outputs=[paul_q, paul_ans])

                with gr.Group():
                    summarize_btn2 = gr.Button("Summarize Paul Graham's Ideas")
                    summarize_output2 = gr.Textbox(label="Summary", lines=6)
                    summarize_btn2.click(fn=summarize_paul, outputs=summarize_output2)

        # 📤 Upload Custom
        with gr.Tab("📤 Upload & Ask Your File"):
            with gr.Column():
                gr.Markdown("### 📤 Upload Your Own Document")
                with gr.Group():
                    file_input = gr.File(label="Upload PDF or TXT", type="filepath")
                    status = gr.Textbox(label="Status", interactive=False)
                    file_input.change(fn=upload_and_refresh, inputs=file_input, outputs=status)

                with gr.Group():
                    gr.Markdown("#### 💬 Ask a Question")
                    custom_q = gr.Textbox(label="Your Query")
                    with gr.Row():
                        custom_ask = gr.Button("Ask")
                        custom_clear = gr.Button("Clear")
                    custom_ans = gr.Textbox(label="Response", lines=6)
                    custom_ask.click(fn=ask_custom, inputs=custom_q, outputs=custom_ans)
                    custom_clear.click(fn=clear_fields, outputs=[custom_q, custom_ans])

                with gr.Group():
                    gr.Markdown("#### 📌 Summarize the File")
                    summarize = gr.Button("Summarize Uploaded Document")
                    summary_output = gr.Textbox(label="Summary", lines=6)
                    summarize.click(fn=summarize_custom, outputs=summary_output)

    app.launch()


# Run the app
launch()