import os
import queue
from threading import Thread
from dotenv import load_dotenv

print(f"Start loading .env")
load_dotenv()
print(f"Finish loading .env")

from langchain.callbacks.base import BaseCallbackHandler
print(f"Start importing from rag_func")
from rag_func import prepare_RAG, retrieve_RAG, generate_RAG
print(f"Finish importing from rag_func")
import gradio as gr

# ----------------- Context Setup -----------------
# For local execution, where the user can select the directory
# user_input = input("Enter a subfolder inside 'context' (press Enter for full 'context'): ").strip()
# if user_input:
#     user_dir = os.path.join("context", user_input)
# else:
#     user_dir = "context"
# print(f"[Info] Using context directory: {user_dir}")

# For gradio deploy. The user cannot choose the directory
user_dir = "context"

pinecone_API = os.getenv("PINECONE_API")
index_name = "regulatorik"
#llm_model = os.getenv("MODELNAME")
llm_model="gpt-5-nano"

index, pc, llm = prepare_RAG(pinecone_API, index_name, llm_model=llm_model, dir_name=user_dir)

# ----------------- Chat Functions -----------------
def add_user_message(message, history):
    history = history or []
    history.append({"role": "user", "content": message})
    return "", history, history

import time

# ----------------- Streaming Handler -----------------
class StreamHandler(BaseCallbackHandler):
    def __init__(self, q: queue.Queue):
        self.q = q
        self.first_token_received = False
        self.ttft = None  # time to first token

    def on_llm_new_token(self, token: str, **kwargs):
        if not self.first_token_received:
            self.ttft = time.time() - self.start_time
            self.first_token_received = True
        self.q.put(token)

    def on_llm_end(self, *args, **kwargs):
        self.total_time = time.time() - self.start_time
        self.q.put("[[END]]")

# ----------------- Chat Functions with timing -----------------
def generate_bot_response(history):
    if not history or history[-1]["role"] != "user":
        yield history, history
        return

    user_msg = history[-1]["content"]
    retrieved_chunks = retrieve_RAG(user_msg, pc, index)

    q = queue.Queue()
    handler = StreamHandler(q)
    handler.start_time = time.time()

    model_name = getattr(llm, "model_name", getattr(llm, "model", None))
    streaming_llm = llm.__class__(model=model_name, streaming=True, callbacks=[handler])

    def run_llm():
        try:
            generate_RAG(user_msg, streaming_llm, retrieved_chunks)
        finally:
            q.put("[[END]]")

    Thread(target=run_llm, daemon=True).start()

    partial = ""
    history.append({"role": "assistant", "content": ""})

    while True:
        token = q.get()
        if token == "[[END]]":
            print(f"[Timing] TTFT: {handler.ttft:.3f} s, Total: {handler.total_time:.3f} s")
            break
        partial += token
        history[-1]["content"] = partial
        yield history, history

# ----------------- Simplified CSS for Default Gradio Font -----------------
custom_css = """
:root {
    --brand-blue: #17428f;
    --brand-orange: #f39719;
    --text-dark: #111827;   /* very dark grey (near black) */
    --text-gray: #4B5563;   /* medium grey for messages */
    color-scheme: only light;
}

body, .gradio-container {
    /* Default Gradio font will be used */
    background: linear-gradient(135deg, var(--brand-blue) 0%, var(--brand-orange) 100%);
    min-height: 100vh;
    color: var(--text-dark);
}

/* Top bar transparent */
#topbar { background: transparent !important; }

/* Header text over gradient */
#header h1, #header h2, #header h3, #header h4, #header h5, #header h6,
#header p {
    color: #ffffff;
    text-align: center;
}

/* Chatbox container */
#chatbot {
    height: 600px;
    border-radius: 14px;
    border: 2px solid var(--brand-blue);
    background-color: #ffffff;
    padding: 8px;
    overflow-y: auto;
}

/* Chat text */
#chatbot, #chatbot * { color: var(--text-gray) !important; }

/* Bubble styling */
#chatbot .message.user {
    background: #fff4e1;
    border-radius: 10px;
    padding: 6px 12px;
    color: var(--text-gray) !important;
    text-align: right;
}
#chatbot .message.bot {
    background: #f0f0f0;
    border-radius: 10px;
    padding: 6px 12px;
    color: var(--text-gray) !important;
    text-align: left;
}

/* Fallback selectors for other Gradio versions */
#chatbot [data-testid*="message"] { border-radius: 10px; padding: 6px 12px; }
#chatbot [data-testid="user-message"] { background: #fff4e1; color: var(--text-gray) !important; text-align: right; }
#chatbot [data-testid="assistant-message"] { background: #f0f0f0; color: var(--text-gray) !important; text-align: left; }

/* Inputs */
input[type="text"], textarea, .gr-text-input input, .gr-textbox textarea {
    border-radius: 10px;
    padding: 10px;
    font-size: 16px;
    border: 2px solid var(--brand-orange);
}
input:focus, textarea:focus, .gr-text-input input:focus, .gr-textbox textarea:focus {
    border-color: var(--brand-blue);
    outline: none;
    box-shadow: 0 0 6px rgba(23, 66, 143, 0.5);
}

/* Buttons (global gradient) */
.gr-button, button {
    border-radius: 10px;
    font-weight: 600;
    background: linear-gradient(90deg, var(--brand-blue), var(--brand-orange));
    color: white;
    border: none;
}
.gr-button:hover, button:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 8px rgba(0,0,0,0.2);
}

/* Chat area: icon-only buttons */
#chatbot {
  --icon-light: #9CA3AF;
  --icon-hover: #6B7280;
}

/* Tint SVG icons */
#chatbot button svg,
#chatbot [role="button"] svg,
#chatbot .icon svg,
#chatbot [class*="icon"] svg,
#chatbot [data-testid*="icon"] svg,
#chatbot [data-testid*="message"] .tools svg,
#chatbot .message-tools svg {
  color: var(--icon-light) !important;
  fill: var(--icon-light) !important;
  stroke: var(--icon-light) !important;
  opacity: 0.95;
}

/* Remove gradient background only on small icon-only buttons */
#chatbot :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient) {
  background: transparent !important;
  background-image: none !important;
  border: none !important;
  box-shadow: none !important;
  padding: 6px !important;
  border-radius: 8px !important;
  color: var(--icon-light) !important;
}

/* Hover/focus/active states */
#chatbot :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient):hover {
  background-color: rgba(0,0,0,0.05) !important;
}
#chatbot :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient):focus-visible {
  outline: none !important;
  box-shadow: 0 0 0 2px rgba(23, 66, 143, 0.35) !important;
  background-color: rgba(0,0,0,0.06) !important;
}
#chatbot :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient):active {
  background-color: rgba(0,0,0,0.08) !important;
}

/* Optional 'danger' icons */
#chatbot .danger svg { color: var(--icon-light) !important; fill: var(--icon-light) !important; stroke: var(--icon-light) !important; }
#chatbot .danger:hover svg { color: #ef4444 !important; fill: #ef4444 !important; stroke: #ef4444 !important; }

#topbar .gr-button.keep-gradient,
#topbar .gr-button:not(:has(svg)) {
  background: linear-gradient(90deg, var(--brand-blue), var(--brand-orange)) !important;
  color: #fff !important;
}

/* Icon-only buttons in topbar: transparent */
#topbar :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient) {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  padding: 6px !important;
  border-radius: 8px !important;
  color: var(--icon-light) !important;
}

/* Tint SVGs in topbar */
#topbar :is(button,[role="button"]):has(> svg) > svg {
  color: var(--icon-light) !important;
  fill: var(--icon-light) !important;
  stroke: var(--icon-light) !important;
  opacity: 0.95;
}

/* Hover/focus/active for topbar icons */
#topbar :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient):hover {
  background-color: rgba(0,0,0,0.05) !important;
}
#topbar :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient):focus-visible {
  outline: none !important;
  box-shadow: 0 0 0 2px rgba(23, 66, 143, 0.35) !important;
  background-color: rgba(0,0,0,0.06) !important;
}
#topbar :is(button,[role="button"]):is([aria-label],[title], :has(> svg)):not(.keep-gradient):active {
  background-color: rgba(0,0,0,0.08) !important;
}

"""

js_force_light = """ function refresh() { const url = new URL(window.location); if (url.searchParams.get('__theme') !== 'light') { url.searchParams.set('__theme', 'light'); window.location.replace(url); } } """

# ----------------- Gradio App -----------------
with gr.Blocks(css=custom_css, fill_height=True, js=js_force_light) as demo:

    gr.Markdown(
        "<h1>📚 RAG Chat Assistant</h1>"
        "<p>Ask questions and get accurate answers from your documents.</p>",
        elem_id="header"
    )

    chatbot = gr.Chatbot(type="messages", label="Conversation", elem_id="chatbot", height=600)
    msg = gr.Textbox(label="Your message", placeholder="Type your question here...")
    state = gr.State([])

    msg.submit(add_user_message, inputs=[msg, state], outputs=[msg, chatbot, state]) \
       .then(generate_bot_response, inputs=[state], outputs=[chatbot, state])

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