import os
import re
import string
import warnings
from collections import Counter
from dataclasses import dataclass, field
from typing import Optional

import faiss
import gradio as gr
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer


warnings.filterwarnings("ignore")


TOOL_NAME = "SourceSure"
TOOL_TAGLINE = "Grounded answers with a clear source trail."


# Config: keep these aligned with the Kaggle notebook defaults.
CHECKPOINT = "selfrag/selfrag_llama2_7b"
DELTA = 0.2
BEAM_WIDTH = 2
K_PASSAGES = 5
MAX_NEW_TOKS = 100
MAX_SEGMENTS = 5
W_REL = 1.0
W_SUP = 1.0
W_USE = 0.5
ABSTENTION_ISREL_THRESHOLD = 0.20
ABSTENTION_MIN_QUERY_COVERAGE = 0.20

# Optional runtime tweak for Hugging Face Spaces.
# Default stays aligned with the notebook: float16 on GPU.
LOAD_IN_4BIT = os.getenv("LOAD_IN_4BIT", "0") == "1"


class RetrieveToken:
    YES = "[Retrieve=Yes]"
    NO = "[Retrieve=No]"
    CONTINUE = "[Retrieve=Continue]"
    ALL = [YES, NO, CONTINUE]


class IsRelToken:
    RELEVANT = "[IsRel=Relevant]"
    IRRELEVANT = "[IsRel=Irrelevant]"
    ALL = [RELEVANT, IRRELEVANT]


class IsSupportToken:
    FULLY = "[IsSupport=Fully Supported]"
    PARTIALLY = "[IsSupport=Partially Supported]"
    NO = "[IsSupport=No Support]"
    ALL = [FULLY, PARTIALLY, NO]


class IsUseToken:
    FIVE = "[IsUse=5]"
    FOUR = "[IsUse=4]"
    THREE = "[IsUse=3]"
    TWO = "[IsUse=2]"
    ONE = "[IsUse=1]"
    ALL = [FIVE, FOUR, THREE, TWO, ONE]
    WEIGHTS = {1: -1.0, 2: -0.5, 3: 0.0, 4: 0.5, 5: 1.0}


ALL_REFLECTION_TOKENS = (
    RetrieveToken.ALL + IsRelToken.ALL + IsSupportToken.ALL + IsUseToken.ALL
)


DOCUMENTS = [
    {
        "id": "doc_00",
        "title": "RAG Overview",
        "text": (
            "Retrieval-Augmented Generation (RAG) improves factual correctness in large "
            "language models by combining a neural generator with a non-parametric retrieval "
            "system. The model is given relevant documents at inference time instead of "
            "relying solely on memorised parametric knowledge. RAG was introduced by Lewis "
            "et al. in 2020."
        ),
    },
    {
        "id": "doc_01",
        "title": "Vanilla RAG Limitations",
        "text": (
            "Vanilla RAG always retrieves a fixed number of top-k documents regardless of "
            "whether retrieval is actually necessary for the query. This indiscriminate "
            "retrieval can introduce noise when the documents are not relevant to the "
            "question, and the model may produce outputs that are inconsistent with "
            "retrieved passages."
        ),
    },
    {
        "id": "doc_02",
        "title": "SELF-RAG",
        "text": (
            "SELF-RAG is a framework that trains a single language model to retrieve, "
            "generate, and critique text through self-reflection. It uses four special "
            "reflection tokens: Retrieve (yes/no/continue), IsRel (relevant/irrelevant), "
            "IsSupport (fully/partially/no support), and IsUse (1 to 5). These tokens "
            "allow the model to decide adaptively when to retrieve and to evaluate the "
            "quality of its own generation at inference time without any external critic "
            "model."
        ),
    },
    {
        "id": "doc_03",
        "title": "CRAG - Corrective RAG",
        "text": (
            "Corrective Retrieval-Augmented Generation (CRAG) addresses the problem of "
            "inaccurate or irrelevant retrieved documents by adding a lightweight evaluator "
            "that grades each retrieved document's relevance. When all retrieved documents "
            "are judged irrelevant, CRAG triggers a corrective web search to obtain better "
            "evidence before generation."
        ),
    },
    {
        "id": "doc_04",
        "title": "RAGAs Evaluation",
        "text": (
            "RAGAs is an automated evaluation framework for retrieval-augmented generation "
            "systems. It provides metrics including faithfulness, which measures whether "
            "the generated answer is supported by the retrieved context, and answer "
            "relevance, which measures whether the answer addresses the original question."
        ),
    },
    {
        "id": "doc_05",
        "title": "Hallucination in LLMs",
        "text": (
            "Hallucination in large language models refers to the generation of fluent and "
            "confident statements that are not supported by evidence or are factually "
            "incorrect. RAG reduces hallucination by grounding generation in retrieved "
            "documents, but does not eliminate it entirely when the model ignores or "
            "misuses the retrieved context."
        ),
    },
    {
        "id": "doc_06",
        "title": "Adaptive Retrieval in SELF-RAG",
        "text": (
            "SELF-RAG uses an adaptive retrieval threshold delta to control how often "
            "retrieval is triggered. The model predicts a Retrieve token and if the "
            "probability of Retrieve=Yes divided by the sum of Retrieve=Yes and Retrieve=No "
            "exceeds the threshold delta, retrieval is triggered. The default threshold is "
            "0.2 for most tasks."
        ),
    },
    {
        "id": "doc_07",
        "title": "Segment-Level Beam Search",
        "text": (
            "SELF-RAG uses segment-level beam search with beam width B=2 during inference. "
            "For each retrieved passage, the model generates a candidate continuation scored "
            "using a weighted sum of IsRel, IsSupport, and IsUse token probabilities "
            "combined with the log-probability of the generated segment. Default weights "
            "are 1.0 for IsRel, 1.0 for IsSupport, and 0.5 for IsUse."
        ),
    },
]


EXAMPLES = [
    ["What is SELF-RAG and how does it use reflection tokens?"],
    ["What is the main limitation of Vanilla RAG?"],
    ["How does RAG improve factual correctness?"],
    ["What is hallucination in language models?"],
    ["How does CRAG handle irrelevant retrieved documents?"],
    [
        "What is the deployment date and version number of the AuthService "
        "microservice in the Pegasus project?"
    ],
]


@dataclass
class Passage:
    id: str
    title: str
    text: str


@dataclass
class CritiqueScores:
    p_relevant: float = 0.0
    p_irrelevant: float = 0.0
    p_fully: float = 0.0
    p_partially: float = 0.0
    p_no_support: float = 0.0
    p_use_1: float = 0.0
    p_use_2: float = 0.0
    p_use_3: float = 0.0
    p_use_4: float = 0.0
    p_use_5: float = 0.0


@dataclass
class SegmentResult:
    text: str
    retrieve_tok: str
    passage: Optional[Passage] = None
    isrel: Optional[str] = None
    issup: Optional[str] = None
    isuse: Optional[str] = None
    score: float = 0.0
    is_sufficient: bool = True
    isrel_score: float = 0.0
    query_coverage: float = 0.0


@dataclass
class SelfRAGOutput:
    query: str
    segments: list = field(default_factory=list)
    abstained: bool = False

    @property
    def answer(self) -> str:
        return " ".join(s.text for s in self.segments if s.is_sufficient).strip()


class LightweightRetriever:
    def __init__(self, device: str = "cpu"):
        self.device = device
        self.model = SentenceTransformer("all-MiniLM-L6-v2", device=device)
        self.passages = []
        self.index = None

    def index_documents(self, docs):
        self.passages = [Passage(**d) for d in docs]
        texts = [f"{p.title} {p.text}".strip() for p in self.passages]
        embs = self.model.encode(
            texts, convert_to_numpy=True, normalize_embeddings=True
        ).astype("float32")
        self.index = faiss.IndexFlatIP(embs.shape[1])
        self.index.add(embs)

    def retrieve(self, query: str, k: int = K_PASSAGES):
        q_emb = self.model.encode(
            [query], convert_to_numpy=True, normalize_embeddings=True
        ).astype("float32")
        _, idxs = self.index.search(q_emb, k)
        return [self.passages[i] for i in idxs[0] if i < len(self.passages)]


retriever = None
gen_tokenizer = None
gen_model = None
TOKEN_IDS = {}
MODEL_LOAD_ERROR = None


def score_isrel(cs: CritiqueScores) -> float:
    denom = cs.p_relevant + cs.p_irrelevant
    return cs.p_relevant / denom if denom > 0 else 0.0


def score_issup(cs: CritiqueScores) -> float:
    denom = cs.p_fully + cs.p_partially + cs.p_no_support
    if denom == 0:
        return 0.0
    return (cs.p_fully / denom) + 0.5 * (cs.p_partially / denom)


def score_isuse(cs: CritiqueScores) -> float:
    probs = {
        1: cs.p_use_1,
        2: cs.p_use_2,
        3: cs.p_use_3,
        4: cs.p_use_4,
        5: cs.p_use_5,
    }
    denom = sum(probs.values())
    if denom == 0:
        return 0.0
    return sum(IsUseToken.WEIGHTS[i] * p for i, p in probs.items()) / denom


def segment_score(log_prob: float, cs: CritiqueScores) -> float:
    return log_prob + W_REL * score_isrel(cs) + W_SUP * score_issup(cs) + W_USE * score_isuse(cs)


def should_retrieve(p_yes: float, p_no: float, delta: float = DELTA) -> bool:
    denom = p_yes + p_no
    return (p_yes / denom) > delta if denom > 0 else False


def setup_models():
    global retriever, gen_tokenizer, gen_model, TOKEN_IDS, MODEL_LOAD_ERROR

    if retriever is not None and gen_tokenizer is not None and gen_model is not None:
        return

    try:
        retriever = LightweightRetriever(device="cpu")
        retriever.index_documents(DOCUMENTS)

        gen_tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)

        if LOAD_IN_4BIT:
            from transformers import BitsAndBytesConfig

            quant_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.float16,
            )
            gen_model = AutoModelForCausalLM.from_pretrained(
                CHECKPOINT,
                device_map="auto",
                quantization_config=quant_config,
                low_cpu_mem_usage=True,
            )
        else:
            if not torch.cuda.is_available():
                raise RuntimeError(
                    "CUDA GPU not available. This notebook-faithful demo expects a GPU "
                    "Space, or you can set LOAD_IN_4BIT=1 for a quantized fallback."
                )
            gen_model = AutoModelForCausalLM.from_pretrained(
                CHECKPOINT,
                torch_dtype=torch.float16,
                device_map="auto",
                low_cpu_mem_usage=True,
            )

        gen_model.eval()
        gen_model.generation_config.do_sample = False
        gen_model.generation_config.temperature = None
        gen_model.generation_config.top_p = None

        def get_tok_id(tok_str: str) -> int:
            ids = gen_tokenizer.encode(tok_str, add_special_tokens=False)
            return ids[0]

        TOKEN_IDS = {tok: get_tok_id(tok) for tok in ALL_REFLECTION_TOKENS}
        MODEL_LOAD_ERROR = None
    except Exception as exc:  # pragma: no cover - startup/runtime guard
        MODEL_LOAD_ERROR = str(exc)
        raise


def _model_device():
    return next(gen_model.parameters()).device


def _encode_input(text: str):
    enc = gen_tokenizer(text, return_tensors="pt", truncation=True, max_length=2048)
    device = _model_device()
    return {k: v.to(device) for k, v in enc.items()}


@torch.no_grad()
def _next_token_probs(prompt: str):
    enc = _encode_input(prompt)
    logits = gen_model(**enc).logits[0, -1, :]
    return torch.softmax(logits, dim=-1)


@torch.no_grad()
def _generate_segment(prompt: str):
    enc = _encode_input(prompt)
    out = gen_model.generate(
        **enc,
        max_new_tokens=MAX_NEW_TOKS,
        do_sample=False,
        return_dict_in_generate=True,
        output_scores=True,
    )

    gen_ids = out.sequences[0, enc["input_ids"].shape[1]:]
    text = gen_tokenizer.decode(gen_ids, skip_special_tokens=False).strip()

    for tok in ALL_REFLECTION_TOKENS:
        if tok in text:
            text = text[:text.index(tok)].strip()

    log_prob = 0.0
    if out.scores:
        for tid, score_t in zip(gen_ids, out.scores):
            log_prob += torch.log_softmax(score_t[0], dim=-1)[tid].item()
        log_prob /= max(len(out.scores), 1)

    return text, log_prob


def _predict_retrieve(prompt: str):
    probs = _next_token_probs(prompt)
    p_yes = probs[TOKEN_IDS[RetrieveToken.YES]].item()
    p_no = probs[TOKEN_IDS[RetrieveToken.NO]].item()
    return should_retrieve(p_yes, p_no, DELTA), p_yes, p_no


def _score_passage_relevance(query: str, passage: Passage) -> float:
    prompt = f"### Instruction:\n{query}\n\n### Input:\n<p>{passage.text}</p>"
    probs = _next_token_probs(prompt)
    p_rel = probs[TOKEN_IDS[IsRelToken.RELEVANT]].item()
    p_irr = probs[TOKEN_IDS[IsRelToken.IRRELEVANT]].item()
    denom = p_rel + p_irr
    return p_rel / denom if denom > 0 else 0.0


def _content_terms(text: str):
    stop = {
        "a", "an", "the", "is", "are", "was", "were", "be", "been", "being",
        "do", "does", "did", "have", "has", "had", "how", "what", "when",
        "where", "why", "which", "who", "whom", "this", "that", "these",
        "those", "and", "or", "but", "for", "with", "into", "from", "about",
        "main", "use", "uses", "using", "used", "number", "version", "date",
        "deployment",
    }
    tokens = re.findall(r"[A-Za-z][A-Za-z0-9_-]+", text.lower())
    return {tok for tok in tokens if tok not in stop and len(tok) > 2}


def _query_passage_coverage(query: str, passage: Passage) -> float:
    q_terms = _content_terms(query)
    if not q_terms:
        return 1.0
    p_terms = _content_terms(f"{passage.title} {passage.text}")
    return len(q_terms & p_terms) / len(q_terms)


def _first_segment_abstention_probe(query: str, context: str = ""):
    passages = retriever.retrieve(query + (" " + context if context else ""), k=K_PASSAGES)
    best_passage = None
    best_isrel = -1.0
    best_coverage = 0.0

    for passage in passages:
        isrel_score = _score_passage_relevance(query, passage)
        coverage = _query_passage_coverage(query, passage)
        if (isrel_score, coverage) > (best_isrel, best_coverage):
            best_passage = passage
            best_isrel = isrel_score
            best_coverage = coverage

    return passages, best_passage, max(best_isrel, 0.0), best_coverage


def _get_critique_scores(base_prompt: str, passage: Passage, gen_text: str):
    cs = CritiqueScores()

    isrel_prompt = f"{base_prompt}<p>{passage.text}</p>"
    probs = _next_token_probs(isrel_prompt)
    cs.p_relevant = probs[TOKEN_IDS[IsRelToken.RELEVANT]].item()
    cs.p_irrelevant = probs[TOKEN_IDS[IsRelToken.IRRELEVANT]].item()

    issup_prompt = f"{isrel_prompt}{IsRelToken.RELEVANT} {gen_text}"
    probs = _next_token_probs(issup_prompt)
    cs.p_fully = probs[TOKEN_IDS[IsSupportToken.FULLY]].item()
    cs.p_partially = probs[TOKEN_IDS[IsSupportToken.PARTIALLY]].item()
    cs.p_no_support = probs[TOKEN_IDS[IsSupportToken.NO]].item()

    isuse_prompt = f"{issup_prompt}{IsSupportToken.FULLY}"
    probs = _next_token_probs(isuse_prompt)
    cs.p_use_5 = probs[TOKEN_IDS[IsUseToken.FIVE]].item()
    cs.p_use_4 = probs[TOKEN_IDS[IsUseToken.FOUR]].item()
    cs.p_use_3 = probs[TOKEN_IDS[IsUseToken.THREE]].item()
    cs.p_use_2 = probs[TOKEN_IDS[IsUseToken.TWO]].item()
    cs.p_use_1 = probs[TOKEN_IDS[IsUseToken.ONE]].item()

    return cs


def _get_isuse_only(prompt: str, text: str):
    cs = CritiqueScores()
    probs = _next_token_probs(f"{prompt}{text}")
    cs.p_use_5 = probs[TOKEN_IDS[IsUseToken.FIVE]].item()
    cs.p_use_4 = probs[TOKEN_IDS[IsUseToken.FOUR]].item()
    cs.p_use_3 = probs[TOKEN_IDS[IsUseToken.THREE]].item()
    cs.p_use_2 = probs[TOKEN_IDS[IsUseToken.TWO]].item()
    cs.p_use_1 = probs[TOKEN_IDS[IsUseToken.ONE]].item()
    return cs


def _argmax_tok(options, probs_map):
    return max(options, key=lambda t: probs_map.get(t, -1.0))


def run_selfrag(query: str, max_segments: int = MAX_SEGMENTS) -> SelfRAGOutput:
    output = SelfRAGOutput(query=query)
    context = ""

    for seg_idx in range(max_segments):
        retrieve_prompt = (
            f"### Instruction:\n{query}\n\n"
            f"### Response:\n{context}"
        )

        if seg_idx == 0:
            _, probe_passage, probe_isrel, probe_coverage = _first_segment_abstention_probe(query, context)
            if (
                probe_passage is None
                or probe_coverage < ABSTENTION_MIN_QUERY_COVERAGE
                or probe_isrel < ABSTENTION_ISREL_THRESHOLD
            ):
                output.abstained = True
                output.segments.append(
                    SegmentResult(
                        text="Insufficient evidence to answer this question.",
                        retrieve_tok=RetrieveToken.NO,
                        passage=probe_passage,
                        isrel=(
                            IsRelToken.RELEVANT
                            if probe_isrel >= 0.5
                            else IsRelToken.IRRELEVANT
                        ),
                        is_sufficient=False,
                        score=-float("inf"),
                        isrel_score=max(probe_isrel, 0.0),
                        query_coverage=probe_coverage,
                    )
                )
                break

        do_retrieve, _, _ = _predict_retrieve(retrieve_prompt)

        if do_retrieve:
            passages = retriever.retrieve(query + " " + context, k=K_PASSAGES)
            candidates = []

            for passage in passages:
                gen_prompt = (
                    f"### Instruction:\n{query}\n\n"
                    f"### Input:\n<p>{passage.text}</p>\n\n"
                    f"### Response:\n{context}"
                )
                text, log_prob = _generate_segment(gen_prompt)
                if not text:
                    continue

                cs = _get_critique_scores(
                    base_prompt=f"### Instruction:\n{query}\n\n### Input:\n",
                    passage=passage,
                    gen_text=text,
                )

                isrel_raw = score_isrel(cs)
                isrel_label = (
                    IsRelToken.RELEVANT if isrel_raw > 0.5 else IsRelToken.IRRELEVANT
                )
                issup_label = (
                    IsSupportToken.FULLY
                    if cs.p_fully >= max(cs.p_partially, cs.p_no_support)
                    else IsSupportToken.PARTIALLY
                    if cs.p_partially >= cs.p_no_support
                    else IsSupportToken.NO
                )
                isuse_label = _argmax_tok(
                    IsUseToken.ALL,
                    {
                        IsUseToken.FIVE: cs.p_use_5,
                        IsUseToken.FOUR: cs.p_use_4,
                        IsUseToken.THREE: cs.p_use_3,
                        IsUseToken.TWO: cs.p_use_2,
                        IsUseToken.ONE: cs.p_use_1,
                    },
                )

                candidates.append(
                    SegmentResult(
                        text=text,
                        retrieve_tok=RetrieveToken.YES,
                        passage=passage,
                        isrel=isrel_label,
                        issup=issup_label,
                        isuse=isuse_label,
                        score=segment_score(log_prob, cs),
                        is_sufficient=(issup_label != IsSupportToken.NO),
                        isrel_score=isrel_raw,
                        query_coverage=_query_passage_coverage(query, passage),
                    )
                )

            if not candidates:
                output.abstained = True
                segment = SegmentResult(
                    text="Insufficient evidence to answer this question.",
                    retrieve_tok=RetrieveToken.YES,
                    is_sufficient=False,
                    score=-float("inf"),
                    isrel_score=0.0,
                )
            else:
                candidates.sort(key=lambda candidate: candidate.score, reverse=True)
                segment = candidates[0]
        else:
            text, log_prob = _generate_segment(retrieve_prompt)
            cs = _get_isuse_only(retrieve_prompt, text)
            isuse_label = _argmax_tok(
                IsUseToken.ALL,
                {
                    IsUseToken.FIVE: cs.p_use_5,
                    IsUseToken.FOUR: cs.p_use_4,
                    IsUseToken.THREE: cs.p_use_3,
                    IsUseToken.TWO: cs.p_use_2,
                    IsUseToken.ONE: cs.p_use_1,
                },
            )
            segment = SegmentResult(
                text=text,
                retrieve_tok=RetrieveToken.NO,
                isuse=isuse_label,
                score=log_prob + W_USE * score_isuse(cs),
                isrel_score=0.0,
            )

        output.segments.append(segment)
        context = (context + " " + segment.text).strip()

        if output.abstained or not segment.text or "</s>" in segment.text:
            break

    return output


def _fix_spaced_acronyms(text: str) -> str:
    pattern = re.compile(r"\b(?:[A-Z]\s+){2,}[A-Z]\b")
    while True:
        updated = pattern.sub(lambda match: match.group(0).replace(" ", ""), text)
        if updated == text:
            return updated
        text = updated


def _detokenize_text(text: str) -> str:
    text = text.replace("\u2581", " ").replace("_", " ")
    text = _fix_spaced_acronyms(text)
    text = re.sub(r"\s+([,.;:!?])", r"\1", text)
    text = re.sub(r"\s*([()\[\]{}])\s*", r" \1 ", text)
    text = re.sub(r"\s*-\s*", "-", text)
    return " ".join(text.split()).strip()


def clean_answer(text: str) -> str:
    for tok in ALL_REFLECTION_TOKENS:
        text = text.replace(tok, "")
    text = re.sub(r"<[^>]+>", "", text)
    text = text.replace("</s>", "").replace("<s>", "")
    text = re.sub(r"\[.*?\]", "", text)
    text = re.sub(r"^\s*\d+[\.\)]\s*", "", text)
    text = re.sub(
        r"^(great question!?|sure!?|here's.*?:)\s*",
        "",
        text,
        flags=re.IGNORECASE,
    )
    return _detokenize_text(text)


def _norm(text: str) -> str:
    text = _detokenize_text(text).lower()
    text = re.sub(r"\b(a|an|the)\b", " ", text)
    text = text.translate(str.maketrans("", "", string.punctuation))
    return " ".join(text.split())


def token_f1(pred: str, gold: str) -> float:
    pred_tokens = _norm(pred).split()
    gold_tokens = _norm(gold).split()
    common = Counter(pred_tokens) & Counter(gold_tokens)
    overlap = sum(common.values())
    if overlap == 0:
        return 0.0
    precision = overlap / len(pred_tokens) if pred_tokens else 0.0
    recall = overlap / len(gold_tokens) if gold_tokens else 0.0
    return 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0.0


def _split_sentences(text: str):
    return [part.strip() for part in re.split(r"(?<=[.!?])\s+", text) if part.strip()]


def _best_source_snippet(question: str, source_text: str) -> str:
    sentences = _split_sentences(source_text)
    if not sentences:
        return source_text.strip()

    q_terms = _content_terms(question)
    ranked = []
    for sent in sentences:
        sent_terms = _content_terms(sent)
        overlap = len(q_terms & sent_terms)
        ranked.append((overlap, len(sent_terms), sent))

    ranked.sort(reverse=True)
    return ranked[0][2].strip()


def _primary_segment(output: SelfRAGOutput) -> Optional[SegmentResult]:
    candidates = [seg for seg in output.segments if seg.passage]
    if not candidates:
        return None
    sufficient = [seg for seg in candidates if seg.is_sufficient]
    chosen = sufficient or candidates
    return max(chosen, key=lambda seg: (seg.score, seg.isrel_score, seg.query_coverage))


def _compact_answer(output: SelfRAGOutput, source_text: str = "") -> str:
    if output.abstained:
        return "I don't have enough evidence in the indexed knowledge base to answer that reliably."

    cleaned = clean_answer(output.answer)
    if not cleaned:
        cleaned = source_text

    if source_text and (len(cleaned) > 260 or "[" in cleaned or "\n" in cleaned):
        cleaned = _best_source_snippet(output.query, source_text)

    sentences = _split_sentences(cleaned)
    compact = " ".join(sentences[:2]) if sentences else cleaned.strip()
    if len(compact) > 260:
        compact = compact[:257].rsplit(" ", 1)[0] + "..."
    return compact


def _result_card(title: str, body: str, tone: str = "normal") -> str:
    palette = {
        "normal": ("#ffffff", "#0f172a", "#dbe4f0"),
        "error": ("#fff4f4", "#9f1239", "#fecdd3"),
        "warn": ("#fff8eb", "#9a3412", "#fed7aa"),
    }
    bg, fg, border = palette[tone]
    return (
        f"<div style='background:{bg};color:{fg};border:1px solid {border};"
        "border-radius:18px;padding:18px 20px;box-shadow:0 10px 30px rgba(15,23,42,0.06)'>"
        f"<div style='font-size:12px;letter-spacing:0.08em;text-transform:uppercase;"
        "opacity:0.65;margin-bottom:8px'>"
        f"{title}</div><div style='font-size:16px;line-height:1.65'>{body}</div></div>"
    )


def render_output(output: SelfRAGOutput) -> str:
    source_segment = _primary_segment(output)
    source_text = source_segment.passage.text if source_segment and source_segment.passage else ""
    final_answer = _compact_answer(output, source_text=source_text)

    source_html = _result_card(
        "Source",
        "No supporting passage was selected."
        if not source_segment or not source_segment.passage
        else (
            f"<div style='font-weight:700;font-size:18px;margin-bottom:8px'>"
            f"{source_segment.passage.title}</div>"
            f"<div style='font-size:13px;opacity:0.7;margin-bottom:10px'>"
            f"{source_segment.passage.id}</div>"
            f"<div>{_best_source_snippet(output.query, source_segment.passage.text)}</div>"
        ),
        tone="warn" if output.abstained else "normal",
    )

    answer_html = _result_card(
        "Final Answer",
        final_answer,
        tone="warn" if output.abstained else "normal",
    )

    question_html = (
        "<div style='margin-bottom:14px;padding:14px 16px;border-radius:16px;"
        "background:linear-gradient(135deg,#eef4ff 0%,#f8fafc 100%);"
        "border:1px solid #d9e6ff;color:#0f172a'>"
        f"<div style='font-size:12px;letter-spacing:0.08em;text-transform:uppercase;"
        "opacity:0.65;margin-bottom:6px'>Question</div>"
        f"<div style='font-size:18px;font-weight:600;line-height:1.5'>{output.query}</div>"
        "</div>"
    )

    return (
        "<div style='font-family:ui-sans-serif,system-ui,sans-serif'>"
        f"{question_html}"
        "<div style='display:grid;grid-template-columns:1fr;gap:14px'>"
        f"{answer_html}{source_html}"
        "</div></div>"
    )


def predict(query: str):
    query = (query or "").strip()
    if not query:
        return _result_card(
            "Ready",
            "Enter a question to see a grounded answer with its source passage.",
        )

    try:
        setup_models()
        output = run_selfrag(query, max_segments=MAX_SEGMENTS)
        return render_output(output)
    except Exception as exc:  # pragma: no cover - runtime guard
        global MODEL_LOAD_ERROR
        MODEL_LOAD_ERROR = str(exc)
        return _result_card(
            "Inference Failed",
            f"{MODEL_LOAD_ERROR}",
            tone="error",
        )


with gr.Blocks(
    css="""
    .gradio-container {
        background:
            radial-gradient(circle at top left, rgba(59,130,246,0.08), transparent 28%),
            radial-gradient(circle at top right, rgba(16,185,129,0.08), transparent 22%),
            linear-gradient(180deg, #f8fbff 0%, #f4f7fb 100%);
    }
    #ask_btn {
        background: linear-gradient(135deg, #0f172a 0%, #1d4ed8 100%) !important;
        color: white !important;
        border: none !important;
    }
    """
) as demo:
    gr.Markdown(
        f"""
        # {TOOL_NAME}
        {TOOL_TAGLINE}
        """
    )

    with gr.Row():
        qbox = gr.Textbox(label="Your Question", scale=5, lines=2)
        ask = gr.Button("Ask", elem_id="ask_btn", scale=1)

    result = gr.HTML(_result_card("Ready", "Enter a question to see a grounded answer with its source passage."))

    gr.Examples(
        examples=EXAMPLES,
        inputs=qbox,
        outputs=result,
        fn=predict,
        cache_examples=False,
    )

    ask.click(predict, inputs=qbox, outputs=result)
    qbox.submit(predict, inputs=qbox, outputs=result)


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