# ── Cell 2: Imports ────────────────────────────────────────────────────────────
import os
import re
import json
import time
import random
import shutil
import unicodedata
import numpy as np
import pandas as pd
from getpass import getpass
from pymilvus import MilvusClient
from groq import Groq
from openai import OpenAI
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi
from sklearn.metrics import roc_auc_score
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from huggingface_hub import hf_hub_download, list_repo_files, HfFileSystem, login
from datasets import load_dataset
import gradio as gr



# ── Cell 3: API Keys ───────────────────────────────────────────────────────────
# Choose your provider: "groq" or "openrouter"


# ── Cell 3: API Keys ───────────────────────────────────────────────────────────

import os
from getpass import getpass
from huggingface_hub import login

# Choose your provider: "groq" or "openrouter"
LLM_PROVIDER = "openrouter"   # change to "groq" if preferred


def get_secret_or_prompt(secret_name, prompt_text=None):
    """
    Try to read secret from Google Colab Secrets.
    If not available, ask user securely using getpass().
    """

    value = None

    # Try Colab Secrets first
    try:
        value = os.getenv(secret_name)
    except Exception:
        value = None

    # Fallback to environment variable
    if not value:
        value = os.environ.get(secret_name)

    # Fallback to manual secure input
    if not value:
        prompt_text = prompt_text or f"Enter {secret_name}: "
        value = getpass(prompt_text)

    return value


# ── HuggingFace Token ─────────────────────────────────────────────────────────

HF_TOKEN = get_secret_or_prompt(
    "HF_TOKEN",
    "Enter HuggingFace Token: "
)

login(token=HF_TOKEN)
os.environ["HF_TOKEN"] = HF_TOKEN

print("✅ HuggingFace token loaded and login completed")


# ── LLM Provider API Key ──────────────────────────────────────────────────────

if LLM_PROVIDER == "groq":

    GROQ_API_KEY = get_secret_or_prompt(
        "GROQ_API_KEY",
        "Enter GROQ API Key: "
    )

    OPENROUTER_API_KEY = None

    os.environ["GROQ_API_KEY"] = GROQ_API_KEY

    print("✅ GROQ API key loaded")

elif LLM_PROVIDER == "openrouter":

    OPENROUTER_API_KEY = get_secret_or_prompt(
        "OPENROUTER_API_KEY",
        "Enter OpenRouter API Key: "
    )

    GROQ_API_KEY = None

    os.environ["OPENROUTER_API_KEY"] = OPENROUTER_API_KEY

    print("✅ OpenRouter API key loaded")

else:
    raise ValueError(f"Unknown LLM_PROVIDER: {LLM_PROVIDER}")



# ── Cell 4: Global configuration ──────────────────────────────────────────────
BUCKET_ID      = "Phani555/IIITH-Cohort26-RAG-Batch37-storage"
BUCKET_PREFIX  = f"hf://buckets/{BUCKET_ID}/milvus_dbs"
MILVUS_DIR     = "/content/milvus_store"
HF_REPO_ID     = "Phani555/IIITH-Cohort26-RAG-Batch37-storage"
HF_REPO_TYPE   = "dataset"
HF_FOLDER      = "ablations"

# ── Model lists per provider ───────────────────────────────────────────────────
GROQ_LLM_CHOICES = [
    "llama-3.1-8b-instant",
    "gemma2-9b-it",
    "llama-3.3-70b-versatile",
    "mixtral-8x7b-32768",
    "qwen/qwen3-32b",
    "qwen-qwq-32b",
    "deepseek-r1-distill-llama-70b",
]
OPENROUTER_LLM_CHOICES = [
    "meta-llama/llama-3.1-8b-instruct",
    "meta-llama/llama-3.3-70b-instruct",
    "openai/gpt-oss-20b",
    "openai/gpt-oss-120b",
    "qwen/qwen3-32b",
    "deepseek/deepseek-r1",
    "moonshotai/kimi-k2-instruct",
    "openai/gpt-oss-safeguard-20b",
]
LLM_CHOICES = OPENROUTER_LLM_CHOICES if LLM_PROVIDER == "openrouter" else GROQ_LLM_CHOICES

EMBEDDING_CHOICES = ["bge_small", "llm_embedder"]
EMBED_MODELS = {
    "bge_small":    "BAAI/bge-small-en-v1.5",
    "llm_embedder": "BAAI/llm-embedder",
}

# ── Runtime globals ────────────────────────────────────────────────────────────
MODEL_NAME       = LLM_CHOICES[0]
MODEL_NAME_BIG   = LLM_CHOICES[4]
EMBEDDING_TYPE   = "llm_embedder"

ENABLE_HYBRID        = False
ENABLE_HYDE          = False
ENABLE_RERANKING     = False
RERANKER_TYPE        = "monot5"
PROMPT_STRATEGY      = "short"
ENABLE_REPACKING     = False
REPACK_STRATEGY      = "sides"
ENABLE_SUMMARIZATION = False
SUMMARIZATION_TYPE   = "recomp"

ENABLE_QUERY_REWRITING     = False
ENABLE_QUERY_DECOMPOSITION = False
ENABLE_QUERY_CLASSIFICATION= False
MAX_SUBQUERIES             = 3
QUERY_REWRITE_MODEL        = None   # falls back to MODEL_NAME
QUERY_DECOMPOSE_MODEL      = None

RETRIEVE_TOP_K         = 10
RERANK_TOP_K           = 3
HYBRID_ALPHA           = 0.5
MONOT5_MODEL           = "castorini/monot5-base-msmarco-10k"
TILDE_MODEL            = "BAAI/bge-reranker-base"

RECOMP_TOP_K_SENTS     = 8
RECOMP_MIN_SCORE       = 0.10
RECOMP_GROUNDING_BOOST = 0.15
RECOMP_MIN_KEEP_RATIO  = 0.60
RECOMP_KEEP_CRITICAL   = True
LLMLINGUA_RATE         = 0.5

milvus_clients       = {}
bm25_indexes         = {}
embed_model          = None
llm_client           = None
monot5_reranker      = None
tilde_reranker       = None
llmlingua_compressor = None
ragbench_by_domain   = {}

DOMAIN_NAMES = [
    "Bio_Medical",
    "General_Knowledge",
    "Customer_Support",
    "Finance",
    "Legal_Contracts",
]

GROUNDING_PATTERNS = [
    r"\b(?:must|should|shall|cannot|can't|never|always|only|except|unless|required|recommended)\b",
    r"\b(?:warning|caution|note|important|attention)\b",
    r"\b(?:do not|don't|does not|did not|not allowed|not recommended|never)\b",
    r"\b\d+(?:\.\d+)?\s*(?:%|percent|seconds?|minutes?|hours?|days?|weeks?|months?|years?)\b",
    r"\b\d+(?:\.\d+)?\s*(?:GB|MB|KB|TB|kg|g|mg|mm|cm|m|km|degrees?|°C|°F)\b",
    r"[$€£¥]\s*\d+(?:,\d{3})*(?:\.\d+)?",
    r"\b\d+(?:,\d{3})*(?:\.\d+)?\s*(?:dollars?|rupees?|crores?|lakhs?|million|billion)\b",
    r"\b(?:19|20)\d{2}\b",
    r"\b\d{2,}\b",
    r"\b[A-Z]{2,}[-_]?\d+[A-Z0-9-]*\b",
    r"\b[A-Z0-9]{3,}[-_][A-Z0-9]{2,}\b",
    r"\b[A-Z]{3,}\b",
]
GROUNDING_REGEX = re.compile("|".join(GROUNDING_PATTERNS), re.IGNORECASE)

print(f"Config loaded.  Provider: {LLM_PROVIDER}  |  Models: {len(LLM_CHOICES)}")


# ── Cell 5: Pipeline functions (from Task1 Final) ─────────────────────────────

# ── Utilities ──────────────────────────────────────────────────────────────────

def _safe_message_content(response):
    """Extract assistant content from OpenAI/OpenRouter/Groq response safely."""
    try:
        msg = response.choices[0].message
        content = getattr(msg, "content", None)
        return str(content).strip() if content else ""
    except Exception:
        return ""

def _sanitize(text):
    """
    Normalise to NFC then replace non-ASCII characters that some HTTP
    transports reject with their closest ASCII equivalent (or a space).
    Handles bullets, arrows, curly quotes, em-dashes, etc.
    """
    if not text:
        return text
    text = unicodedata.normalize("NFC", str(text))
    return text.encode("ascii", errors="replace").decode("ascii")

def get_domain(dataset):
    if dataset in ("covidqa", "pubmedqa"):           return "Bio_Medical"
    elif dataset in ("expertqa","hagrid","hotpotqa","msmarco"): return "General_Knowledge"
    elif dataset in ("delucionqa","emanual","techqa"): return "Customer_Support"
    elif dataset in ("finqa","tatqa"):                return "Finance"
    else:                                             return "Legal_Contracts"

def get_db_path(domain_name, embedding_type=None):
    embedding_type = embedding_type or EMBEDDING_TYPE
    if embedding_type == "bge_small":
        return os.path.join(MILVUS_DIR, f"{domain_name}.db")
    elif embedding_type == "llm_embedder":
        llm_dir = os.path.join(MILVUS_DIR, "llm_embedder")
        os.makedirs(llm_dir, exist_ok=True)
        return os.path.join(llm_dir, f"{domain_name}.db")
    raise ValueError(f"Unknown EMBEDDING_TYPE={embedding_type}")

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

def _tokenize(text):
    return re.findall(r'\w+', str(text).lower())

def _normalize(scores):
    arr = np.array(scores, dtype=float)
    if len(arr) == 0 or arr.max() == arr.min():
        return np.zeros_like(arr)
    return (arr - arr.min()) / (arr.max() - arr.min())

def _count_grounding_signals(sentence):
    return len(GROUNDING_REGEX.findall(str(sentence)))

def _is_critical_sentence(sentence):
    pat = re.compile(
        r"\b(?:warning|caution|important|must|must not|cannot|can't|do not|don't|never|only|except|unless|required)\b",
        re.IGNORECASE)
    return bool(pat.search(str(sentence)))


# ── LLM client factory ─────────────────────────────────────────────────────────

def get_llm_client():
    if LLM_PROVIDER == "groq":
        return Groq(api_key=GROQ_API_KEY)
    elif LLM_PROVIDER == "openrouter":
        return OpenAI(api_key=OPENROUTER_API_KEY, base_url="https://openrouter.ai/api/v1")
    raise ValueError(f"Unknown LLM_PROVIDER: {LLM_PROVIDER}")


# ── Query Classification ───────────────────────────────────────────────────────

def classify_query(query, domain_name=None):
    """Returns 'RAG' or 'LLM'. Benchmark domains always force 'RAG'."""
    if not ENABLE_QUERY_CLASSIFICATION:
        return "RAG"
    rag_domains = {"Bio_Medical","General_Knowledge","Customer_Support","Finance","Legal_Contracts"}
    if domain_name in rag_domains:
        return "RAG"
    llm_keywords = [
        "who is","what is","when was","where is","define","explain",
        "tell me about","what are","why is","how does","what does",
    ]
    if any(kw in str(query).lower() for kw in llm_keywords):
        return "LLM"
    return "RAG"


# ── Query Rewriting ────────────────────────────────────────────────────────────

def rewrite_query(query, domain_name, llm_client, model_name=None):
    """
    Rewrites the query to improve retrieval quality.
    Returns the original query on any failure or if disabled.
    """
    if not ENABLE_QUERY_REWRITING:
        return query

    model = model_name or QUERY_REWRITE_MODEL or MODEL_NAME

    prompt = f"""Rewrite the question to improve document retrieval. Apply only when needed.

Domain: {domain_name}

Rules:
- Preserve the complete meaning and question form.
- Fix grammar and resolve ambiguity.
- Expand abbreviations if their full form aids retrieval.
- Preserve all names, product names, dates, numbers, legal, biomedical, financial and technical terms.
- Do not answer the question.
- Do not add unsupported information.
- If the query is already clear and specific, return it unchanged.
- Return ONLY the rewritten query — no explanation, no prefix, no quotes.

Original question:
{query}""".strip()

    try:
        resp = llm_client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You rewrite questions to improve semantic document retrieval. Return only the rewritten question."},
                {"role": "user",   "content": _sanitize(prompt)},
            ],
            temperature=0.0,
            max_tokens=150,
        )
        rewritten = _safe_message_content(resp).strip()
        # Reject if empty, too long, or suspiciously different length
        if not rewritten or len(rewritten) > 600:
            return query
        return rewritten
    except Exception as e:
        print(f"Query rewriting failed: {e}")
        return query


# ── Query Decomposition helpers ────────────────────────────────────────────────

def _clean_subquery_text(text):
    if text is None: return ""
    text = str(text).strip()
    text = text.replace("```json","").replace("```","").strip()
    text = text.rstrip(",").strip('"').strip("'").strip()
    text = re.sub(r"^\s*[-*]\s*", "", text)
    text = re.sub(r"^\s*\d+[\).\:\-]\s*", "", text)
    return text.strip()

def _looks_like_explanation_line(text):
    if not text: return True
    text_l = text.lower().strip()
    bad_prefixes = ["here are","here is","decomposed","search queries","the decomposed",
                    "queries:","subqueries:","output:","json:","answer:"]
    if any(text_l.startswith(p) for p in bad_prefixes): return True
    if text_l in {"queries","subqueries","search queries","decomposed search queries"}: return True
    return False

def _parse_json_object_line(line):
    line = _clean_subquery_text(line)
    if not line: return None
    try:
        obj = json.loads(line)
        if isinstance(obj, dict):
            for key in ["query","question","subquery","search_query"]:
                if key in obj and str(obj[key]).strip():
                    return str(obj[key]).strip()
        if isinstance(obj, str): return obj.strip()
    except Exception: pass
    m = re.search(r'"(?:query|question|subquery|search_query)"\s*:\s*"([^"]+)"', line)
    if m: return m.group(1).strip()
    return None

def _split_multi_question_locally(query, max_subqueries=None):
    max_subqueries = max_subqueries or MAX_SUBQUERIES
    query = str(query).strip()
    parts = [p.strip() for p in re.split(r"\?\s*", query) if p.strip()]
    if len(parts) <= 1: return None
    return [(p + "?" if not p.endswith("?") else p) for p in parts[:max_subqueries]]

def _parse_subqueries(raw_text, original_query, max_subqueries=None):
    """Robustly parse subqueries from any LLM output format."""
    max_subqueries = max_subqueries or MAX_SUBQUERIES
    if not raw_text: return [original_query]
    text = str(raw_text).strip().replace("```json","").replace("```","").strip()

    # Try full JSON first
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            subs = []
            for item in parsed:
                if isinstance(item, dict):
                    for key in ["query","question","subquery","search_query"]:
                        if key in item and str(item[key]).strip():
                            subs.append(str(item[key]).strip()); break
                elif isinstance(item, str):
                    subs.append(item.strip())
            subs = [_clean_subquery_text(q) for q in subs if _clean_subquery_text(q)]
            return subs[:max_subqueries] or [original_query]
        elif isinstance(parsed, dict):
            raw_list = (parsed.get("subqueries") or parsed.get("queries") or
                        parsed.get("questions") or parsed.get("search_queries") or [])
            if isinstance(raw_list, list):
                subs = [_clean_subquery_text(q) for q in raw_list if _clean_subquery_text(q)]
                return subs[:max_subqueries] or [original_query]
    except Exception: pass

    # Line-by-line fallback
    subqueries = []
    for raw_line in text.splitlines():
        line = _clean_subquery_text(raw_line)
        if not line or _looks_like_explanation_line(line): continue
        obj_q = _parse_json_object_line(line)
        if obj_q:
            obj_q = _clean_subquery_text(obj_q)
            if obj_q and not _looks_like_explanation_line(obj_q):
                subqueries.append(obj_q)
            continue
        if line.startswith("{") or line.endswith("}") or line in {"[","]","{","}"}: continue
        subqueries.append(line)

    deduped = []
    for q in subqueries:
        q = _clean_subquery_text(q)
        if q and q not in deduped: deduped.append(q)
    return deduped[:max_subqueries] or [original_query]


# ── Query Decomposition ────────────────────────────────────────────────────────

def decompose_query(query, llm_client, domain=None, model=None, max_subqueries=None):
    """
    Decompose a complex query into focused retrieval subqueries.
    Handles multi-question input locally before calling the LLM.
    Uses robust parsing to handle messy LLM output.
    """
    if not ENABLE_QUERY_DECOMPOSITION:
        return [query]

    max_subqueries = max_subqueries or MAX_SUBQUERIES

    # Handle obvious multi-question input locally (no LLM call needed)
    local_split = _split_multi_question_locally(query, max_subqueries)
    if local_split:
        return local_split

    model = model or QUERY_DECOMPOSE_MODEL or MODEL_NAME
    if not model:
        return [query]

    prompt = f"""Decompose the question into at most {max_subqueries} retrieval-focused search queries.

Return ONLY a valid JSON list of strings. No explanations. No markdown. No object notation.

Valid output example:
["What caused the 2008 financial crisis?", "Which banks failed in 2008?"]

Rules:
- If the question is already simple, return a JSON list with the original question only.
- Do not answer the question.
- Preserve all names, dates, numbers, legal, biomedical, financial and technical terms.

Domain: {domain}

Question: {query}""".strip()

    try:
        resp = llm_client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You decompose complex questions into retrieval subqueries and return only a JSON list of strings."},
                {"role": "user",   "content": _sanitize(prompt)},
            ],
            temperature=0.0,
            max_tokens=300,
        )
        raw = _safe_message_content(resp)
        return _parse_subqueries(raw, original_query=query, max_subqueries=max_subqueries)
    except Exception as e:
        print(f"Query decomposition failed: {e}")
        return [query]


# ── Reranking ──────────────────────────────────────────────────────────────────

class MonoT5Reranker:
    def __init__(self, model_name=None):
        model_name = model_name or MONOT5_MODEL
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device); self.model.eval()
        self.true_id  = self.tokenizer.convert_tokens_to_ids("▁true")
        self.false_id = self.tokenizer.convert_tokens_to_ids("▁false")
        if not self.true_id  or self.true_id  < 0: self.true_id  = self.tokenizer.encode("true",  add_special_tokens=False)[0]
        if not self.false_id or self.false_id < 0: self.false_id = self.tokenizer.encode("false", add_special_tokens=False)[0]
        print(f"MonoT5 loaded: {model_name} on {self.device}")

    def score(self, query, document):
        text = f"Query: {query} Document: {document} Relevant:"
        enc  = self.tokenizer(text, return_tensors="pt", max_length=512, truncation=True).to(self.device)
        with torch.no_grad():
            out = self.model.generate(**enc, max_new_tokens=1, return_dict_in_generate=True, output_scores=True)
        logits = out.scores[0][0]
        probs  = torch.softmax(torch.stack([logits[self.false_id], logits[self.true_id]]), dim=0)
        return float(probs[1].item())

    def compute_scores(self, query, texts):
        return np.array([self.score(query, t) for t in texts], dtype=float)

def get_monot5_reranker():
    global monot5_reranker
    if monot5_reranker is None:
        monot5_reranker = MonoT5Reranker(MONOT5_MODEL)
    return monot5_reranker

def get_tilde_reranker():
    global tilde_reranker
    if tilde_reranker is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        tilde_reranker = CrossEncoder(TILDE_MODEL, device=device)
        print(f"TILDE reranker loaded: {TILDE_MODEL} on {device}")
    return tilde_reranker

def rerank_documents(query, documents, top_k=3):
    if not documents: return []
    texts  = [d["text"] if isinstance(d, dict) else d for d in documents]
    rtype  = RERANKER_TYPE.lower().strip()
    scores = get_monot5_reranker().compute_scores(query, texts) if rtype == "monot5" \
             else np.asarray(get_tilde_reranker().predict([(query, t) for t in texts], show_progress_bar=False), dtype=float).reshape(-1)
    ranked_idx = np.argsort(scores)[::-1][:top_k]
    reranked = []
    for i in ranked_idx:
        item = dict(documents[i]) if isinstance(documents[i], dict) else {"text": documents[i]}
        item["base_score"]    = item.get("score")
        item["score"]         = float(scores[i])
        item["rerank_score"]  = float(scores[i])
        item["reranker_type"] = rtype
        reranked.append(item)
    return reranked


# ── BM25 ───────────────────────────────────────────────────────────────────────

def build_bm25_index(domain_name, clients):
    client = clients[domain_name]
    col    = domain_name.lower()
    try:
        if client.get_load_state(col) != "Loaded": client.load_collection(col)
    except Exception: pass
    try:
        n = int(client.get_collection_stats(col).get("row_count", 0))
    except Exception: return
    if n == 0: return
    rows  = client.query(collection_name=col, filter="", limit=n, output_fields=["text"])
    texts = [r["text"] for r in rows if r.get("text")]
    if not texts: return
    bm25_indexes[domain_name] = {"bm25": BM25Okapi([_tokenize(t) for t in texts]), "texts": texts}
    print(f"  BM25 built: {len(texts)} docs  [{domain_name}]")

def build_all_bm25_indexes(clients):
    bm25_indexes.clear()
    for d in clients: build_bm25_index(d, clients)


# ── HyDE ───────────────────────────────────────────────────────────────────────

def generate_hyde(query, llm_client, model_name=None):
    model  = model_name or MODEL_NAME
    prompt = f"Write a brief factual passage answering this question (under 4 sentences).\nQuestion: {query}\nPassage:"
    try:
        resp = llm_client.chat.completions.create(
            model=model,
            messages=[{"role": "system", "content": "You write hypothetical answer passages for retrieval."},
                      {"role": "user",   "content": _sanitize(prompt)}],
            temperature=0.2, max_tokens=300,
        )
        return _safe_message_content(resp)
    except Exception as e:
        print(f"HyDE failed: {e}"); return ""


# ── Hybrid search ──────────────────────────────────────────────────────────────

def hybrid_search(query, domain_name, em, top_k=20, alpha=0.5):
    client = milvus_clients[domain_name]
    col    = domain_name.lower()
    try:
        if client.get_load_state(col) != "Loaded": client.load_collection(col)
    except Exception: pass
    q_emb = em.encode([query], normalize_embeddings=True).astype("float32")
    hits  = client.search(collection_name=col, data=q_emb.tolist(), limit=top_k,
                          output_fields=["text"], search_params={"metric_type":"IP","params":{}})
    dense = {h.entity.get("text",""): float(h.distance) for h in hits[0] if h.entity.get("text","")}
    bm25_obj = bm25_indexes.get(domain_name)
    if not bm25_obj:
        return [{"text":t,"score":s,"dense_score":s,"bm25_score":0.0} for t,s in sorted(dense.items(),key=lambda x:-x[1])[:top_k]]
    bm25_scores = bm25_obj["bm25"].get_scores(_tokenize(query))
    top_idx     = np.argsort(bm25_scores)[::-1][:top_k]
    sparse      = {bm25_obj["texts"][i]: float(bm25_scores[i]) for i in top_idx}
    all_texts   = sorted(set(dense) | set(sparse))
    d_vals = [dense.get(t,0.0) for t in all_texts]
    b_vals = [sparse.get(t,0.0) for t in all_texts]
    d_norm, b_norm = _normalize(d_vals), _normalize(b_vals)
    combined = [{"text":t, "score":float(alpha*d_norm[i]+(1-alpha)*b_norm[i]),
                 "dense_score":float(d_vals[i]), "bm25_score":float(b_vals[i])}
                for i,t in enumerate(all_texts)]
    combined.sort(key=lambda x: -x["score"])
    return combined[:top_k]


# ── Repacking ──────────────────────────────────────────────────────────────────

def repack_documents(docs, strategy="sides"):
    if not docs: return []
    if strategy == "forward": return docs
    if strategy == "reverse": return docs[::-1]
    if strategy == "sides":
        n, result, left, right = len(docs), [None]*len(docs), 0, len(docs)-1
        for i, doc in enumerate(docs):
            if i % 2 == 0: result[left]  = doc; left  += 1
            else:           result[right] = doc; right -= 1
        return result
    raise ValueError(f"Unknown REPACK_STRATEGY: {strategy}")


# ── Summarization ──────────────────────────────────────────────────────────────

def recomp_summarize(query, docs, em, top_k=6, min_score=0.0,
                     grounding_boost=0.15, min_keep_ratio=0.50, keep_critical=True):
    texts     = [d.get("text","") if isinstance(d,dict) else d for d in docs]
    sentences = [s.strip() for doc in texts for s in split_into_sentences(doc) if s.strip()]
    if not sentences: return ""
    q_emb   = em.encode([query], normalize_embeddings=True)
    s_emb   = em.encode(sentences, normalize_embeddings=True)
    scores  = (q_emb @ s_emb.T).flatten() + np.array([_count_grounding_signals(s)*grounding_boost for s in sentences])
    crits   = {i for i,s in enumerate(sentences) if keep_critical and _is_critical_sentence(s)}
    valid   = np.where(scores >= min_score)[0]
    if len(valid) == 0: valid = np.array([int(np.argmax(scores))])
    keep    = min(max(top_k, int(np.ceil(len(sentences)*min_keep_ratio))), len(sentences))
    chosen  = sorted(set(list(valid[np.argsort(scores[valid])[::-1][:keep]])) | crits)
    return " ".join(sentences[i] for i in chosen)

def _get_llmlingua():
    global llmlingua_compressor
    if llmlingua_compressor is None:
        from llmlingua import PromptCompressor
        device = "cuda" if torch.cuda.is_available() else "cpu"
        llmlingua_compressor = PromptCompressor(
            model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",
            use_llmlingua2=True, device_map=device)
        print(f"LLMLingua loaded on {device}")
    return llmlingua_compressor

def llmlingua_compress(query, docs, rate=0.5):
    texts = [d.get("text","") if isinstance(d,dict) else d for d in docs]
    comp  = _get_llmlingua()
    parts = []
    for t in texts:
        if not t or not t.strip(): continue
        try:    parts.append(comp.compress_prompt(t, question=query, rate=rate)["compressed_prompt"])
        except Exception as e: print(f"LLMLingua chunk failed: {e}"); parts.append(t)
    return "\n\n".join(parts)

def summarize_docs(query, docs, em=None, llm_client=None):
    if not ENABLE_SUMMARIZATION:
        return [d.get("text","") if isinstance(d,dict) else d for d in docs]
    em = em or embed_model
    if SUMMARIZATION_TYPE == "recomp":
        s = recomp_summarize(query, docs, em, RECOMP_TOP_K_SENTS, RECOMP_MIN_SCORE,
                             RECOMP_GROUNDING_BOOST, RECOMP_MIN_KEEP_RATIO, RECOMP_KEEP_CRITICAL)
        return [s] if s else []
    elif SUMMARIZATION_TYPE == "longllmlingua":
        c = llmlingua_compress(query, docs, LLMLINGUA_RATE)
        return [c] if c else []
    raise ValueError(f"Unknown SUMMARIZATION_TYPE: {SUMMARIZATION_TYPE}")


# ── Main retrieve ──────────────────────────────────────────────────────────────

def retrieve(query, domain_name, embed_model=None, llm_client=None, top_k=3, rewritten_query=None):
    if domain_name not in milvus_clients:
        raise ValueError(f"Domain '{domain_name}' not loaded.")
    em  = embed_model or globals().get("embed_model")
    llm = llm_client  or globals().get("llm_client")
    if em is None: raise ValueError("embed_model is required")
    fetch_k = RETRIEVE_TOP_K if (ENABLE_HYBRID or ENABLE_RERANKING or ENABLE_SUMMARIZATION or ENABLE_REPACKING) else top_k
    eff_q = rewritten_query or query
    search_q = eff_q
    if ENABLE_HYDE and llm:
        hyde = generate_hyde(eff_q, llm)
        if hyde: search_q = f"{eff_q} {hyde}"
    if ENABLE_HYBRID:
        retrieved = hybrid_search(search_q, domain_name, em, top_k=fetch_k, alpha=HYBRID_ALPHA)
    else:
        client = milvus_clients[domain_name]; col = domain_name.lower()
        try:
            if client.get_load_state(col) != "Loaded": client.load_collection(col)
        except Exception: pass
        q_emb = em.encode([search_q], normalize_embeddings=True).astype("float32")
        hits  = client.search(collection_name=col, data=q_emb.tolist(), limit=fetch_k,
                              output_fields=["text"], search_params={"metric_type":"IP","params":{}})
        seen, retrieved = set(), []
        for h in hits[0]:
            t = h.entity.get("text","")
            if t and t not in seen: retrieved.append({"text":t,"score":float(h.distance)}); seen.add(t)
    if not retrieved: return []
    if ENABLE_RERANKING: retrieved = rerank_documents(eff_q, retrieved, top_k)
    retrieved = retrieved[:top_k]
    if ENABLE_SUMMARIZATION:
        summarized = summarize_docs(eff_q, retrieved, em=em, llm_client=llm)
        avg = float(np.mean([d.get("score",0) for d in retrieved])) if retrieved else 1.0
        retrieved = [{"text":s,"score":avg,"summarized":True,"summary_type":SUMMARIZATION_TYPE} for s in summarized]
    if ENABLE_REPACKING: retrieved = repack_documents(retrieved, REPACK_STRATEGY)
    return retrieved


# ── Prompt / generation ────────────────────────────────────────────────────────

def _build_prompt(context, question, strategy="short"):
    context  = _sanitize(context)
    question = _sanitize(question)
    if strategy == "short":
        return f"Answer the question using the provided context.\n\nContext:\n{context}\n\nQuestion:\n{question}".strip()
    elif strategy == "long":
        return (
            "You are a chatbot providing answers to user queries. Use the context documents to answer the question.\n"
            'If the documents do not provide enough information, say "The documents are missing some of the information required to answer the question."\n'
            "Do not use external knowledge. Do not make up an answer.\n\n"
            f"Context Documents:\n{context}\n\nQuestion: {question}"
        ).strip()
    elif strategy == "long_cot":
        return (
            "You are a chatbot providing answers to user queries. Use the context documents to answer the question.\n"
            'If the documents do not provide enough information, say "The documents are missing some of the information required to answer the question."\n'
            "Do not use external knowledge. Do not make up an answer.\n"
            "Think step by step and quote documents when necessary.\n\n"
            f"Context Documents:\n{context}\n\nQuestion: {question}"
        ).strip()
    raise ValueError(f"Unknown PROMPT_STRATEGY: {strategy}")

def ask_rag(context, question, llm_client, strategy=None):
    strategy = strategy or PROMPT_STRATEGY
    resp = llm_client.chat.completions.create(
        model=MODEL_NAME,
        messages=[{"role":"system","content":"You are a helpful RAG assistant"},
                  {"role":"user","content":_build_prompt(context, question, strategy)}],
        temperature=0.3,
    )
    return _safe_message_content(resp)


print("Pipeline functions defined.")


# ── Cell 8: Download Milvus DBs from HF Bucket - Clean V2 ─────────────────────

import os
import shutil
import tempfile
from huggingface_hub import HfFileSystem

MILVUS_DIR = "/content/milvus_store/milvus_dbs"


def get_llm_embedder_folder(index_version=None):
    index_version = index_version or INDEX_VERSION

    return (
        "llm_embedder"
        if index_version == "default"
        else f"llm_embedder_{index_version}"
    )


def get_db_path(domain_name, embedding_type=None, index_version=None):
    embedding_type = embedding_type or EMBEDDING_TYPE
    index_version = index_version or INDEX_VERSION

    if embedding_type == "bge_small":
        db_dir = MILVUS_DIR

    elif embedding_type == "llm_embedder":
        db_dir = os.path.join(
            MILVUS_DIR,
            get_llm_embedder_folder(index_version)
        )

    else:
        raise ValueError(f"Unknown EMBEDDING_TYPE={embedding_type}")

    os.makedirs(db_dir, exist_ok=True)

    return os.path.join(db_dir, f"{domain_name}.db")


def get_remote_bucket_dir(embedding_type=None, index_version=None):
    embedding_type = embedding_type or EMBEDDING_TYPE
    index_version = index_version or INDEX_VERSION

    if "BUCKET_PREFIX" not in globals():
        raise ValueError("BUCKET_PREFIX is not defined")

    bucket_prefix = str(BUCKET_PREFIX).rstrip("/")

    if embedding_type == "bge_small":
        return bucket_prefix

    elif embedding_type == "llm_embedder":
        return f"{bucket_prefix}/{get_llm_embedder_folder(index_version)}"

    else:
        raise ValueError(f"Unknown EMBEDDING_TYPE={embedding_type}")


def get_local_download_dir(embedding_type=None, index_version=None):
    embedding_type = embedding_type or EMBEDDING_TYPE
    index_version = index_version or INDEX_VERSION

    if embedding_type == "bge_small":
        return MILVUS_DIR

    elif embedding_type == "llm_embedder":
        return os.path.join(
            MILVUS_DIR,
            get_llm_embedder_folder(index_version)
        )

    else:
        raise ValueError(f"Unknown EMBEDDING_TYPE={embedding_type}")


def get_path_size_mb(path):
    if os.path.isfile(path):
        return os.path.getsize(path) / 1e6

    total = 0

    for root, _, files in os.walk(path):
        for file in files:
            fp = os.path.join(root, file)

            if os.path.exists(fp):
                total += os.path.getsize(fp)

    return total / 1e6


def find_db_dirs_or_files(root_dir, domain_names=None):
    """
    Find both:
      1. Directories ending with .db
      2. Files ending with .db

    Milvus Lite DBs are usually directories ending with .db.
    """

    expected_names = None

    if domain_names is not None:
        expected_names = {f"{d}.db" for d in domain_names}

    found = []

    for root, dirs, files in os.walk(root_dir):
        for d in dirs:
            if not d.endswith(".db"):
                continue

            if expected_names is not None and d not in expected_names:
                continue

            found.append(os.path.join(root, d))

        for f in files:
            if not f.endswith(".db"):
                continue

            if expected_names is not None and f not in expected_names:
                continue

            found.append(os.path.join(root, f))

    return sorted(set(found))


def copy_db_object(src, dst):
    """
    Copy .db directory or .db file.
    """

    if os.path.isdir(src):
        if os.path.exists(dst):
            shutil.rmtree(dst)

        shutil.copytree(src, dst)

    elif os.path.isfile(src):
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        shutil.copy2(src, dst)

    else:
        raise FileNotFoundError(f"Source DB object not found: {src}")


def print_staging_debug(staging_root, max_items=80):
    print("\nStaging debug tree sample:")

    shown = 0

    for root, dirs, files in os.walk(staging_root):
        for d in dirs:
            print("  DIR :", os.path.join(root, d))
            shown += 1

            if shown >= max_items:
                return

        for f in files:
            print("  FILE:", os.path.join(root, f))
            shown += 1

            if shown >= max_items:
                return


def download_milvus_dbs_from_bucket_v2(
    embedding_type=None,
    index_version=None,
    domain_names=None,
    force_download=False,
    debug=True,
):
    """
    Correct downloader for HF bucket Milvus Lite DBs.

    Important:
      Milvus Lite .db is usually a DIRECTORY, not a single file.

    This function:
      1. Lists remote bucket folder.
      2. Downloads recursively into staging.
      3. Detects .db directories/files.
      4. Copies each .db object into final expected local path.
    """

    embedding_type = embedding_type or EMBEDDING_TYPE
    index_version = index_version or INDEX_VERSION

    remote_dir = get_remote_bucket_dir(
        embedding_type=embedding_type,
        index_version=index_version,
    ).rstrip("/")

    final_local_dir = get_local_download_dir(
        embedding_type=embedding_type,
        index_version=index_version,
    )

    print("=" * 100)
    print("DOWNLOAD MILVUS DBS FROM HF BUCKET V2")
    print("=" * 100)
    print(f"Embedding Type : {embedding_type}")
    print(f"Index Version  : {index_version}")
    print(f"Remote Dir     : {remote_dir}")
    print(f"Final Local Dir: {final_local_dir}")

    fs_token = globals().get("HF_TOKEN", None)
    fs = HfFileSystem(token=fs_token) if fs_token else HfFileSystem()

    try:
        remote_items = fs.ls(remote_dir, detail=False)
    except Exception as e:
        print("\nCould not list remote dir:")
        print(f"  {remote_dir}")
        print(f"Error: {e}")
        return {}

    if debug:
        print("\nRemote listing check:")
        print(f"Found {len(remote_items)} remote items under:")
        print(f"  {remote_dir}")

        for item in remote_items[:50]:
            print(f"  - {item}")

    remote_db_names = [
        os.path.basename(str(item))
        for item in remote_items
        if os.path.basename(str(item)).endswith(".db")
    ]

    if domain_names is not None:
        expected_names = {f"{d}.db" for d in domain_names}
        remote_db_names = [
            name for name in remote_db_names
            if name in expected_names
        ]

    remote_db_names = sorted(set(remote_db_names))

    print(f"\nRemote .db entries detected: {len(remote_db_names)}")

    for name in remote_db_names:
        print(f"  - {name}")

    if not remote_db_names:
        print("\nWARNING: No remote .db entries detected.")
        return {}

    if force_download and os.path.exists(final_local_dir):
        shutil.rmtree(final_local_dir)

    os.makedirs(final_local_dir, exist_ok=True)

    staging_root = tempfile.mkdtemp(prefix="hf_milvus_download_v2_")

    print("\nDownloading recursively using fs.get()...")
    print(f"Remote : {remote_dir}")
    print(f"Staging: {staging_root}")

    try:
        fs.get(
            remote_dir,
            staging_root,
            recursive=True,
        )
    except Exception as e:
        print("\nRecursive download failed.")
        print(f"Error: {e}")
        shutil.rmtree(staging_root, ignore_errors=True)
        return {}

    staging_db_paths = find_db_dirs_or_files(
        root_dir=staging_root,
        domain_names=domain_names,
    )

    print(f"\nLocal .db paths found in staging: {len(staging_db_paths)}")

    for p in staging_db_paths:
        kind = "DIR" if os.path.isdir(p) else "FILE"
        size_mb = get_path_size_mb(p)
        print(f"  - [{kind}] {p} ({size_mb:.2f} MB)")

    if not staging_db_paths:
        print("\nWARNING: No .db directory/file was found in staging.")
        print_staging_debug(staging_root, max_items=100)
        shutil.rmtree(staging_root, ignore_errors=True)
        return {}

    report = {}

    print("\nCopying DB objects into final local directory...")

    for src in staging_db_paths:
        db_name = os.path.basename(src)
        dst = os.path.join(final_local_dir, db_name)

        try:
            copy_db_object(src, dst)

            size_mb = get_path_size_mb(dst)
            kind = "DIR" if os.path.isdir(dst) else "FILE"

            print(f"  OK [{kind}] {db_name} -> {dst} ({size_mb:.2f} MB)")

            report[db_name] = {
                "status": "downloaded",
                "kind": kind,
                "source": src,
                "local_path": dst,
                "size_mb": size_mb,
            }

        except Exception as e:
            print(f"  FAILED {db_name}: {e}")

            report[db_name] = {
                "status": "failed",
                "source": src,
                "local_path": dst,
                "error": str(e),
            }

    shutil.rmtree(staging_root, ignore_errors=True)

    print("\n" + "=" * 100)
    print("DOWNLOAD SUMMARY")
    print("=" * 100)

    for name, info in sorted(report.items()):
        status = info.get("status", "unknown")
        size_mb = info.get("size_mb", 0.0)
        local_path = info.get("local_path")
        kind = info.get("kind", "")

        print(
            f"{name:30s} "
            f"{status:15s} "
            f"{kind:5s} "
            f"{size_mb:10.2f} MB  -> {local_path}"
        )

    return report


def verify_milvus_dbs_v2(domain_names=None, embedding_type=None, index_version=None):
    embedding_type = embedding_type or EMBEDDING_TYPE
    index_version = index_version or INDEX_VERSION

    if domain_names is None:
        domain_names = DOMAIN_NAMES

    print("\n" + "=" * 100)
    print("LOCAL MILVUS DB VERIFY V2")
    print("=" * 100)
    print(f"Embedding Type : {embedding_type}")
    print(f"Index Version  : {index_version}")

    found = []
    missing = []

    for domain in domain_names:
        p = get_db_path(
            domain_name=domain,
            embedding_type=embedding_type,
            index_version=index_version,
        )

        if os.path.exists(p):
            kind = "DIR" if os.path.isdir(p) else "FILE"
            size_mb = get_path_size_mb(p)

            print(f"  OK      {domain:20s} -> [{kind}] {p} ({size_mb:.2f} MB)")
            found.append(domain)

        else:
            print(f"  MISSING {domain:20s} -> {p}")
            missing.append(domain)

    print("\nSummary:")
    print(f"  Found   : {len(found)}")
    print(f"  Missing : {len(missing)}")

    if missing:
        print(f"  Missing domains: {missing}")

    return {
        "found": found,
        "missing": missing,
    }

EMBEDDING_TYPE = "llm_embedder"
INDEX_VERSION = "default"

download_report = download_milvus_dbs_from_bucket_v2(
    embedding_type=EMBEDDING_TYPE,
    index_version=INDEX_VERSION,
    domain_names=DOMAIN_NAMES,
    force_download=False,
    debug=True,
)

verify_report = verify_milvus_dbs_v2(
    domain_names=DOMAIN_NAMES,
    embedding_type=EMBEDDING_TYPE,
    index_version=INDEX_VERSION,
)

llm_client = get_llm_client()
print(f"LLM client ready.  Provider: {LLM_PROVIDER}")

# ── Cell 7: Load embedding model (LLM-Embedder by default) ────────────────────
embed_model_name = EMBED_MODELS[EMBEDDING_TYPE]
print(f"Loading embedding model: {embed_model_name}")
embed_model = SentenceTransformer(embed_model_name)
print("Embedding model loaded.")

# ── Cell 9: Open Milvus clients for all domains ────────────────────────────────
def load_milvus_clients(embedding_type=None):
    global milvus_clients
    etype = embedding_type or EMBEDDING_TYPE
    milvus_clients = {}
    for domain in DOMAIN_NAMES:
        db_path = get_db_path(domain, etype)
        if not os.path.exists(db_path):
            print(f"  DB not found, skipping: {db_path}")
            continue
        try:
            client = MilvusClient(db_path)
            collections = client.list_collections()
            print(f"  {domain}: {collections}")
            if collections:
                milvus_clients[domain] = client
        except Exception as e:
            print(f"  Failed to open {domain}: {e}")
    print(f"Loaded {len(milvus_clients)} domain clients: {list(milvus_clients.keys())}")

load_milvus_clients()

# Build BM25 indexes (needed for hybrid search)
build_all_bm25_indexes(milvus_clients)
print("BM25 indexes ready.")

# ── Cell 10: Load RAGBench (test split only) + sample catalogue ───────────────
DATASET_BY_DOMAIN = {
    "Bio_Medical":       ["covidqa", "pubmedqa"],
    "General_Knowledge": ["expertqa", "hagrid", "hotpotqa", "msmarco"],
    "Customer_Support":  ["delucionqa", "emanual", "techqa"],
    "Finance":           ["finqa", "tatqa"],
    "Legal_Contracts":   ["cuad"],
}

# sample_store[domain][dataset] = list of row dicts from the test split
sample_store = {}

def load_ragbench(domains=None):
    global ragbench_by_domain, sample_store
    domains = domains or list(DATASET_BY_DOMAIN.keys())
    for domain in domains:
        ragbench_by_domain[domain] = {}
        sample_store[domain]       = {}
        for ds_name in DATASET_BY_DOMAIN.get(domain, []):
            try:
                ds = load_dataset("rungalileo/ragbench", ds_name)
                ragbench_by_domain[domain][ds_name] = ds

                if "test" not in ds:
                    print(f"  WARNING: no 'test' split for {domain}/{ds_name}, skipping")
                    continue

                rows = []
                for idx, row in enumerate(ds["test"]):
                    rows.append({
                        "idx":       idx,
                        "question":  row.get("question", ""),
                        "response":  row.get("response", ""),
                        "documents": row.get("documents", []),
                        "gold_relevance":    row.get("relevance_score"),
                        "gold_utilization":  row.get("utilization_score"),
                        "gold_completeness": row.get("completeness_score"),
                        "gold_adherence":    row.get("adherence_score"),
                    })
                sample_store[domain][ds_name] = rows
                print(f"  Loaded: {domain}/{ds_name}  test rows={len(rows)}")
            except Exception as e:
                print(f"  Failed: {domain}/{ds_name}: {e}")
    print(f"\nRAGBench loaded (test only). Domains: {list(sample_store.keys())}")

load_ragbench()


# ── Helpers for cascading dropdowns ───────────────────────────────────────────

def get_datasets_for_domain(domain):
    return list(sample_store.get(domain, {}).keys())

def get_sample_ids_for_dataset(domain, dataset):
    """Return 'idx – first 80 chars of question' labels for the test split."""
    rows = sample_store.get(domain, {}).get(dataset, [])
    labels = []
    for r in rows:
        q = r["question"]
        labels.append(f"{r['idx']}  –  {q[:80]}{'…' if len(q) > 80 else ''}")
    return labels

def get_row_by_label(domain, dataset, label):
    """Retrieve a stored row dict from a label string."""
    if not label: return None
    idx_str = label.split("–")[0].strip()
    try:
        idx = int(idx_str)
    except ValueError:
        return None
    rows = sample_store.get(domain, {}).get(dataset, [])
    return next((r for r in rows if r["idx"] == idx), None)

print("Sample catalogue ready.")

# ── Cell 11: Evaluation helpers + all Gradio handlers ─────────────────────────

# ── Judge / evaluation ─────────────────────────────────────────────────────────

def build_keyed_response(answer):
    return {f"r_{i}": s for i, s in enumerate(split_into_sentences(answer))}

def build_sentence_keyed_docs(retrieved_docs):
    keyed = {}
    for di, doc in enumerate(retrieved_docs):
        text = doc.get("text","") if isinstance(doc, dict) else doc
        for si, s in enumerate(split_into_sentences(text)):
            keyed[f"{di}_{si}"] = s
    return keyed

def build_evaluation_prompt(documents_text, question, answer_text):
    return f"""Evaluate the RAG response using the provided documents.

Documents (sentence-keyed):
{documents_text}

Question:
{question}

Response (sentence-keyed):
{answer_text}

Return ONLY valid JSON:
{{
  "overall_supported": true,
  "all_relevant_sentence_keys": ["0_0"],
  "all_utilized_sentence_keys": ["0_0"],
  "sentence_support_information": [
    {{"response_sentence_key": "r_0", "supporting_sentence_keys": ["0_0"], "fully_supported": true}}
  ]
}}
Rules: document keys look like 0_0; response keys like r_0. Return only JSON.""".strip()

def ask_judge(prompt, llm_client, judge_model, max_retries=5):
    last_error = None
    for attempt in range(max_retries):
        try:
            resp = llm_client.chat.completions.create(
                model=judge_model,
                messages=[
                    {"role":"system","content":"You are a strict RAG evaluation judge. Return ONLY valid JSON. No markdown. No <think> tags."},
                    {"role":"user","content":_sanitize(prompt)},
                ],
                temperature=0.0, max_tokens=3000,
            )
            return _safe_message_content(resp)
        except Exception as e:
            last_error = e; msg = str(e)
            wait = 2**attempt
            if "429" in msg or "rate_limit" in msg:
                m = re.search(r"try again in ([\\d.]+)s", msg)
                if m: wait = float(m.group(1))
            elif not any(x in msg for x in ["503","502","504","over capacity","gateway"]):
                raise
            time.sleep(wait + random.uniform(0.1, 0.5))
    raise RuntimeError(f"Judge failed after {max_retries} retries: {last_error}")

def parse_judge_json(raw):
    if not raw: raise ValueError("Judge output empty")
    cleaned = re.sub(r"<think>.*?</think>","",str(raw),flags=re.DOTALL).strip()
    cleaned = cleaned.replace("```json","").replace("```","").strip()
    s, e = cleaned.find("{"), cleaned.rfind("}")
    if s == -1 or e == -1: raise ValueError(f"No JSON: {cleaned[:300]}")
    cleaned = cleaned[s:e+1]
    cleaned = re.sub(r"}\s*{","}, {",cleaned)
    cleaned = re.sub(r",\s*([}\]])",r"\1",cleaned)
    return json.loads(cleaned)

def evaluate_ragbench_json(judge_json, keyed_docs):
    vk  = set(keyed_docs.keys())
    rel = set(judge_json.get("all_relevant_sentence_keys", [])) & vk
    utl = set(judge_json.get("all_utilized_sentence_keys", [])) & vk
    ovl = rel & utl; n = len(vk)
    return {
        "adherence_score":    int(bool(judge_json.get("overall_supported", False))),
        "hallucination_flag": 1 - int(bool(judge_json.get("overall_supported", False))),
        "relevance_score":    float(np.clip(len(rel)/n        if n   else 0, 0, 1)),
        "utilization_score":  float(np.clip(len(utl)/n        if n   else 0, 0, 1)),
        "completeness_score": float(np.clip(len(ovl)/len(rel) if rel else 0, 0, 1)),
    }


# ── Source badge helper ────────────────────────────────────────────────────────

def _source_badge(source, model, extra=None):
    parts = [f"[Source: {source}  |  model: {model}"]
    if extra:
        parts += [f"  |  {k}: {v}" for k, v in extra.items()]
    parts.append("]")
    return "".join(parts)


# ── DB status helper ───────────────────────────────────────────────────────────

def db_status_md():
    if not milvus_clients:
        return (
            "> **No vector DBs loaded.**  "
            "Re-run **Cell 8** (download DBs) then **Cell 9** (open clients), "
            "then re-run this cell."
        )
    loaded = ", ".join(f"`{d}`" for d in sorted(milvus_clients.keys()))
    return f"> **Loaded domains ({len(milvus_clients)}):** {loaded}"


# ── Config applier ─────────────────────────────────────────────────────────────

def apply_config(llm_choice, embed_choice,
                 enable_hybrid, enable_hyde, enable_reranking, reranker_type,
                 enable_repacking, repack_strategy,
                 enable_summarization, summarization_type,
                 prompt_strategy, hybrid_alpha, top_k,
                 enable_query_classification, enable_query_rewriting, enable_query_decomp):
    global MODEL_NAME, EMBEDDING_TYPE, embed_model
    global ENABLE_HYBRID, ENABLE_HYDE, ENABLE_RERANKING, RERANKER_TYPE
    global ENABLE_REPACKING, REPACK_STRATEGY, ENABLE_SUMMARIZATION, SUMMARIZATION_TYPE
    global PROMPT_STRATEGY, HYBRID_ALPHA
    global ENABLE_QUERY_CLASSIFICATION, ENABLE_QUERY_REWRITING, ENABLE_QUERY_DECOMPOSITION

    MODEL_NAME                  = llm_choice
    ENABLE_HYBRID               = enable_hybrid
    ENABLE_HYDE                 = enable_hyde
    ENABLE_RERANKING            = enable_reranking
    RERANKER_TYPE               = reranker_type
    ENABLE_REPACKING            = enable_repacking
    REPACK_STRATEGY             = repack_strategy
    ENABLE_SUMMARIZATION        = enable_summarization
    SUMMARIZATION_TYPE          = summarization_type
    PROMPT_STRATEGY             = prompt_strategy
    HYBRID_ALPHA                = float(hybrid_alpha)
    ENABLE_QUERY_CLASSIFICATION = enable_query_classification
    ENABLE_QUERY_REWRITING      = enable_query_rewriting
    ENABLE_QUERY_DECOMPOSITION  = enable_query_decomp

    if embed_choice != EMBEDDING_TYPE:
        EMBEDDING_TYPE = embed_choice
        print(f"Reloading embedding model: {EMBED_MODELS[embed_choice]}")
        embed_model = SentenceTransformer(EMBED_MODELS[embed_choice])
        download_milvus_dbs_from_bucket_v2(embed_choice)
        load_milvus_clients(embed_choice)
        build_all_bm25_indexes(milvus_clients)


# ── Cascading dropdown callbacks ───────────────────────────────────────────────

_NONE_DOMAIN = "None (direct LLM, no retrieval)"

def on_domain_change(domain):
    if domain == _NONE_DOMAIN:
        return gr.update(choices=[], value=None), gr.update(choices=[], value=None), gr.update()
    datasets   = get_datasets_for_domain(domain)
    ds         = datasets[0] if datasets else None
    sample_ids = get_sample_ids_for_dataset(domain, ds) if ds else []
    return (
        gr.update(choices=datasets,   value=ds),
        gr.update(choices=sample_ids, value=None),
        gr.update(value=""),
    )

def on_dataset_change(domain, dataset):
    if domain == _NONE_DOMAIN or not dataset:
        return gr.update(choices=[], value=None), gr.update(value="")
    sample_ids = get_sample_ids_for_dataset(domain, dataset)
    return gr.update(choices=sample_ids, value=None), gr.update(value="")

def on_sample_select(domain, dataset, label):
    if domain == _NONE_DOMAIN or not label:
        return gr.update()
    row = get_row_by_label(domain, dataset, label)
    if row is None: return gr.update()
    return gr.update(value=row["question"])


# ── Chunk display helpers ──────────────────────────────────────────────────────

def _format_chunks(docs, title="Retrieved"):
    if not docs:
        return f"_No documents for {title}._"
    parts = []
    for i, doc in enumerate(docs):
        if isinstance(doc, dict):
            text  = doc.get("text", str(doc))
            score = doc.get("rerank_score", doc.get("score", 0.0))
            tags  = []
            if doc.get("summarized"):    tags.append(f"summarized/{doc.get('summary_type','')}")
            if doc.get("reranker_type"): tags.append(f"reranked/{doc.get('reranker_type','')}")
            if ENABLE_HYBRID:            tags.append(f"dense={doc.get('dense_score',0):.3f} bm25={doc.get('bm25_score',0):.3f}")
            tag_str = f"  `{'  |  '.join(tags)}`" if tags else ""
        else:
            text, score, tag_str = str(doc), 0.0, ""
        parts.append(f"**{title} Chunk {i+1}** — score: `{score:.4f}`{tag_str}\n\n{text}")
    return "\n\n---\n\n".join(parts)

def _format_gt_docs(doc_list):
    if not doc_list:
        return "_No ground-truth documents stored for this sample._"
    parts = []
    for i, text in enumerate(doc_list):
        parts.append(f"**GT Doc {i+1}**\n\n{str(text)}")
    return "\n\n---\n\n".join(parts)


# ── Main run handler ───────────────────────────────────────────────────────────

def run_query(
    query, domain,
    dataset_sel, sample_label,
    llm_choice, judge_llm_choice, embed_choice,
    enable_hybrid, enable_hyde, enable_reranking, reranker_type,
    enable_repacking, repack_strategy,
    enable_summarization, summarization_type,
    prompt_strategy, hybrid_alpha, top_k,
    enable_query_classification, enable_query_rewriting, enable_query_decomp,
    run_judge,
):
    query = _sanitize(query)

    if not query.strip():
        return ("Please enter a query.",) + ("",)*4

    if llm_client is None:
        return ("LLM client is not initialised.\n\nRe-run Cell 6, then re-run Cell 12.",) + ("",)*4

    apply_config(
        llm_choice, embed_choice,
        enable_hybrid, enable_hyde, enable_reranking, reranker_type,
        enable_repacking, repack_strategy,
        enable_summarization, summarization_type,
        prompt_strategy, float(hybrid_alpha), int(top_k),
        enable_query_classification, enable_query_rewriting, enable_query_decomp,
    )

    # ── Domain = None → direct LLM, skip all retrieval ───────────────────────
    if domain == _NONE_DOMAIN or not domain:
        try:
            direct_ans = _safe_message_content(llm_client.chat.completions.create(
                model=MODEL_NAME,
                messages=[{"role":"system","content":"You are a helpful assistant."},
                          {"role":"user","content":query}],
                temperature=0.3, max_tokens=800,
            ))
        except Exception as e:
            direct_ans = f"Direct LLM error: {e}"
        badge = _source_badge("Direct LLM (no retrieval)", MODEL_NAME)
        note  = "_[Domain set to None — answered directly by LLM without vector DB retrieval]_"
        return f"{badge}\n\n{direct_ans}", note, note, note, note

    # ── Guard: domain must be loaded ──────────────────────────────────────────
    if not milvus_clients:
        return ("No vector DBs are loaded.\n\nRe-run Cell 8 then Cell 9, then re-run Cell 12.",) + ("",)*4

    if domain not in milvus_clients:
        return (
            f"Domain '{domain}' is not loaded.\n"
            f"Loaded domains: {list(milvus_clients.keys())}\n\nRe-run Cell 8 and Cell 9.",
        ) + ("",)*4

    # ── Query Classification ──────────────────────────────────────────────────
    route = classify_query(query, domain_name=domain)
    if route == "LLM":
        try:
            direct_ans = _safe_message_content(llm_client.chat.completions.create(
                model=MODEL_NAME,
                messages=[{"role":"system","content":"You are a concise factual assistant."},
                          {"role":"user","content":query}],
                temperature=0.2, max_tokens=500,
            ))
        except Exception as e:
            direct_ans = f"Direct LLM error: {e}"
        badge = _source_badge("Direct LLM", MODEL_NAME)
        note  = "_[Query Classifier routed this to direct LLM — no retrieval performed]_"
        return f"{badge}\n\n{direct_ans}", note, note, note, note

    # ── Query Rewriting + Decomposition ──────────────────────────────────────
    rewritten  = rewrite_query(query, domain, llm_client) if ENABLE_QUERY_REWRITING else query
    subqueries = decompose_query(rewritten, llm_client, domain=domain) if ENABLE_QUERY_DECOMPOSITION else [rewritten]

    # ── Retrieve + Generate ───────────────────────────────────────────────────
    all_retrieved, all_answers = [], []
    for sq in subqueries:
        try:
            docs = retrieve(sq, domain, embed_model=embed_model, llm_client=llm_client, top_k=int(top_k))
        except Exception as e:
            return (f"Retrieval error: {e}",) + ("",)*4
        if not docs: continue
        all_retrieved.extend(docs)
        ctx = _sanitize("\n\n".join(d.get("text","") if isinstance(d,dict) else d for d in docs))
        sq  = _sanitize(sq)
        try:
            all_answers.append(ask_rag(ctx, sq, llm_client, strategy=PROMPT_STRATEGY))
        except Exception as e:
            return (f"Generation error: {e}",) + ("",)*4

    if not all_retrieved:
        return ("No documents retrieved.",) + ("",)*4

    raw_answer = "\n\n".join(all_answers)

    # ── Source badge ──────────────────────────────────────────────────────────
    active = {"prompt": PROMPT_STRATEGY, "chunks": len(all_retrieved)}
    if ENABLE_HYBRID:        active["hybrid"]  = f"alpha={HYBRID_ALPHA}"
    if ENABLE_HYDE:          active["hyde"]    = "on"
    if ENABLE_RERANKING:     active["rerank"]  = RERANKER_TYPE
    if ENABLE_SUMMARIZATION: active["summ"]    = SUMMARIZATION_TYPE
    if ENABLE_REPACKING:     active["repack"]  = REPACK_STRATEGY
    if len(subqueries) > 1:  active["subq"]    = len(subqueries)
    rag_response = f"{_source_badge('RAG', MODEL_NAME, extra=active)}\n\n{raw_answer}"

    # ── Ground truth lookup ───────────────────────────────────────────────────
    row = get_row_by_label(domain, dataset_sel, sample_label) if sample_label else None
    if row is None:
        for ds_name, rows in sample_store.get(domain, {}).items():
            match = next((r for r in rows if r["question"].strip().lower() == query.strip().lower()), None)
            if match: row = match; break

    ground_truth = row["response"]  if row else "_(no matching sample found)_"
    gt_docs_md   = _format_gt_docs(row["documents"] if row else [])
    rag_docs_md  = _format_chunks(all_retrieved, title="RAG")

    # ── Judge evaluation ──────────────────────────────────────────────────────
    metrics_md = "_Judge evaluation not requested._"
    if run_judge:
        try:
            keyed_docs   = build_sentence_keyed_docs(all_retrieved)
            keyed_answer = build_keyed_response(raw_answer)
            docs_text = _sanitize("\n".join(f"{k}: {v}" for k,v in keyed_docs.items()))
            ans_text  = _sanitize("\n".join(f"{k}: {v}" for k,v in keyed_answer.items()))
            raw       = ask_judge(build_evaluation_prompt(docs_text, query, ans_text), llm_client, judge_llm_choice)
            pred      = evaluate_ragbench_json(parse_judge_json(raw), keyed_docs)
            gold = {k: row.get(f"gold_{k}") for k in ("relevance","utilization","completeness","adherence")} if row else {}
            def _f(v): return f"{v:.3f}" if isinstance(v, float) else (str(v) if v is not None else "—")
            metrics_md = "\n".join([
                "| Metric | Predicted | Gold |",
                "|--------|-----------|------|",
                f"| Relevance    | {_f(pred['relevance_score'])}    | {_f(gold.get('relevance'))} |",
                f"| Utilization  | {_f(pred['utilization_score'])}  | {_f(gold.get('utilization'))} |",
                f"| Completeness | {_f(pred['completeness_score'])} | {_f(gold.get('completeness'))} |",
                f"| Adherence    | {_f(pred['adherence_score'])}    | {_f(gold.get('adherence'))} |",
                f"| Hallucination| {_f(pred['hallucination_flag'])} | — |",
            ])
        except Exception as e:
            metrics_md = f"Judge error: {e}"

    return ground_truth, rag_response, gt_docs_md, rag_docs_md, metrics_md


print("Handlers ready.")

# ── Cell 12: Gradio UI ────────────────────────────────────────────────────────

AVAILABLE_DOMAINS = list(milvus_clients.keys())
DEFAULT_DOMAIN    = AVAILABLE_DOMAINS[0] if AVAILABLE_DOMAINS else None

_init_datasets = get_datasets_for_domain(DEFAULT_DOMAIN) if DEFAULT_DOMAIN else []
_init_ds       = _init_datasets[0] if _init_datasets else None
_init_samples  = get_sample_ids_for_dataset(DEFAULT_DOMAIN, _init_ds) if _init_ds else []

CSS = """
footer { display: none !important; }
.gr-button-primary { font-size: 1.1rem !important; }
"""

with gr.Blocks(title="RAG Capstone Demo") as demo:

    # ── Header ────────────────────────────────────────────────────────────────
    gr.Markdown("# 🔍 RAG Capstone — Interactive Demo")
    gr.Markdown(
        f"Provider: **{LLM_PROVIDER.upper()}** &nbsp;|&nbsp; "
        "Type any question in the **Query** box. Pick a **Domain** to search its vector DB, "
        "or leave Domain as **None** to get a direct LLM answer without retrieval. "
        "Expand **Sample Selector** to load a test-split example."
    )
    gr.Markdown(db_status_md())

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 1 — Query + Domain  (always visible)
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Row():
        query_input = gr.Textbox(
            lines=3,
            placeholder="Type any question here… or expand Sample Selector below to auto-fill from the dataset.",
            label="Query",
            scale=4,
        )
        domain_dd = gr.Dropdown(
            choices=["None (direct LLM, no retrieval)"] + AVAILABLE_DOMAINS,
            value="None (direct LLM, no retrieval)" if not AVAILABLE_DOMAINS else DEFAULT_DOMAIN,
            label="Domain",
            info="None = direct LLM answer; pick a domain to run full RAG retrieval",
            scale=1,
        )

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 2 — Sample Selector  (collapsed by default — optional)
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Accordion("📋  Sample Selector  (optional — expand to load a test-split example)", open=False):
        gr.Markdown(
            "_Select a preloaded sample to auto-fill the Query box and Domain above. "
            "Leave collapsed to ask your own question._"
        )
        with gr.Row():
            dataset_dd = gr.Dropdown(choices=_init_datasets, value=_init_ds,  label="Dataset", scale=1)
            sample_dd  = gr.Dropdown(choices=_init_samples,  value=None,
                                     label="Sample ID  (idx – question preview)", scale=4)

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 3 — Control Panel  (collapsed by default)
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Accordion("⚙️  Control Panel", open=False):
        with gr.Tabs():

            with gr.Tab("🤖 Models"):
                with gr.Row():
                    llm_choice = gr.Dropdown(
                        choices=LLM_CHOICES, value=LLM_CHOICES[0],
                        label="Generator LLM", info="Produces the RAG answer")
                    judge_llm_choice = gr.Dropdown(
                        choices=LLM_CHOICES, value=LLM_CHOICES[4] if len(LLM_CHOICES) > 4 else LLM_CHOICES[-1],
                        label="Judge LLM", info="Used for evaluation scoring")
                    embed_choice = gr.Dropdown(
                        choices=EMBEDDING_CHOICES, value="llm_embedder",
                        label="Embedding Model", info="Changing this reloads the vector DB")

            with gr.Tab("🔄 Query Processing"):
                gr.Markdown("Applied **before** retrieval, in order: Classify → Rewrite → Decompose")
                with gr.Row():
                    enable_query_classification = gr.Checkbox(
                        label="Query Classification", value=False,
                        info="Route simple factual queries directly to LLM, skip retrieval. "
                             "All benchmark domain queries always use RAG regardless.")
                with gr.Row():
                    enable_query_rewriting = gr.Checkbox(
                        label="Query Rewriting", value=False,
                        info="LLM rewrites the query to improve retrieval")
                    enable_query_decomp = gr.Checkbox(
                        label="Query Decomposition", value=False,
                        info="Break multi-part queries into subqueries")

            with gr.Tab("🔎 Retrieval"):
                with gr.Row():
                    top_k = gr.Slider(minimum=1, maximum=10, step=1, value=3,
                                      label="Top-K chunks returned")
                    hybrid_alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5,
                                             label="Hybrid Alpha  (1=dense, 0=BM25)")
                with gr.Row():
                    enable_hybrid = gr.Checkbox(label="Hybrid Search  (Dense + BM25)", value=False)
                    enable_hyde   = gr.Checkbox(label="HyDE  (query expansion)", value=False)

            with gr.Tab("↕️ Reranking"):
                with gr.Row():
                    enable_reranking = gr.Checkbox(label="Enable Reranking", value=False)
                    reranker_type    = gr.Radio(choices=["monot5","tilde"], value="monot5",
                                               label="Reranker", info="MonoT5: seq2seq | TILDE: cross-encoder")

            with gr.Tab("📦 Repacking"):
                with gr.Row():
                    enable_repacking = gr.Checkbox(label="Enable Repacking", value=False)
                    repack_strategy  = gr.Radio(choices=["forward","reverse","sides"], value="sides",
                                               label="Strategy", info="forward | reverse | U-shape sides")

            with gr.Tab("📝 Summarization"):
                with gr.Row():
                    enable_summarization = gr.Checkbox(label="Enable Summarization", value=False)
                    summarization_type   = gr.Radio(choices=["recomp","longllmlingua"], value="recomp",
                                                   label="Method", info="RECOMP: extractive | LLMLingua: token compression")

            with gr.Tab("💬 Prompt"):
                prompt_strategy = gr.Radio(
                    choices=["short","long","long_cot"], value="short",
                    label="Prompt Strategy",
                    info="short: minimal | long: strict no-hallucination | long_cot: step-by-step")

            with gr.Tab("⚖️ Judge"):
                run_judge = gr.Checkbox(
                    label="Run Judge evaluation after generation", value=False,
                    info="~1 extra LLM call. Gold scores shown only for preloaded samples.")
                gr.Markdown("_Judge LLM is configured in the **Models** tab._")

    # ── Run button ────────────────────────────────────────────────────────────
    run_btn = gr.Button("▶  Run Query", variant="primary", size="lg")

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 4 — Responses  (Ground Truth LEFT, RAG RIGHT)
    # ══════════════════════════════════════════════════════════════════════════
    gr.Markdown("## 💬 Responses")
    with gr.Row(equal_height=True):
        gt_out  = gr.Textbox(label="Ground Truth Response", lines=10, interactive=False, scale=1)
        rag_out = gr.Textbox(label="RAG Response",          lines=10, interactive=False, scale=1)

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 5 — Retrieved Documents  (GT LEFT, RAG RIGHT)
    # ══════════════════════════════════════════════════════════════════════════
    gr.Markdown("## 📄 Retrieved Documents")
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            gr.Markdown("### Ground Truth Documents")
            gt_docs_out  = gr.Markdown(value="_Select a preloaded sample to see GT documents._")
        with gr.Column(scale=1):
            gr.Markdown("### RAG Retrieved Documents")
            rag_docs_out = gr.Markdown(value="_Run a query to see RAG retrieved chunks._")

    # ══════════════════════════════════════════════════════════════════════════
    # SECTION 6 — Metrics
    # ══════════════════════════════════════════════════════════════════════════
    with gr.Accordion("📊  Metrics  (Gold vs Predicted)", open=False):
        metrics_out = gr.Markdown(value="_Enable the Judge in the Control Panel and run a query._")

    # ── Sample Selector cascade ───────────────────────────────────────────────
    domain_dd.change(
        fn=on_domain_change, inputs=[domain_dd],
        outputs=[dataset_dd, sample_dd, query_input],
    )
    dataset_dd.change(
        fn=on_dataset_change, inputs=[domain_dd, dataset_dd],
        outputs=[sample_dd, query_input],
    )
    sample_dd.change(
        fn=on_sample_select, inputs=[domain_dd, dataset_dd, sample_dd],
        outputs=[query_input],
    )

    # ── Run wiring ────────────────────────────────────────────────────────────
    _config_inputs = [
        llm_choice, judge_llm_choice, embed_choice,
        enable_hybrid, enable_hyde, enable_reranking, reranker_type,
        enable_repacking, repack_strategy,
        enable_summarization, summarization_type,
        prompt_strategy, hybrid_alpha, top_k,
        enable_query_classification, enable_query_rewriting, enable_query_decomp,
        run_judge,
    ]
    _all_inputs  = [query_input, domain_dd, dataset_dd, sample_dd] + _config_inputs
    _all_outputs = [gt_out, rag_out, gt_docs_out, rag_docs_out, metrics_out]

    run_btn.click(fn=run_query, inputs=_all_inputs, outputs=_all_outputs)
    query_input.submit(fn=run_query, inputs=_all_inputs, outputs=_all_outputs)

demo.launch(
    share=True,
    debug=True,
    theme=gr.themes.Soft(),
    css=CSS,
)