"""
Gradio UI для RAG QA.
- ключ Groq береться із секрету GROQ_API_KEY; поле в UI — fallback;
- перемикачі: BM25 / dense / reranker;
- вивід відповіді з цитатами [n] + перелік пронумерованих чанків.
"""
import os

import gradio as gr

import rag

SECRET_KEY = os.environ.get("GROQ_API_KEY", "")

print("Loading corpus + building indexes...")
CHUNKS = rag.load_corpus()
RETR = rag.Retrievers(CHUNKS)
RERANKER = rag.Reranker()
print(f"Ready. {len(CHUNKS)} chunks indexed.")


def run(query, api_key, use_bm25, use_dense, use_reranker):
    query = (query or "").strip()
    if not query:
        return "Введіть запитання.", ""
    # секрет має пріоритет; поле UI — лише запасний варіант
    key = SECRET_KEY or (api_key or "").strip()
    if not key:
        return "Groq-ключ не заданий (ні в секреті, ні в полі).", ""

    candidates = RETR.search(query, k=8, use_bm25=use_bm25, use_dense=use_dense)
    if not candidates:
        return "Увімкніть хоча б один метод пошуку.", ""

    final = RERANKER.rerank(query, candidates, top_n=4) if use_reranker else candidates[:4]

    try:
        ans = rag.answer(query, final, api_key=key)
    except Exception as e:
        return f"Помилка виклику Groq: {e}", ""

    sources = "\n\n".join(
        f"**[{i + 1}]** ({c.source})\n{c.text}" for i, c in enumerate(final)
    )
    return ans, sources


with gr.Blocks(title="RAG Question Answering") as demo:
    gr.Markdown(
        "# RAG Question Answering\n"
        "Система відповідає на запитання, спираючись на статті з "
        "`rag-mini-wikipedia`. Пошук релевантних чанків — BM25 (ключові слова) "
        "та семантичний (dense), з опційним cross-encoder реранкером. "
        "Відповідь генерує Groq LLM і цитує джерела у форматі `[n]`."
    )
    with gr.Row():
        with gr.Column(scale=2):
            query = gr.Textbox(label="Запитання", placeholder="Was Abraham Lincoln a president?")
            api_key = gr.Textbox(
                label="Groq API key", type="password", placeholder="gsk_...",
                visible=not SECRET_KEY,
            )
            with gr.Row():
                use_bm25 = gr.Checkbox(value=True, label="BM25 (ключові слова)")
                use_dense = gr.Checkbox(value=True, label="Dense (семантичний)")
                use_reranker = gr.Checkbox(value=True, label="Reranker")
            btn = gr.Button("Запитати", variant="primary")
        with gr.Column(scale=3):
            answer_box = gr.Markdown(label="Відповідь")
            sources_box = gr.Markdown(label="Джерела")

    btn.click(run, [query, api_key, use_bm25, use_dense, use_reranker],
              [answer_box, sources_box])

    gr.Examples(
        examples=[
            ["How did Abraham Lincoln's life come to an end?"],
            ["Who was Lincoln's Secretary of War?"],
            ["What did Lincoln do about slavery during his presidency?"],
            ["What programming language was used to build this app?"],
        ],
        inputs=query,
    )

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