"""Gradio UI for the Rust documentation assistant."""

import logging
from dataclasses import dataclass, field

import gradio as gr
from llama_index.core.base.llms.types import ChatMessage, MessageRole

from rag.config import DEFAULT_PROVIDER, PROVIDERS, RETRIEVAL_MODE, RETRIEVAL_TOP_K
from rag.index import download_index_if_missing, load_retriever
from rag.pipeline import get_response_stream
from rag.prompts import render_search, render_sources, source_count

logging.basicConfig(level=logging.INFO)
logging.getLogger("httpx").setLevel(logging.WARNING)

download_index_if_missing()


def get_retriever(book: str | None = None):
    """The retriever for one search, scoped to a book when the agent asked for one.

    `load_retriever` is cached on its arguments, so we build the unscoped
    retriever once and a scoped one at most once per book.
    """
    return load_retriever(RETRIEVAL_MODE, RETRIEVAL_TOP_K, book)


get_retriever()  # build the default eagerly to reduce latency for answer to first question 

DESCRIPTION = """Ask a question about Rust and get an answer grounded in the official documentation
— The Book, the Reference, Rust by Example, the Rustonomicon, and the async book.

Paste a key for one provider below. It is used only to generate answers, and is never stored or logged.
Search runs locally with an open source embedding model, so your key pays only for the answer."""

EXAMPLES = [
    "Why can't I have two mutable references to the same value?",
    "What does the ? operator do?",
    "How do I share mutable state between threads?",
    "error[E0502]: cannot borrow as mutable",
    "When should I use Box<dyn Error> instead of a custom error enum?",
    "How can I append a value to a vector?"
]


def on_provider_change(provider: str):
    spec = PROVIDERS[provider]
    return (
        gr.Dropdown(choices=list(spec.models), value=spec.default_model),
        # We restate `type` because this replaces the component, and a
        # Textbox built without it defaults to plain text, which would put a key
        # the visitor had already pasted on screen.
        gr.Textbox(
            label=spec.key_label,
            placeholder=f"Paste your {spec.key_label}",
            type="password",
        ),
    )


TITLE_LIMIT = 72


@dataclass
class Turn:
    question: str
    messages: list[gr.ChatMessage] = field(default_factory=list)
    answer: str = ""


def display_messages(turns: list[Turn], pending: Turn | None = None) -> list[gr.ChatMessage]:
    shown: list[gr.ChatMessage] = []
    for turn in [*turns, *([pending] if pending else [])]:
        shown.append(gr.ChatMessage(role="user", content=turn.question))
        shown.extend(turn.messages)
    return shown


def model_messages(turns: list[Turn]) -> list[ChatMessage]:
    history: list[ChatMessage] = []
    for turn in turns:
        if not turn.answer:
            continue
        history.append(ChatMessage(role=MessageRole.USER, content=turn.question))
        history.append(ChatMessage(role=MessageRole.ASSISTANT, content=turn.answer))
    return history


def answer_or_spinner(answer: str) -> gr.ChatMessage:
    if answer:
        return gr.ChatMessage(role="assistant", content=answer)
    return gr.ChatMessage(
        role="assistant", content="", metadata={"title": "Thinking…", "status": "pending"}
    )


def loop_messages(state) -> list[gr.ChatMessage]:
    messages = []
    for search in state.searches:
        query, nodes = search.query, search.nodes
        short = query if len(query) <= TITLE_LIMIT else query[:TITLE_LIMIT].rstrip() + "…"
        scope = f" · {search.book}" if search.book else ""
        counts = (
            f" — {len(nodes)} excerpts, {source_count(nodes)} sources" if nodes else ""
        )
        messages.append(
            gr.ChatMessage(
                role="assistant",
                content=render_search(query, search.thought, nodes),
                metadata={"title": f"🔍 {short}{scope}{counts}", "status": "done"},
            )
        )
    return messages


async def on_submit(
    question: str | None,
    turns: list[Turn],
    provider: str,
    model: str,
    api_key: str | None,
):
    question = (question or "").strip()
    api_key = api_key or ""

    if not question:
        yield display_messages(turns), "", turns
        return

    pending = Turn(question=question)
    yield display_messages(turns, pending), "", turns

    stream = get_response_stream(
        question, get_retriever, provider, api_key, model, model_messages(turns)
    )

    state = None
    answer = ""
    try:
        async for state in stream:
            answer = state.answer
            pending.messages = loop_messages(state) + [answer_or_spinner(answer)]
            yield display_messages(turns, pending), "", turns
    except Exception as exc:
        detail = f"**{type(exc).__name__}:** {exc}"
        body = f"{answer}\n\n{detail}".strip()

        pending.messages = (loop_messages(state) if state else []) + [
            gr.ChatMessage(role="assistant", content=body)
        ]
        yield display_messages(turns, pending), "", turns + [pending]
        return

    sources = render_sources(state.nodes, answer) if state else ""
    pending.messages = loop_messages(state) + [
        gr.ChatMessage(role="assistant", content=answer + sources)
    ]
    pending.answer = answer
    yield display_messages(turns, pending), "", turns + [pending]


def build_ui() -> gr.Blocks:
    default = PROVIDERS[DEFAULT_PROVIDER]

    with gr.Blocks(title="Rust Docs Assistant", fill_height=True) as chat:
        gr.Markdown("# Rust Docs Assistant 🦀")
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            provider = gr.Dropdown(
                choices=[(spec.label, slug) for slug, spec in PROVIDERS.items()],
                value=DEFAULT_PROVIDER,
                label="Provider",
            )
            model = gr.Dropdown(
                choices=list(default.models), value=default.default_model, label="Model"
            )

        api_key = gr.Textbox(
            label=default.key_label,
            placeholder=f"Paste your {default.key_label}",
            type="password",
        )

        turns = gr.State([])

        chatbot = gr.Chatbot(
            height=480,
            label="Conversation",
            buttons=["copy"],
            placeholder="Ask a question about Rust to get started.",
        )

        question = gr.Textbox(
            placeholder="Ask about ownership, lifetimes, traits, async…",
            show_label=False,
            submit_btn=True,
        )

        gr.Examples(examples=EXAMPLES, inputs=question, label="Try one")
         
        clear = gr.Button("Clear conversation", variant="secondary")

        provider.change(on_provider_change, inputs=provider, outputs=[model, api_key])

        question.submit(
            on_submit,
            inputs=[question, turns, provider, model, api_key],
            outputs=[chatbot, question, turns],
        )
         
        clear.click(lambda: ([], "", []), outputs=[chatbot, question, turns])

    return chat


if __name__ == "__main__":
    build_ui().queue(default_concurrency_limit=8).launch(footer_links=["settings"])