import os
from llama_index.core import VectorStoreIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.groq import Groq
from llama_index.core import Settings
from llama_index.readers.file import PyMuPDFReader
from llama_index.core.memory import ChatMemoryBuffer
import gradio as gr

os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")

Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.llm = Groq(model="llama-3.1-8b-instant")

chat_engine = None

INDUSTRY_PROMPTS = {
    "🏠 Real Estate": "You are a professional real estate assistant. Answer questions about properties, contracts, valuations and listings accurately.",
    "⚖️ Legal": "You are a legal document assistant. Answer questions about contracts and legal terms precisely. Always remind users to consult a qualified lawyer.",
    "💰 Finance & DeFi": "You are a financial and DeFi expert. Answer questions about financial reports, tokenomics and investments clearly.",
    "🏥 Medical": "You are a medical document assistant. Answer clearly and always remind users to consult a doctor.",
    "🚗 Car Dealership": "You are a car dealership assistant. Answer questions about vehicles, pricing and warranties helpfully.",
    "📄 General": "You are a helpful assistant. Answer questions from the uploaded document clearly."
}

css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

* { box-sizing: border-box; margin: 0; padding: 0; }

body {
    font-family: 'Inter', sans-serif !important;
    background: linear-gradient(135deg, #0f0c29, #302b63, #24243e) !important;
    min-height: 100vh;
}

.gradio-container {
    max-width: 860px !important;
    margin: 0 auto !important;
    padding: 16px !important;
    background: transparent !important;
}

.header-box {
    text-align: center;
    padding: 24px 16px 16px;
    color: white;
}

.mascot {
    font-size: 72px;
    display: block;
    margin-bottom: 8px;
    animation: float 3s ease-in-out infinite;
}

@keyframes float {
    0% { transform: translateY(0px); }
    50% { transform: translateY(-10px); }
    100% { transform: translateY(0px); }
}

.app-title {
    font-size: clamp(22px, 5vw, 40px) !important;
    font-weight: 700 !important;
    background: linear-gradient(90deg, #a78bfa, #60a5fa, #34d399);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 6px;
}

.app-subtitle {
    font-size: clamp(13px, 2vw, 15px);
    color: #94a3b8;
}

.card {
    background: rgba(255,255,255,0.06) !important;
    border: 1px solid rgba(255,255,255,0.12) !important;
    border-radius: 16px !important;
    padding: 16px !important;
    margin-bottom: 14px !important;
    backdrop-filter: blur(12px);
}

label {
    color: #e2e8f0 !important;
    font-weight: 500 !important;
    font-size: 13px !important;
}

footer { display: none !important; }

@media (max-width: 640px) {
    .gradio-container { padding: 8px !important; }
    .mascot { font-size: 52px; }
}
"""

def process_pdf(pdf_file, industry):
    global chat_engine
    if pdf_file is None:
        return "⚠️ Please upload a PDF first!"
    loader = PyMuPDFReader()
    documents = loader.load(file_path=pdf_file.name)
    index = VectorStoreIndex.from_documents(documents)
    memory = ChatMemoryBuffer.from_defaults(token_limit=3000)
    system_prompt = INDUSTRY_PROMPTS.get(industry, list(INDUSTRY_PROMPTS.values())[-1])
    chat_engine = index.as_chat_engine(
        chat_mode="context",
        memory=memory,
        system_prompt=system_prompt
    )
    return f"✅ Document loaded! Ask me anything."

def respond(message, chat_history):
    if not message.strip():
        return "", chat_history
    if chat_engine is None:
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": "⚠️ Please upload a PDF and select an industry first!"})
        return "", chat_history
    response = chat_engine.chat(message)
    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": str(response)})
    return "", chat_history

with gr.Blocks(css=css) as demo:

    gr.HTML("""
    <div class="header-box">
        <span class="mascot">🤖</span>
        <div class="app-title">Smart RAG Assistant</div>
        <div class="app-subtitle">Select your industry • Upload your document • Ask anything!</div>
    </div>
    """)

    with gr.Group(elem_classes="card"):
       industry_select = gr.Radio(
    choices=list(INDUSTRY_PROMPTS.keys()),
    value="📄 General",
    label="🏢 Select Your Industry",
)

    with gr.Group(elem_classes="card"):
        pdf_input = gr.File(
            label="📎 Upload Your PDF Document",
            file_types=[".pdf"]
        )
        status = gr.Textbox(
            label="Status",
            interactive=False,
            placeholder="Waiting for document..."
        )
        pdf_input.change(
            fn=process_pdf,
            inputs=[pdf_input, industry_select],
            outputs=status
        )

    with gr.Group(elem_classes="card"):
        chatbot = gr.Chatbot(
            height=380,
            label="💬 Chat"
        )
        msg = gr.Textbox(
            placeholder="Ask a question about your document...",
            label="Your Message",
            lines=1,
            max_lines=4
        )
        with gr.Row():
            submit = gr.Button("Send 🚀", variant="primary")
            clear = gr.Button("Clear 🗑️")

    submit.click(fn=respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
    msg.submit(fn=respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
    clear.click(lambda: [], outputs=chatbot)

demo.launch(share=True)