import os
import typing
import base64

import gradio as gr
import spaces

from openai import OpenAI
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np


# ---------- Secrets / OpenAI (HF Router) ----------
HF_AUTH = os.getenv("HF_AUTH", "")
if not HF_AUTH:
    raise RuntimeError(
        "Missing HF auth secret in environment. "
        "In Spaces, create a secret with name HF_AUTH (or change app.py to match your secret name)."
    )

client = OpenAI(
    base_url="https://router.huggingface.co/v1",
    api_key=HF_AUTH,
)



# ---------- Dataset config ----------
DATASET_ID = "gitrelief/Florida_Legislation_and_Summaries"
PREFERRED_FIELDS = [
    "full_text",
    "text",
    "content",
    "article",
    "body",
    "legislation",
    "summary",
    "summaries",
    "passage",
]

# Lazy-loaded globals (so the app can start immediately)
INDEX_READY = False
documents: list[str] = []
index = None
embed_model = None
text_field = None


def pick_text_field(example: dict) -> str:
    for f in PREFERRED_FIELDS:
        if f in example and isinstance(example[f], str) and example[f].strip():
            return f
    for k, v in example.items():
        if isinstance(v, str) and v.strip():
            return k
    raise ValueError(f"Could not find a usable text field. Keys: {list(example.keys())}")


def build_index(progress=gr.Progress(track_tqdm=True)):
    global INDEX_READY, documents, index, embed_model, text_field

    # 1) Load dataset
    progress(0.05, desc="Loading dataset…")
    ds = load_dataset(DATASET_ID, split="train")

    # 2) Decide which field to use
    text_field = pick_text_field(ds[0])

    # 3) Extract documents
    progress(0.15, desc=f"Building document list from field: {text_field}")
    documents.clear()
    for ex in ds:
        t = ex.get(text_field, "")
        if isinstance(t, str) and t.strip():
            documents.append(t)

    if not documents:
        raise RuntimeError("No documents found after extracting the text field.")

    # 4) Embed + FAISS
    progress(0.35, desc="Loading embedding model…")
    embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    progress(0.45, desc="Embedding documents…")
    embeddings = embed_model.encode(
        documents,
        batch_size=32,
        show_progress_bar=False,
        normalize_embeddings=True,
    )
    embeddings = np.asarray(embeddings, dtype="float32")

    progress(0.80, desc="Building FAISS index…")
    index = faiss.IndexFlatIP(embeddings.shape[1])
    index.add(embeddings)

    INDEX_READY = True
    progress(1.0, desc="Index ready")


def ensure_index():
    global INDEX_READY
    if not INDEX_READY:
        build_index()


def retrieve(query: str, k: int = 5) -> list[str]:
    q = embed_model.encode([query], normalize_embeddings=True)
    q = np.asarray(q, dtype="float32")
    scores, ids = index.search(q, k)

    hits: list[str] = []
    for i in ids[0]:
        if 0 <= i < len(documents):
            hits.append(documents[i])
    return hits


def build_messages(user_message: str, retrieved: typing.List[str]) -> list:
    system_prompt = (
        "You are Kimi K2.6. Answer the user using the provided Florida legislation text as context. "
        "If the context is insufficient, say so and ask a clarifying question."
    )

    context_block = "\n\n".join([f"[Context {i+1}]\n{txt}" for i, txt in enumerate(retrieved)])

    return [
        {"role": "system", "content": system_prompt},
        {
            "role": "user",
            "content": f"Retrieved context:\n{context_block}\n\nUser question: {user_message}",
        },
    ]


# ---------- UI handlers ----------
with gr.Blocks() as demo:
    gr.Markdown("## Kimi-K2.6 + Florida Legislation (RAG, full text preferred)")

    status = gr.Textbox(label="Status", interactive=False)
    inp = gr.Textbox(label="Your question", lines=4)
    out = gr.Textbox(label="Answer", lines=12)

    submit = gr.Button("Submit")
    build_btn = gr.Button("Build/Reload Index")

    def on_build():
        nonlocal_status = "Building index (this may take a while)…"
        try:
            build_index()
            return "Index ready."
        except Exception as e:
            return f"Index build failed: {type(e).__name__}: {e}"

    build_btn.click(fn=on_build, inputs=None, outputs=status)

    @spaces.GPU(duration=120)
    def chat_handler(message: str) -> str:
        # Ensure index exists before doing GPU work
        if not message or not message.strip():
            return "Please enter a question."

        if not INDEX_READY:
            # Build on first ask (still better than failing silently)
            build_index()

        retrieved = retrieve(message, k=5)
        messages = build_messages(message, retrieved)

        completion = client.chat.completions.create(
            model="moonshotai/Kimi-K2.6:fireworks-ai",
            messages=messages,
            stream=False,
        )
        return completion.choices[0].message.content

    submit.click(fn=chat_handler, inputs=inp, outputs=out)


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