"""Gradio app for the Zephyr docs RAG.

Deployed at https://huggingface.co/spaces/eoinedge/zephyrproject-rag and
runnable locally with `python app.py`.

Retrieval and generation are separate tabs on purpose. Retrieval alone is fast,
runs on CPU, and answers "where is this documented?" — which is often the whole
question. Generation is opt-in because on a free CPU Space it is slow, and a
slow wrong answer is worse than a fast citation.
"""

from __future__ import annotations

import json
import os
import sys
from functools import lru_cache
from pathlib import Path

import gradio as gr

sys.path.insert(0, str(Path(__file__).resolve().parent / "scripts"))
import retrieval  # noqa: E402

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

ROOT = Path(__file__).resolve().parent
INDEX_DIR = Path(os.environ.get("ZEPHYR_INDEX", ROOT / "data" / "index"))
GEN_MODEL = os.environ.get("ZEPHYR_GEN_MODEL", "Qwen/Qwen2.5-Coder-1.5B-Instruct")
ENABLE_GENERATION = os.environ.get("ZEPHYR_ENABLE_GENERATION", "0") == "1"

SYSTEM_PROMPT = (
    "You are a Zephyr RTOS assistant. Answer only from the documentation excerpts "
    "provided. Zephyr is large and moves fast, and a confident wrong answer costs an "
    "engineer hours on hardware — if the excerpts do not cover the question, say so "
    "and name what is missing rather than filling the gap from general RTOS knowledge. "
    "Cite the source path for anything you assert. Show configuration as it would "
    "actually appear: a Kconfig symbol, a devicetree node, a west command."
)

EXAMPLES = [
    "How do I define a devicetree binding for a new sensor?",
    "What is the difference between a work queue and a thread?",
    "How do I enable logging and set the log level for one module?",
    "How do I add a custom board definition?",
    "What does CONFIG_MAIN_STACK_SIZE control?",
    "How do I use west to build for a specific board?",
]


def resources():
    return retrieval.load(str(INDEX_DIR))


def search(question: str, k: int) -> list[dict]:
    """Shared with the CLI via scripts/retrieval.py — same index, same
    re-ranking, so a result here matches a result from `scripts/ask.py`."""
    try:
        return retrieval.search(question, int(k), INDEX_DIR)
    except FileNotFoundError as error:
        raise gr.Error(str(error)) from error


def render(hits: list[dict]) -> str:
    if not hits:
        return "Nothing retrieved."
    blocks = ["### Retrieved passages\n"]
    for position, hit in enumerate(hits, start=1):
        blocks.append(
            f"**{position}. [{hit['source']}]({hit['url']})** — {hit['section']}  \n"
            f"`similarity {hit['score']:.3f}`\n\n"
            f"```\n{hit['text'][:900]}\n```\n"
        )
    return "\n".join(blocks)


def do_retrieve(question: str, k: int) -> str:
    if not question.strip():
        return "Ask something about Zephyr."
    return render(search(question, k))


@lru_cache(maxsize=1)
def generator():
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL)
    model = AutoModelForCausalLM.from_pretrained(
        GEN_MODEL,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None,
    )
    return tokenizer, model


def do_answer(question: str, k: int) -> tuple[str, str]:
    if not question.strip():
        return "Ask something about Zephyr.", ""
    hits = search(question, k)
    if not hits:
        return "Nothing retrieved — the index may not cover this.", ""
    if not ENABLE_GENERATION:
        return (
            "Generation is off on this Space (CPU hardware makes it slow enough to be "
            "misleading). The retrieved passages below are the grounded answer; set "
            "`ZEPHYR_ENABLE_GENERATION=1` to turn the model on locally or on GPU "
            "hardware.",
            render(hits),
        )

    tokenizer, model = generator()
    context = "\n\n".join(
        f"[{index + 1}] {hit['source']} - {hit['section']}\n{hit['text']}"
        for index, hit in enumerate(hits)
    )
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Documentation excerpts:\n\n{context}\n\nQuestion: {question}"},
    ]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = model.generate(
        **inputs, max_new_tokens=512, do_sample=False, pad_token_id=tokenizer.eos_token_id
    )
    answer = tokenizer.decode(
        output[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
    ).strip()
    return answer, render(hits)


def corpus_note() -> str:
    try:
        _, _, _, meta = resources()
    except Exception:
        return "Index not built yet."
    source = str(meta.get("source", "")).replace("\n", " · ").strip(" ·")
    return (
        f"**{meta['chunks']:,} chunks** from **{meta['documents']} documents** · "
        f"embeddings `{meta['embedding_model']}` · {meta['index']}  \n"
        f"<sub>{source}</sub>"
    )


with gr.Blocks(title="Zephyr Docs RAG", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Zephyr RTOS documentation assistant")
    gr.Markdown(
        "Retrieval over the Zephyr documentation source, built from the RST in "
        "`zephyrproject-rtos/zephyr`. Every passage links back to the page it came "
        "from, so an answer can always be checked against the docs."
    )
    gr.Markdown(corpus_note())

    with gr.Tab("Ask"):
        question = gr.Textbox(label="Question", placeholder=EXAMPLES[0], lines=2)
        with gr.Row():
            top_k = gr.Slider(1, 12, value=6, step=1, label="Passages to retrieve")
            ask_button = gr.Button("Ask", variant="primary")
        answer_box = gr.Markdown(label="Answer")
        sources_box = gr.Markdown(label="Sources")
        gr.Examples(examples=EXAMPLES, inputs=question)
        ask_button.click(do_answer, [question, top_k], [answer_box, sources_box])
        question.submit(do_answer, [question, top_k], [answer_box, sources_box])

    with gr.Tab("Search only"):
        gr.Markdown(
            "Retrieval with no model in the loop. Fast, and often the whole answer — "
            "the passage you need with a link to the page."
        )
        search_question = gr.Textbox(label="Question", lines=2)
        search_k = gr.Slider(1, 20, value=8, step=1, label="Passages")
        search_button = gr.Button("Search", variant="primary")
        search_output = gr.Markdown()
        search_button.click(do_retrieve, [search_question, search_k], search_output)
        search_question.submit(do_retrieve, [search_question, search_k], search_output)


if __name__ == "__main__":
    demo.launch(
        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
    )