model_utils.py

# model_utils.py
from typing import List, Optional, Any, Dict, Tuple, Mapping, Sequence
from difflib import SequenceMatcher
import hashlib
import html
import json
import re
import os
import shutil

# This app uses PyTorch only; disable TensorFlow path in transformers.
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("TRANSFORMERS_NO_TF", "1")

import gradio as gr
import numpy as np
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    StoppingCriteria,
    StoppingCriteriaList,
)
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

import qa_store
from answer_repair import (
    build_consensus_structure_fallback,
    build_extractive_evidence_answer as _build_repaired_extractive_answer,
    component_is_supported,
    evidence_span_for_component,
    select_manual_near_miss_candidates,
)
from routing_policy import (
    GLOSSARY_SEMANTIC_MARGIN,
    GLOSSARY_SEMANTIC_THRESHOLD,
    QA_GLOBAL_SEMANTIC_MARGIN,
    QA_GLOBAL_SEMANTIC_THRESHOLD,
    QA_SCOPED_NEAR_MATCH_EVIDENCE_MIN,
    QA_SCOPED_SEMANTIC_MARGIN,
    QA_SCOPED_SEMANTIC_THRESHOLD,
    RAG_GLOBAL_MIN_SIMILARITY,
    RAG_SCOPED_MIN_SIMILARITY,
    SAFE_LAO_REFUSAL,
    assess_retrieval_confidence,
    extract_compound_components,
    glossary_question_eligibility,
    is_compound_question,
    item_in_scope,
    normalize_scope,
    prepare_generated_answer_text,
    scope_label,
    semantic_acceptance,
    significant_tokens,
    split_compound_question,
    strong_evidence_for_extractive_fallback,
    validate_generated_answer,
)
from loader import (
    load_curriculum,
    load_manual_qa,
    manual_qa_write_lock,
    rebuild_combined_qa,
    load_glossary,
    sync_download_manual_qa,
    sync_download_cache,
    sync_upload_data_tree,
    sync_upload_cache,
    CACHE_PATH,
)

# -----------------------------
# Base chat model
# -----------------------------
MODEL_NAME = "SeaLLMs/SeaLLMs-v3-1.5B-Chat"
EMBED_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CACHE_FILE = os.path.join(BASE_DIR, "data", "cached_embeddings.pt")
CACHE_DOWNLOAD_FILE = os.path.join(BASE_DIR, "data", "cached_embeddings.download.pt")
CACHE_SCHEMA_VERSION = 2
GENERATION_PROMPT_VERSION = "natural-science-rag-v3-output-guarded"
MAX_GENERATED_TOKENS = 160
GENERATION_DO_SAMPLE = False
GENERATION_TEMPERATURE = 1.0
GENERATION_TOP_P = 1.0
OUTPUT_MIN_QUESTION_SEMANTIC_SIMILARITY = 0.32

device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=dtype)
model.to(device)
model.eval()
model.generation_config.do_sample = GENERATION_DO_SAMPLE
model.generation_config.temperature = GENERATION_TEMPERATURE
model.generation_config.top_p = GENERATION_TOP_P

embed_model = SentenceTransformer(EMBED_MODEL_NAME)
embed_model = embed_model.to(device)

# Number of textbook entries to include in the RAG context
MAX_CONTEXT_ENTRIES = 4
MIN_QA_QUERY_CHARS = 12
GUIDE_SUGGESTION_LIMIT: Optional[int] = None
BROAD_QA_LIMIT = 4
BROAD_CONTEXT_LIMIT = 2
BROAD_MIN_EMBED_SIM = 0.46

GUIDE_WELCOME_MESSAGE = (
    "ເພື່ອຄວາມໄວ ແລະ ຕອບໄດ້ຕົງກັບປຶ້ມແບບຮຽນ ກະລຸນາລະບຸຊັ້ນຮຽນຂອງທ່ານ:\n\n"
    "ມ.1\n"
    "ມ.2\n"
    "ມ.3\n"
    "ມ.4"
)

_LAO_TO_ARABIC_DIGITS = str.maketrans("໐໑໒໓໔໕໖໗໘໙", "0123456789")


# ------------------------------------------------------------
# Helpers: YES/NO follow-up handling (context linking)
# ------------------------------------------------------------
_YES_SET = {
    "ແມ່ນ", "ແມ່ນແລ້ວ", "ແມ່ນນະ", "ແມ່ນເດ", "ແມ່ນເດີ",
    "ok", "okay", "yes", "y"
}
_NO_SET = {
    "ບໍ່", "ບໍ່ແມ່ນ", "ບໍ່ແມ່ນເດ", "no", "n"
}

def _is_yes_no_reply(text: str) -> bool:
    t = (text or "").strip().lower()
    return t in _YES_SET or t in _NO_SET

def _extract_last_user_text(history: Any) -> Optional[str]:
    """
    Gradio ChatInterface usually passes history as:
      List[Tuple[user_msg, bot_msg]]
    But we support a few other shapes just in case.
    """
    if not history:
        return None

    if isinstance(history, list):
        for item in reversed(history):
            if item is None:
                continue
            if isinstance(item, (tuple, list)) and len(item) >= 1:
                u = item[0]
                if isinstance(u, str) and u.strip():
                    return u.strip()
            if isinstance(item, dict):
                if item.get("role") == "user":
                    c = item.get("content")
                    if isinstance(c, str) and c.strip():
                        return c.strip()
    return None


# -----------------------------
# Embedding builders
# -----------------------------
def _entry_embedding_texts() -> List[str]:
    texts: List[str] = []
    for e in qa_store.ENTRIES:
        chapter = e.get("chapter_title", "") or e.get("chapter", "") or ""
        section = e.get("section_title", "") or e.get("section", "") or ""
        text = e.get("text", "") or ""
        texts.append(f"{chapter}\n{section}\n{text}")
    return texts


def _glossary_embedding_texts() -> List[str]:
    return [
        f"{item.get('term', '')} :: {item.get('definition', '')}"
        for item in qa_store.GLOSSARY
    ]


def _hash_texts(texts: List[str]) -> str:
    digest = hashlib.sha256()
    for text in texts:
        digest.update((text or "").encode("utf-8", errors="surrogatepass"))
        digest.update(b"\0")
    return digest.hexdigest()


def _cache_metadata(textbook_texts: List[str], glossary_texts: List[str]) -> dict:
    return {
        "schema_version": CACHE_SCHEMA_VERSION,
        "embedding_model": EMBED_MODEL_NAME,
        "textbook_count": len(textbook_texts),
        "textbook_hash": _hash_texts(textbook_texts),
        "glossary_count": len(glossary_texts),
        "glossary_hash": _hash_texts(glossary_texts),
    }


def _cache_len(value: Any) -> int:
    if value is None:
        return 0
    return len(value)


def _cache_matches(cache: dict, expected: dict) -> bool:
    metadata = cache.get("metadata")
    if not isinstance(metadata, dict):
        return False

    keys = (
        "schema_version",
        "embedding_model",
        "textbook_count",
        "textbook_hash",
        "glossary_count",
        "glossary_hash",
    )
    if any(metadata.get(key) != expected.get(key) for key in keys):
        return False

    return (
        _cache_len(cache.get("textbook")) == expected["textbook_count"]
        and _cache_len(cache.get("glossary")) == expected["glossary_count"]
    )


def _load_cache_from_path(path: str, expected: dict, label: str) -> bool:
    if not os.path.exists(path):
        return False

    try:
        print(f"[INFO] Loading {label} cached embeddings from {path}...")
        # Cache is generated by this app, so loading full object graph is intentional.
        cache = torch.load(path, map_location=device, weights_only=False)
    except Exception as e:
        print(f"[WARN] Failed to load {label} cache: {e}")
        return False

    if not isinstance(cache, dict) or not _cache_matches(cache, expected):
        textbook_len = _cache_len(cache.get("textbook")) if isinstance(cache, dict) else 0
        glossary_len = _cache_len(cache.get("glossary")) if isinstance(cache, dict) else 0
        print(
            "[WARN] Ignoring stale cache "
            f"({label}: textbook={textbook_len}/{expected['textbook_count']}, "
            f"glossary={glossary_len}/{expected['glossary_count']})."
        )
        return False

    textbook = cache.get("textbook")
    glossary = cache.get("glossary")
    qa_store.TEXT_EMBEDDINGS = textbook.to(device) if textbook is not None else None
    if isinstance(glossary, torch.Tensor):
        glossary = glossary.detach().cpu().numpy()
    qa_store.GLOSSARY_EMBEDDINGS = glossary
    print("[INFO] Cached embeddings loaded successfully.")
    return True


def _save_embedding_cache(metadata: dict) -> None:
    try:
        os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True)
        tmp_path = f"{CACHE_PATH}.tmp"
        torch.save(
            {
                "metadata": metadata,
                "textbook": qa_store.TEXT_EMBEDDINGS,
                "glossary": qa_store.GLOSSARY_EMBEDDINGS,
            },
            tmp_path,
        )
        os.replace(tmp_path, CACHE_PATH)
        print(f"[INFO] Saved embedding cache to {CACHE_PATH}.")
    except Exception as e:
        print(f"[WARN] Could not save embedding cache: {e}")


def _compute_embeddings(textbook_texts: List[str], glossary_texts: List[str]) -> None:
    if textbook_texts:
        print("[INFO] Computing textbook embeddings from scratch...")
        qa_store.TEXT_EMBEDDINGS = embed_model.encode(
            textbook_texts,
            convert_to_tensor=True,
            show_progress_bar=False,
        )
    else:
        qa_store.TEXT_EMBEDDINGS = None

    if glossary_texts:
        print("[INFO] Computing glossary embeddings from scratch...")
        qa_store.GLOSSARY_EMBEDDINGS = embed_model.encode(
            glossary_texts,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
    else:
        qa_store.GLOSSARY_EMBEDDINGS = None


def admin_force_rebuild_cache() -> str:
    """
    Force recalculation of all embeddings and upload to cloud.
    Triggered by Teacher Panel button.
    """
    _refresh_runtime_data()

    print("[ADMIN] Rebuilding Textbook Embeddings...")
    textbook_texts = _entry_embedding_texts()
    print("[ADMIN] Rebuilding Glossary Embeddings...")
    glossary_texts = _glossary_embedding_texts()

    _compute_embeddings(textbook_texts, glossary_texts)

    print("[ADMIN] Saving to disk...")
    _save_embedding_cache(_cache_metadata(textbook_texts, glossary_texts))

    data_upload_status = sync_upload_data_tree()
    upload_status = sync_upload_cache()
    upload_status = f"{data_upload_status} | {upload_status}"
    return (
        "Cache rebuilt: "
        f"textbook={len(textbook_texts)}, glossary={len(glossary_texts)} | "
        f"{upload_status}"
    )


def _build_cached_embeddings() -> None:
    """
    Load matching cached embeddings, otherwise build and save a fresh cache.
    """
    textbook_texts = _entry_embedding_texts()
    glossary_texts = _glossary_embedding_texts()
    expected = _cache_metadata(textbook_texts, glossary_texts)

    if _load_cache_from_path(CACHE_FILE, expected, "local"):
        return

    if sync_download_cache(CACHE_DOWNLOAD_FILE):
        if _load_cache_from_path(CACHE_DOWNLOAD_FILE, expected, "downloaded"):
            shutil.copy(CACHE_DOWNLOAD_FILE, CACHE_FILE)
            return
        print("[WARN] Downloaded cache does not match current data; keeping local data authoritative.")

    _compute_embeddings(textbook_texts, glossary_texts)
    _save_embedding_cache(expected)


def _refresh_runtime_data() -> None:
    """
    Re-read local curriculum/manual QA/glossary files so runtime state matches disk.
    """
    with manual_qa_write_lock():
        load_curriculum()
        load_manual_qa()
        load_glossary()
        rebuild_combined_qa()


# -----------------------------
# Load data once at import time
# -----------------------------
if os.getenv("CHATBOT_SKIP_CLOUD_SYNC", "").strip().lower() not in {"1", "true", "yes", "on"}:
    sync_download_manual_qa()
else:
    print("[INFO] Skipping startup Manual Q&A cloud sync for a frozen local evaluation.")
load_curriculum()
load_manual_qa()
load_glossary()
rebuild_combined_qa()
_build_cached_embeddings()


# -----------------------------
# System prompt (Natural Science)
# -----------------------------
SYSTEM_PROMPT = (
    "ທ່ານແມ່ນຜູ້ຊ່ວຍເຫຼືອດ້ານວິທະຍາສາດທໍາມະຊາດ "
    "ສໍາລັບນັກຮຽນຊັ້ນ ມ.1-ມ.4. "
    "ຕອບແຕ່ພາສາລາວ ໃຫ້ຕອບສັ້ນໆ 2–3 ປະໂຫຍກ ແລະເຂົ້າໃຈງ່າຍ. "
    "ໃຫ້ອີງຈາກຂໍ້ມູນອ້າງອີງຂ້າງລຸ່ມນີ້ເທົ່ານັ້ນ. "
    "ຕ້ອງຕອບໃຫ້ຄົບທຸກສ່ວນຂອງຄຳຖາມ ແລະ ຫ້າມທວນຄຳຖາມ. "
    f"ຖ້າຂໍ້ມູນບໍ່ພຽງພໍ ໃຫ້ຕອບພຽງວ່າ: {SAFE_LAO_REFUSAL}"
)


# -----------------------------
# Helper: history formatting
# -----------------------------
def _format_history(history: Optional[List]) -> str:
    """
    Convert last few chat turns into a Lao conversation snippet
    to give the model context for follow-up questions.
    Gradio history format: [[user_msg, bot_msg], [user_msg, bot_msg], ...]
    """
    if not history:
        return ""

    recent = history[-3:]
    lines: List[str] = []
    for turn in recent:
        if not isinstance(turn, (list, tuple)) or len(turn) != 2:
            continue
        user_msg, bot_msg = turn
        lines.append(f"ນັກຮຽນ: {user_msg}")
        lines.append(f"ອາຈານ AI: {bot_msg}")

    if not lines:
        return ""

    return "\n".join(lines) + "\n\n"


def _to_ascii_digits(text: str) -> str:
    return (text or "").translate(_LAO_TO_ARABIC_DIGITS)


def _extract_user_text(item: Any) -> Optional[str]:
    if isinstance(item, (tuple, list)) and len(item) >= 1:
        user = item[0]
        if isinstance(user, str):
            return user.strip()
    if isinstance(item, dict) and item.get("role") == "user":
        content = item.get("content")
        if isinstance(content, str):
            return content.strip()
    return None


def _extract_bot_text(item: Any) -> Optional[str]:
    if isinstance(item, (tuple, list)) and len(item) >= 2:
        bot = item[1]
        if isinstance(bot, str):
            return bot.strip()
    if isinstance(item, dict) and item.get("role") == "assistant":
        content = item.get("content")
        if isinstance(content, str):
            return content.strip()
    return None


def _parse_grade_selection(text: str) -> Optional[str]:
    t = _to_ascii_digits((text or "").lower())
    compact = re.sub(r"\s+", "", t)
    m = re.search(r"(?:m|ມ)[\._-]?([1-4])", compact)
    if not m:
        return None
    return f"M_{int(m.group(1))}"


def _unit_number(unit: str) -> Optional[int]:
    m = re.search(r"(\d+)", unit or "")
    if not m:
        return None
    try:
        return int(m.group(1))
    except ValueError:
        return None


def _parse_unit_selection(text: str, history: Optional[List]) -> Optional[str]:
    t = _to_ascii_digits((text or "").strip())
    if not t:
        return None

    lower_t = t.lower()
    m = re.search(r"\bu[_\-\s]?(\d{1,2})\b", lower_t)
    if m:
        return f"U_{int(m.group(1)):02d}"

    m = re.search(r"(?:ບົດທີ|บทที่|chapter)\s*([0-9]{1,2})", lower_t)
    if m:
        return f"U_{int(m.group(1)):02d}"

    # Accept plain numbers only right after chapter menu was shown.
    if re.fullmatch(r"\d{1,2}", lower_t):
        recent_bot_msgs: List[str] = []
        if isinstance(history, list):
            for item in history[-3:]:
                bot_txt = _extract_bot_text(item)
                if bot_txt:
                    recent_bot_msgs.append(bot_txt)
        if any("ມີຢູ່" in msg and "ບົດ" in msg for msg in recent_bot_msgs):
            return f"U_{int(lower_t):02d}"

    return None


def _grade_display(grade: str) -> str:
    m = re.search(r"(\d+)", grade or "")
    return f"ມ.{m.group(1)}" if m else grade


def _load_unit_title_from_textbook(grade: str, unit: str) -> str:
    """
    Read the unit title directly from textbook.jsonl on disk.
    This covers cases where new curriculum files were added after app startup.
    """
    if not grade or not unit or grade == "LEGACY":
        return ""

    textbook_path = os.path.join(BASE_DIR, "data", grade, unit, "textbook.jsonl")
    if not os.path.exists(textbook_path):
        return ""

    try:
        with open(textbook_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                return (
                    str(obj.get("title") or "").strip()
                    or str(obj.get("chapter_title") or "").strip()
                )
    except Exception:
        return ""

    return ""


def _chapter_label(title: str, unit: str, chapter_num: Optional[int]) -> str:
    clean_title = re.sub(r"\s+", " ", (title or "").strip())
    if clean_title:
        if re.search(r"(?:ບົດທີ|บทที่|chapter)\s*\d+", clean_title, flags=re.IGNORECASE):
            return clean_title
        if chapter_num is not None:
            return f"ບົດທີ {chapter_num}: {clean_title}"
        return clean_title

    if chapter_num is not None:
        return f"ບົດທີ {chapter_num}"
    return unit


def _chapter_catalog_for_grade(grade: str) -> List[dict]:
    by_unit: dict = {}

    for e in qa_store.ENTRIES:
        if e.get("grade") != grade:
            continue
        unit = str(e.get("unit") or "").strip()
        if not unit or unit in by_unit:
            continue

        chapter_num: Optional[int] = None
        try:
            if e.get("chapter") is not None:
                chapter_num = int(e.get("chapter"))
        except Exception:
            chapter_num = None
        if chapter_num is None:
            chapter_num = _unit_number(unit)

        title = (
            str(e.get("title") or "").strip()
            or str(e.get("chapter_title") or "").strip()
        )
        if not title:
            title = _load_unit_title_from_textbook(grade, unit)
        by_unit[unit] = {
            "unit": unit,
            "chapter_num": chapter_num,
            "label": _chapter_label(title, unit, chapter_num),
        }

    # Ensure units that only exist in manual QA are still listed.
    for e in qa_store.MANUAL_QA_LIST:
        if e.get("grade") != grade:
            continue
        unit = str(e.get("unit") or "").strip()
        if not unit or unit in by_unit:
            continue
        unit_num = _unit_number(unit)
        title = _load_unit_title_from_textbook(grade, unit)
        by_unit[unit] = {
            "unit": unit,
            "chapter_num": unit_num,
            "label": _chapter_label(title, unit, unit_num),
        }

    items = list(by_unit.values())
    items.sort(key=lambda x: (x.get("chapter_num") is None, x.get("chapter_num") or 999, x["unit"]))
    return items


def _build_chapter_menu_reply(grade: str) -> str:
    chapters = _chapter_catalog_for_grade(grade)
    grade_name = _grade_display(grade)
    if not chapters:
        return (
            f"ຍັງບໍ່ພົບຂໍ້ມູນບົດຮຽນຂອງ {grade_name}.\n"
            "ກະລຸນາຖາມໂດຍກົງ ຫຼື ກວດສອບໄຟລ໌ textbook.jsonl."
        )

    lines = [item["label"] for item in chapters]
    menu = "\n".join(lines)
    return (
        f"ປຶ້ມແບບຮຽນ ວິທະຍາສາດທຳມະຊາດ ຊັ້ນ {grade_name} ມີຢູ່ {len(lines)} ບົດ:\n"
        f"{menu}\n\n"
        "ກະລຸນາເລືອກບົດທີ່ຕ້ອງການ (ຕົວຢ່າງ: ບົດທີ 5 ຫຼື U_05)."
    )


def _looks_like_chapter_menu(text: str) -> bool:
    return "\u0ea1\u0eb5\u0ea2\u0eb9\u0ec8" in (text or "") and "\u0e9a\u0ebb\u0e94" in (text or "")


def _latest_selected_grade(history: Optional[List]) -> Optional[str]:
    if not isinstance(history, list):
        return None
    for item in reversed(history):
        user_text = _extract_user_text(item)
        if user_text:
            grade = _parse_grade_selection(user_text)
            if grade:
                return grade

        bot_text = _extract_bot_text(item)
        if bot_text and _looks_like_chapter_menu(bot_text):
            grade = _parse_grade_selection(bot_text)
            if grade:
                return grade
    return None


def _grade_from_chapter_label(text: str, unit: Optional[str]) -> Optional[str]:
    if not text or not unit:
        return None

    candidate = re.sub(r"\s+", " ", text.strip())
    matches: List[str] = []
    grades = sorted(
        {
            str(e.get("grade") or "")
            for e in [*qa_store.ENTRIES, *qa_store.MANUAL_QA_LIST]
            if e.get("grade")
        },
        key=lambda g: int(re.search(r"\d+", g).group(0)) if re.search(r"\d+", g) else 999,
    )

    for grade in grades:
        for item in _chapter_catalog_for_grade(grade):
            if item.get("unit") != unit:
                continue
            label = re.sub(r"\s+", " ", str(item.get("label") or "").strip())
            if candidate == label or candidate.startswith(label) or label.startswith(candidate):
                matches.append(grade)

    unique_matches = sorted(set(matches))
    return unique_matches[0] if len(unique_matches) == 1 else None


def _collect_unit_questions(grade: str, unit: str) -> List[str]:
    seen = set()
    out: List[str] = []

    def add_question(q: str) -> None:
        question = (q or "").strip()
        if not question:
            return
        key = qa_store.normalize_question(question)
        if key in seen:
            return
        seen.add(key)
        out.append(question)

    # Manual QA first (teacher curated).
    for e in qa_store.MANUAL_QA_LIST:
        if e.get("grade") == grade and e.get("unit") == unit:
            add_question(str(e.get("q") or ""))

    # Then textbook auto-qa if present.
    for e in qa_store.ENTRIES:
        if e.get("grade") != grade or e.get("unit") != unit:
            continue
        for pair in e.get("qa", []) or []:
            if isinstance(pair, dict):
                add_question(str(pair.get("q") or ""))

    if GUIDE_SUGGESTION_LIMIT is None:
        return out
    return out[:GUIDE_SUGGESTION_LIMIT]


def _build_suggestion_reply(grade: str, unit: str) -> str:
    catalog = _chapter_catalog_for_grade(grade)
    label = next((c["label"] for c in catalog if c["unit"] == unit), unit)
    questions = _collect_unit_questions(grade, unit)

    if not questions:
        return (
            f"ຍັງບໍ່ພົບລາຍການຄຳຖາມແນະນຳສຳລັບ {label}.\n"
            "ທ່ານສາມາດຖາມຄຳຖາມໂດຍກົງໄດ້."
        )

    question_lines = "\n".join(f"- {q}" for q in questions)
    return (
        "ຄຳຖາມທີ່ຖາມແລ້ວໄດ້ຮັບຄຳຕອບຕົງຕາມປຶ້ມແບບຮຽນ:\n"
        f"{question_lines}\n\n"
        f"(ຫົວຂໍ້: {label})"
    )


def _guided_reply(message: str, history: Optional[List]) -> Optional[str]:
    grade = _parse_grade_selection(message)
    if grade:
        return _build_chapter_menu_reply(grade)

    unit = _parse_unit_selection(message, history)
    if not unit:
        return None

    selected_grade = _latest_selected_grade(history) or _grade_from_chapter_label(message, unit) or "M_1"

    chapter_units = {c["unit"] for c in _chapter_catalog_for_grade(selected_grade)}
    if unit not in chapter_units:
        return None
    return _build_suggestion_reply(selected_grade, unit)


# -----------------------------
# RAG: retrieve textbook context
# -----------------------------
def retrieve_context_details(
    question: str,
    max_entries: int = MAX_CONTEXT_ENTRIES,
    *,
    grade: Any = None,
    unit: Any = None,
    allow_global_fallback: bool = True,
    decision_trace: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
    """Return scoped textbook candidates for the production RAG route.

    Grade/unit constraints are strict for the first pass.  A global fallback is
    used only when scoped evidence is absent or substantially weaker than a
    strong global candidate.  Unscored arbitrary entries are never treated as
    retrieval evidence.
    """
    if decision_trace is not None:
        decision_trace.clear()
    entries = list(getattr(qa_store, "ENTRIES", None) or [])
    limit = min(max(int(max_entries or 0), 0), len(entries))
    if not entries or limit == 0:
        if decision_trace is not None:
            decision_trace.update(
                {
                    "scope_mode": scope_label(grade, unit),
                    "fallback_used": False,
                    "reason": "no_textbook_entries",
                }
            )
        return []

    if qa_store.TEXT_EMBEDDINGS is None:
        if decision_trace is not None:
            decision_trace.update(
                {
                    "scope_mode": scope_label(grade, unit),
                    "fallback_used": False,
                    "reason": "missing_similarity_scores",
                }
            )
        return []

    q_vec = embed_model.encode(
        question,
        convert_to_tensor=True,
        show_progress_bar=False,
    )
    sims = cos_sim(q_vec, qa_store.TEXT_EMBEDDINGS)[0]

    def rank_indices(indices: Sequence[int], count: int) -> List[Tuple[int, float]]:
        if not indices or count <= 0:
            return []
        candidate_indices = list(indices)
        candidate_scores = sims[candidate_indices]
        k = min(count, len(candidate_indices))
        values, local_indices = torch.topk(candidate_scores, k=k)
        return [
            (candidate_indices[int(local_idx)], float(score))
            for local_idx, score in zip(local_indices.tolist(), values.tolist())
        ]

    all_indices = list(range(len(entries)))
    scoped = bool(normalize_scope(grade) or normalize_scope(unit))
    scoped_indices = [
        idx for idx, entry in enumerate(entries) if item_in_scope(entry, grade, unit)
    ]
    ranked = rank_indices(scoped_indices if scoped else all_indices, limit)
    selected_scope = scope_label(grade, unit) if scoped else "global"
    fallback_used = False
    fallback_reason = ""

    # Controlled fallback: require a strong global candidate that improves on
    # the scoped top score by a meaningful margin.  This avoids the historical
    # behavior where a merely reordered cross-unit result still won.
    if scoped and allow_global_fallback:
        global_ranked = rank_indices(all_indices, limit)
        scoped_top = ranked[0][1] if ranked else None
        global_top = global_ranked[0][1] if global_ranked else None
        if not ranked and global_top is not None and global_top >= 0.78:
            ranked = global_ranked
            fallback_used = True
            fallback_reason = "no_scoped_entries_strong_global_candidate"
            selected_scope = "global_fallback"
        elif (
            scoped_top is not None
            and global_top is not None
            and scoped_top < 0.57
            and global_top >= 0.78
            and global_top - scoped_top >= 0.12
        ):
            ranked = global_ranked
            fallback_used = True
            fallback_reason = "weak_scoped_evidence_strong_global_candidate"
            selected_scope = "global_fallback"

    details: List[Dict[str, Any]] = []
    for rank, (entry_index, score) in enumerate(ranked, start=1):
        entry = entries[entry_index]
        details.append(
            {
                "rank": rank,
                "source_id": str(entry.get("id") or ""),
                "source_type": "textbook",
                "grade": str(entry.get("grade") or ""),
                "unit": str(entry.get("unit") or ""),
                "source_file": str(entry.get("_source_file") or ""),
                "similarity_score": score,
                "retrieved_text": str(entry.get("text") or ""),
                # Keep the prompt-facing fields separate so formatting stays
                # byte-for-byte compatible with the existing chatbot path.
                "chapter_title": str(entry.get("chapter_title") or ""),
                "section_title": str(entry.get("section_title") or ""),
                "title": str(entry.get("title") or entry.get("chapter") or ""),
                "section": str(entry.get("section") or ""),
                "retrieval_scope": selected_scope,
            }
        )
    if decision_trace is not None:
        decision_trace.update(
            {
                "requested_scope": scope_label(grade, unit),
                "scope_mode": selected_scope,
                "requested_grades": list(normalize_scope(grade)),
                "requested_units": list(normalize_scope(unit)),
                "scoped_candidate_count": len(scoped_indices)
                if scoped
                else len(entries),
                "fallback_used": fallback_used,
                "fallback_reason": fallback_reason,
                "reason": "ranked_candidates" if details else "no_ranked_candidates",
                "top_candidates": [
                    {
                        "rank": item["rank"],
                        "source_id": item["source_id"],
                        "grade": item["grade"],
                        "unit": item["unit"],
                        "similarity_score": item["similarity_score"],
                    }
                    for item in details
                ],
            }
        )
    return details


def format_retrieved_contexts(retrieved_contexts: List[Dict[str, Any]]) -> str:
    """Format structured production retrieval records for the SeaLLMs prompt."""
    context_blocks: List[str] = []
    for item in retrieved_contexts:
        source_type = str(item.get("source_type") or "textbook")
        component_labels = [
            str(value)
            for value in (item.get("component_labels") or [])
            if str(value)
        ]
        component_text = (
            f", ອົງປະກອບ {', '.join(component_labels)}"
            if component_labels
            else ""
        )
        header = (
            f"[ແຫຼ່ງ {source_type}, "
            f"ລະຫັດ {item.get('source_id','')}, "
            f"ຊັ້ນ {item.get('grade','')}, "
            f"ໜ່ວຍ {item.get('unit','')}, "
            f"ບົດ {item.get('chapter_title','')}, "
            f"ຫົວຂໍ້ {item.get('section_title','')}{component_text}]"
        )
        context_blocks.append(f"{header}\n{item.get('retrieved_text','')}")
    return "\n\n".join(context_blocks)


def retrieve_context(
    question: str,
    max_entries: int = MAX_CONTEXT_ENTRIES,
    *,
    grade: Any = None,
    unit: Any = None,
) -> str:
    """
    Embedding-based retrieval over textbook entries.
    Falls back to concatenated raw knowledge if embeddings are missing.
    """
    # Preserve the historical explicit-zero behavior (an empty context) while
    # still falling back to RAW_KNOWLEDGE when no textbook entries exist.
    if getattr(qa_store, "ENTRIES", None) and int(max_entries or 0) == 0:
        return ""
    details = retrieve_context_details(
        question,
        max_entries=max_entries,
        grade=grade,
        unit=unit,
    )
    if not details:
        return getattr(qa_store, "RAW_KNOWLEDGE", "")
    return format_retrieved_contexts(details)


# -----------------------------
# Glossary-based answering
# -----------------------------
def normalize_lao_text(text: str) -> str:
    """
    Clean Lao text for accurate matching.
    Removes punctuation and extra spaces.
    """
    if not text:
        return ""

    text = text.lower().strip()
    text = re.sub(r"[?.!,;։:'\"“”‘’]", "", text)
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def _record_source_match(
    match_trace: Optional[Dict[str, Any]],
    item: Optional[Dict[str, Any]],
    *,
    source_type: str,
    method: str,
    score: Optional[float] = None,
) -> None:
    """Record evaluation provenance without changing the returned answer."""
    if match_trace is None:
        return
    match_trace.update(
        {
            "accepted": True,
            "rejection_reason": "",
            "match_method": method,
            "rank": 1,
            "source_id": str((item or {}).get("id") or ""),
            "source_type": source_type,
            "grade": str((item or {}).get("grade") or ""),
            "unit": str((item or {}).get("unit") or ""),
            "source_file": str((item or {}).get("_source_file") or ""),
            "similarity_score": score,
            "matched_text": str(
                (item or {}).get("q")
                or (item or {}).get("term")
                or ""
            ),
        }
    )


def answer_from_glossary(
    message: str,
    *,
    grade: Any = None,
    unit: Any = None,
    match_trace: Optional[Dict[str, Any]] = None,
) -> Optional[str]:
    """
    Try to answer using the glossary index.
    Tier 1: Exact/Substring match (Sorted by Length to fix overlap bugs).
    Tier 2: Vector embedding match (Fallback).
    """
    if match_trace is not None:
        match_trace.clear()
        match_trace.update(
            {
                "accepted": False,
                "route": "glossary",
                "scope": scope_label(grade, unit),
                "requested_grades": list(normalize_scope(grade)),
                "requested_units": list(normalize_scope(unit)),
                "top_candidates": [],
            }
        )
    eligible, eligibility_reason = glossary_question_eligibility(message)
    if match_trace is not None:
        match_trace.update(
            {
                "question_eligible": eligible,
                "eligibility_reason": eligibility_reason,
            }
        )
    if not eligible:
        if match_trace is not None:
            match_trace["rejection_reason"] = eligibility_reason
        return None
    if not getattr(qa_store, "GLOSSARY", None):
        if match_trace is not None:
            match_trace["rejection_reason"] = "no_glossary_entries"
        return None

    norm_msg = normalize_lao_text(message)

    sorted_glossary = sorted(
        [
            item
            for item in qa_store.GLOSSARY
            if item_in_scope(item, grade, unit)
        ],
        key=lambda x: len(normalize_lao_text(x.get("term", ""))),
        reverse=True,
    )

    for item in sorted_glossary:
        term_raw = item.get("term", "")
        norm_term = normalize_lao_text(term_raw)
        if not norm_term:
            continue

        is_exact = (norm_msg == norm_term)
        is_substring = (norm_term in norm_msg) and (len(norm_msg) < len(norm_term) + 20)

        if is_exact or is_substring:
            _record_source_match(
                match_trace,
                item,
                source_type="glossary",
                method="exact" if is_exact else "substring",
                score=1.0 if is_exact else None,
            )
            definition = (item.get("definition", "") or "").strip()
            example = (item.get("example", "") or "").strip()
            if example:
                return f"{definition} ຕົວຢ່າງ: {example}"
            return definition

    if qa_store.GLOSSARY_EMBEDDINGS is None:
        if match_trace is not None:
            match_trace["rejection_reason"] = "glossary_embeddings_unavailable"
        return None

    q_emb = embed_model.encode(
        [message],
        convert_to_numpy=True,
        normalize_embeddings=True,
    )[0]

    sims = np.dot(qa_store.GLOSSARY_EMBEDDINGS, q_emb)
    ranked_indices = np.argsort(sims)[::-1].tolist()
    if match_trace is not None:
        match_trace["top_candidates"] = [
            {
                "rank": rank,
                "source_id": str(qa_store.GLOSSARY[idx].get("id") or ""),
                "grade": str(qa_store.GLOSSARY[idx].get("grade") or ""),
                "unit": str(qa_store.GLOSSARY[idx].get("unit") or ""),
                "similarity_score": float(sims[idx]),
                "in_scope": item_in_scope(
                    qa_store.GLOSSARY[idx], grade, unit
                ),
                "rejection_reason": ""
                if item_in_scope(qa_store.GLOSSARY[idx], grade, unit)
                else "scope_mismatch",
            }
            for rank, idx in enumerate(ranked_indices[:5], start=1)
        ]

    valid_indices = [
        idx
        for idx in ranked_indices
        if item_in_scope(qa_store.GLOSSARY[idx], grade, unit)
    ]
    if not valid_indices:
        if match_trace is not None:
            match_trace["rejection_reason"] = "no_candidates_in_scope"
        return None
    best_idx = valid_indices[0]
    best_sim = float(sims[best_idx])
    second_sim = float(sims[valid_indices[1]]) if len(valid_indices) > 1 else None
    accepted, acceptance_reason, margin = semantic_acceptance(
        best_sim,
        second_sim,
        threshold=GLOSSARY_SEMANTIC_THRESHOLD,
        minimum_margin=GLOSSARY_SEMANTIC_MARGIN,
    )
    if match_trace is not None:
        match_trace.update(
            {
                "semantic_threshold": GLOSSARY_SEMANTIC_THRESHOLD,
                "minimum_margin": GLOSSARY_SEMANTIC_MARGIN,
                "similarity_margin": margin,
            }
        )
    if not accepted:
        if match_trace is not None:
            match_trace["rejection_reason"] = acceptance_reason
        return None

    item = qa_store.GLOSSARY[best_idx]
    _record_source_match(
        match_trace,
        item,
        source_type="glossary",
        method="semantic",
        score=best_sim,
    )
    definition = (item.get("definition", "") or "").strip()
    example = (item.get("example", "") or "").strip()
    if example:
        return f"{definition} ຕົວຢ່າງ: {example}"
    return definition


# -----------------------------
# Prompt + LLM generation
# -----------------------------
def build_prompt(
    question: str,
    history: Optional[List] = None,
    *,
    retrieved_contexts: Optional[List[Dict[str, Any]]] = None,
) -> str:
    if retrieved_contexts is None:
        context = retrieve_context(question, max_entries=MAX_CONTEXT_ENTRIES)
    elif retrieved_contexts:
        context = format_retrieved_contexts(retrieved_contexts)
    else:
        # An explicit empty list means the retriever supplied no evidence.
        # Never silently replace it with the entire global knowledge base.
        context = "(ບໍ່ມີຂໍ້ມູນອ້າງອີງທີ່ຜ່ານເກນ)"
    history_block = _format_history(history)

    return f"""{SYSTEM_PROMPT}

{history_block}ຂໍ້ມູນອ້າງອີງ:
{context}

ຄຳຖາມ: {question}

ຄຳຕອບດ້ວຍພາສາລາວ:"""


class _StopOnGeneratedTokenSequences(StoppingCriteria):
    """Stop batch-size-one generation after prompt scaffolding reappears."""

    def __init__(
        self,
        prompt_length: int,
        token_sequences: Sequence[Sequence[int]],
    ) -> None:
        self.prompt_length = int(prompt_length)
        self.token_sequences = [
            [int(token) for token in sequence]
            for sequence in token_sequences
            if sequence
        ]

    def __call__(self, input_ids: Any, scores: Any, **kwargs: Any) -> bool:
        del scores, kwargs
        generated = input_ids[0, self.prompt_length :].tolist()
        return any(
            len(generated) >= len(sequence)
            and generated[-len(sequence) :] == sequence
            for sequence in self.token_sequences
        )


def build_evidence_only_retry_prompt(
    question: str,
    retrieved_contexts: Sequence[Mapping[str, Any]],
) -> str:
    """Build the deliberately short, deterministic second-attempt prompt."""
    context = format_retrieved_contexts([dict(item) for item in retrieved_contexts])
    return f"""ຄຳສັ່ງສຳລັບການຕອບຄືນໃໝ່:
- ໃຊ້ສະເພາະຫຼັກຖານທີ່ໃຫ້ໄວ້
- ຕອບເປັນພາສາລາວເທົ່ານັ້ນ
- ຕອບກົງຄຳຖາມ 1 ຫາ 3 ປະໂຫຍກສັ້ນໆ
- ຫ້າມໃຊ້ຫົວຂໍ້ ຫ້າມທວນຄຳຖາມ
- ຫ້າມໃຫ້ຕົວຢ່າງ ນອກຈາກຄຳຖາມຮ້ອງຂໍ
- ຫ້າມເພີ່ມຄຳອະທິບາຍທີ່ບໍ່ກ່ຽວຂ້ອງ
- ຫ້າມກ່າວວ່າຂໍ້ມູນບໍ່ພຽງພໍ ເມື່ອຫຼັກຖານມີຄຳຕອບ

ຫຼັກຖານ:
{context}

ຄຳຖາມ: {question}

ຄຳຕອບ:"""


def _split_answer_sentences(text: str) -> List[str]:
    clean = re.sub(r"\s+", " ", str(text or "")).strip()
    if not clean:
        return []
    chunks = re.split(r"(?<=[.!?…。])\s+|\n+", clean)
    return [chunk.strip() for chunk in chunks if chunk.strip()]


def _context_answer_payload(
    context: Mapping[str, Any],
    question: str,
) -> str:
    """Extract an answer-bearing span without copying a stored question."""
    text = str(context.get("retrieved_text") or context.get("text") or "").strip()
    if not text:
        return ""
    for marker in ("ຄຳຕອບ:", "ຄໍາຕອບ:", "ນິຍາມ:"):
        if marker in text:
            text = text.split(marker, 1)[1]
            break
    text = re.split(r"\n\s*(?:ຄຳຖາມ|ຄໍາຖາມ)\s*:", text, maxsplit=1)[0]
    text = re.sub(r"\s+", " ", text).strip(" -•\t\r\n")
    if not text:
        return ""

    if str(context.get("source_type") or "") == "textbook":
        question_tokens = set(significant_tokens(question))
        sentences = _split_answer_sentences(text)
        if sentences:
            ranked = sorted(
                enumerate(sentences),
                key=lambda pair: (
                    len(
                        question_tokens
                        & set(significant_tokens(pair[1]))
                    ),
                    -pair[0],
                ),
                reverse=True,
            )
            best_index = ranked[0][0]
            chosen = sentences[best_index : best_index + 2]
            text = " ".join(chosen).strip()
    return text


def _formula_from_payload(text: str) -> str:
    matches = re.findall(
        r"(?:1\s*/\s*R|[IUR])\s*=\s*[^.;\n]+",
        str(text or ""),
        flags=re.IGNORECASE,
    )
    if matches:
        return matches[-1].strip(" :")
    if ":" in str(text or ""):
        return str(text).rsplit(":", 1)[-1].strip()
    return str(text or "").strip()


def build_extractive_evidence_answer(
    question: str,
    retrieved_contexts: Sequence[Mapping[str, Any]],
    *,
    compound: bool = False,
) -> str:
    """Delegate narrow evidence extraction to the pure repair policy."""
    return _build_repaired_extractive_answer(
        question, retrieved_contexts, compound=compound
    )


def generate_answer(
    question: str,
    history: Optional[List] = None,
    *,
    retrieved_contexts: Optional[List[Dict[str, Any]]] = None,
    validation_trace: Optional[Dict[str, Any]] = None,
    allow_extractive_fallback: bool = False,
    evidence_confidence: Optional[Mapping[str, Any]] = None,
    compound: bool = False,
) -> str:
    """Generate from evidence, use a short retry, then a gated extraction."""
    if validation_trace is not None:
        validation_trace.clear()
        validation_trace["attempts"] = []

    contexts = list(retrieved_contexts or [])
    evidence_text = "\n".join(
        str(item.get("retrieved_text") or item.get("text") or "")
        for item in contexts
    )

    def generate_once(
        *, attempt_type: str
    ) -> Tuple[str, int, bool, str, List[str]]:
        evidence_only_retry = attempt_type == "evidence_only_retry"
        prompt = (
            build_evidence_only_retry_prompt(question, contexts)
            if evidence_only_retry
            else build_prompt(
                question,
                history,
                retrieved_contexts=contexts,
            )
        )
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        token_budget = 96 if evidence_only_retry else MAX_GENERATED_TOKENS
        generation_kwargs: Dict[str, Any] = {
            "max_new_tokens": token_budget,
            "do_sample": False if evidence_only_retry else GENERATION_DO_SAMPLE,
        }
        stop_sequences = [
            tokenizer.encode(marker, add_special_tokens=False)
            for marker in (
                "\n\nຄຳຖາມ:",
                "\nຄຳຖາມ:",
                " ຄຳຖາມ:",
                "ຄຳຖາມ:",
            )
        ]
        generation_kwargs["stopping_criteria"] = StoppingCriteriaList(
            [
                _StopOnGeneratedTokenSequences(
                    int(inputs["input_ids"].shape[1]),
                    stop_sequences,
                )
            ]
        )
        if evidence_only_retry:
            generation_kwargs.update(
                {
                    "repetition_penalty": 1.15,
                    "no_repeat_ngram_size": 4,
                }
            )
        with torch.no_grad():
            outputs = model.generate(**inputs, **generation_kwargs)

        generated_ids = outputs[0][inputs["input_ids"].shape[1] :]
        generated_count = int(generated_ids.shape[0])
        eos_ids = getattr(tokenizer, "eos_token_id", None)
        if eos_ids is None:
            eos_values: set = set()
        elif isinstance(eos_ids, (list, tuple, set)):
            eos_values = {int(value) for value in eos_ids}
        else:
            eos_values = {int(eos_ids)}
        last_token = int(generated_ids[-1]) if generated_count else None
        limit_reached = bool(
            generated_count >= token_budget
            and (last_token is None or last_token not in eos_values)
        )
        decoded = tokenizer.decode(
            generated_ids, skip_special_tokens=True
        ).strip()
        prepared, safety_transformations = prepare_generated_answer_text(decoded)
        return (
            prepared,
            generated_count,
            limit_reached,
            decoded,
            safety_transformations,
        )

    for attempt_number, attempt_type in enumerate(
        ("normal", "evidence_only_retry"), start=1
    ):
        (
            answer,
            generated_count,
            limit_reached,
            raw_generated_text,
            safety_transformations,
        ) = generate_once(attempt_type=attempt_type)
        validation = validate_generated_answer(
            question,
            answer,
            evidence_text=evidence_text,
            token_limit_reached=limit_reached,
        )
        if validation.get("valid") and evidence_text and answer:
            semantic_vectors = embed_model.encode(
                [answer, question],
                convert_to_numpy=True,
                normalize_embeddings=True,
                show_progress_bar=False,
            )
            semantic_similarity = float(
                np.dot(semantic_vectors[0], semantic_vectors[1])
            )
            validation["semantic_question_similarity"] = semantic_similarity
            validation[
                "minimum_semantic_question_similarity"
            ] = OUTPUT_MIN_QUESTION_SEMANTIC_SIMILARITY
            if (
                answer != SAFE_LAO_REFUSAL
                and semantic_similarity
                < OUTPUT_MIN_QUESTION_SEMANTIC_SIMILARITY
            ):
                validation["valid"] = False
                validation.setdefault("reasons", []).append(
                    "low_semantic_relevance_to_question"
                )
        attempt_record = {
            "attempt": attempt_number,
            "attempt_type": attempt_type,
            "generated_token_count": generated_count,
            "token_budget": 96 if attempt_type == "evidence_only_retry" else MAX_GENERATED_TOKENS,
            "repair_attempt": attempt_type == "evidence_only_retry",
            "generated_text": answer,
            "raw_generated_text": raw_generated_text,
            "safety_transformations": safety_transformations,
            **validation,
        }
        if validation_trace is not None:
            validation_trace["attempts"].append(attempt_record)
        if validation.get("valid"):
            if validation_trace is not None:
                validation_trace.update(
                    {
                        "retry_count": attempt_number - 1,
                        "final_action": "accepted_generated_answer",
                        "final_attempt_type": attempt_type,
                        "returned_safe_refusal": False,
                        "used_extractive_fallback": False,
                    }
                )
            return answer

    fallback_allowed, fallback_reason = strong_evidence_for_extractive_fallback(
        contexts,
        evidence_confidence,
        compound=compound,
    )
    if allow_extractive_fallback and fallback_allowed:
        extractive_answer = build_extractive_evidence_answer(
            question,
            contexts,
            compound=compound,
        )
        extractive_validation = validate_generated_answer(
            question,
            extractive_answer,
            evidence_text=evidence_text,
            token_limit_reached=False,
        )
        if extractive_validation.get("valid") and extractive_answer:
            if validation_trace is not None:
                validation_trace.update(
                    {
                        "retry_count": 1,
                        "final_action": "accepted_extractive_fallback",
                        "final_attempt_type": "extractive_fallback",
                        "returned_safe_refusal": False,
                        "used_extractive_fallback": True,
                        "fallback_reason": fallback_reason,
                        "extractive_fallback": {
                            "attempt_type": "extractive_fallback",
                            "generated_text": extractive_answer,
                            **extractive_validation,
                        },
                    }
                )
            return extractive_answer
        fallback_reason = "extractive_fallback_failed_validation"
        if validation_trace is not None:
            validation_trace["extractive_fallback"] = {
                "attempt_type": "extractive_fallback",
                "generated_text": extractive_answer,
                **extractive_validation,
            }

    consensus = build_consensus_structure_fallback(
        question, contexts, evidence_confidence
    )
    if allow_extractive_fallback and consensus.get("eligible"):
        consensus_answer = str(consensus.get("answer") or "")
        consensus_validation = validate_generated_answer(
            question,
            consensus_answer,
            evidence_text=evidence_text,
            token_limit_reached=False,
        )
        if consensus_answer and consensus_validation.get("valid"):
            if validation_trace is not None:
                validation_trace.update(
                    {
                        "retry_count": 1,
                        "final_action": "accepted_consensus_structure_fallback",
                        "final_attempt_type": "consensus_structure_fallback",
                        "returned_safe_refusal": False,
                        "used_extractive_fallback": True,
                        "used_consensus_structure_fallback": True,
                        "fallback_reason": "consensus_supported_structure_evidence",
                        "standard_extractive_fallback_reason": fallback_reason,
                        "consensus_structure_fallback": {
                            **consensus,
                            **consensus_validation,
                        },
                    }
                )
            return consensus_answer

    if validation_trace is not None:
        validation_trace.update(
            {
                "retry_count": 1,
                "final_action": "safe_refusal_after_failed_validation",
                "final_attempt_type": "safe_refusal",
                "returned_safe_refusal": True,
                "used_extractive_fallback": False,
                "used_consensus_structure_fallback": False,
                "fallback_reason": fallback_reason,
                "consensus_structure_fallback": consensus,
            }
        )
    return SAFE_LAO_REFUSAL


# -----------------------------
# QA lookup (exact + fuzzy + partial + embeddings)
# -----------------------------
def _ensure_qa_embeddings() -> None:
    """
    Build (or rebuild) embeddings for ALL_QA_KNOWLEDGE questions.
    Cached on qa_store to avoid recomputing every chat turn.
    """
    while True:
        knowledge = getattr(qa_store, "ALL_QA_KNOWLEDGE", None) or []
        corpus_version = getattr(qa_store, "QA_CORPUS_VERSION", 0)
        n = len(knowledge)

        if not knowledge:
            qa_store.QA_Q_EMBEDDINGS = None
            qa_store.QA_Q_EMBED_N = 0
            qa_store.QA_Q_EMBED_VERSION = corpus_version
            return

        if (
            getattr(qa_store, "QA_Q_EMBEDDINGS", None) is not None
            and getattr(qa_store, "QA_Q_EMBED_N", 0) == n
            and getattr(qa_store, "QA_Q_EMBED_VERSION", -1) == corpus_version
        ):
            return

        texts = [(item.get("norm_q", "") or "") for item in knowledge]
        embeddings = embed_model.encode(
            texts,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )

        # A reload may complete while encoding is running. Never publish vectors
        # unless they still match the current corpus generation and list.
        if (
            corpus_version != getattr(qa_store, "QA_CORPUS_VERSION", 0)
            or knowledge is not getattr(qa_store, "ALL_QA_KNOWLEDGE", None)
        ):
            continue

        qa_store.QA_Q_EMBEDDINGS = embeddings
        qa_store.QA_Q_EMBED_N = n
        qa_store.QA_Q_EMBED_VERSION = corpus_version
        return


def _qa_embedding_snapshot():
    """Return a corpus/embedding pair from the same Q&A generation."""
    while True:
        _ensure_qa_embeddings()
        corpus_version = getattr(qa_store, "QA_CORPUS_VERSION", 0)
        knowledge = list(getattr(qa_store, "ALL_QA_KNOWLEDGE", None) or [])
        embeddings = getattr(qa_store, "QA_Q_EMBEDDINGS", None)
        if (
            corpus_version == getattr(qa_store, "QA_CORPUS_VERSION", 0)
            and getattr(qa_store, "QA_Q_EMBED_VERSION", -1) == corpus_version
            and (embeddings is None or len(embeddings) == len(knowledge))
        ):
            return knowledge, embeddings


def _escape_md_cell(text: str) -> str:
    return (text or "").replace("|", "\\|").replace("\n", "<br>")


def _build_markdown_table(headers: List[str], rows: List[List[str]]) -> str:
    header_line = "| " + " | ".join(_escape_md_cell(h) for h in headers) + " |"
    sep_line = "| " + " | ".join("---" for _ in headers) + " |"
    row_lines = [
        "| " + " | ".join(_escape_md_cell(cell) for cell in row) + " |"
        for row in rows
    ]
    return "\n".join([header_line, sep_line, *row_lines])


def _escape_answer_text_html(text: str) -> str:
    return html.escape(text or "").replace("\n", "<br>")


def _format_answer_text_html(text: str) -> str:
    safe_text = _escape_answer_text_html(text)
    safe_text = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", safe_text)
    return safe_text


def _render_plain_text_block(text: str) -> str:
    return (
        '<div class="teacher-answer-preview teacher-answer-preview--text">'
        f"{_format_answer_text_html(text)}"
        "</div>"
    )


def _render_html_answer_component(answer: str) -> gr.HTML:
    table_pattern = re.compile(r"(?is)<table[\s>].*?</table>")
    html_parts = ['<div class="teacher-answer-preview chat-answer-rich">']
    last_end = 0

    for match in table_pattern.finditer(answer or ""):
        prefix = (answer or "")[last_end:match.start()].strip()
        if prefix:
            html_parts.append(_render_plain_text_block(prefix))
        html_parts.append(match.group(0))
        last_end = match.end()

    suffix = (answer or "")[last_end:].strip()
    if suffix:
        html_parts.append(_render_plain_text_block(suffix))

    html_parts.append("</div>")
    return gr.HTML(value="".join(html_parts))


def _strip_html_tags(text: str) -> str:
    text = re.sub(r"(?i)<br\s*/?>", "\n", text or "")
    text = re.sub(r"(?is)<[^>]+>", "", text)
    return html.unescape(text).strip()


def _parse_html_span(attrs: str, name: str) -> int:
    match = re.search(rf"{name}\s*=\s*[\"']?(\d+)", attrs or "", flags=re.IGNORECASE)
    if not match:
        return 1
    return max(int(match.group(1)), 1)


def _parse_html_table_rows(section_html: str) -> List[List[dict]]:
    rows: List[List[dict]] = []
    for row_match in re.finditer(r"(?is)<tr\b[^>]*>(.*?)</tr>", section_html or ""):
        row_html = row_match.group(1)
        row_cells: List[dict] = []
        for cell_match in re.finditer(r"(?is)<(th|td)\b([^>]*)>(.*?)</\1>", row_html):
            tag = cell_match.group(1).lower()
            attrs = cell_match.group(2) or ""
            inner_html = cell_match.group(3) or ""
            row_cells.append(
                {
                    "tag": tag,
                    "text": _strip_html_tags(inner_html),
                    "rowspan": _parse_html_span(attrs, "rowspan"),
                    "colspan": _parse_html_span(attrs, "colspan"),
                }
            )
        if row_cells:
            rows.append(row_cells)
    return rows


def _build_html_table_grid(rows: List[List[dict]]) -> List[List[Optional[dict]]]:
    grid: List[List[Optional[dict]]] = []
    next_id = 1

    for row_idx, row in enumerate(rows):
        while len(grid) <= row_idx:
            grid.append([])

        col_idx = 0
        for cell in row:
            while col_idx < len(grid[row_idx]) and grid[row_idx][col_idx] is not None:
                col_idx += 1

            cell_ref = {"id": next_id, "text": cell.get("text", "")}
            next_id += 1

            rowspan = max(int(cell.get("rowspan", 1) or 1), 1)
            colspan = max(int(cell.get("colspan", 1) or 1), 1)

            for fill_row in range(row_idx, row_idx + rowspan):
                while len(grid) <= fill_row:
                    grid.append([])
                while len(grid[fill_row]) < col_idx + colspan:
                    grid[fill_row].append(None)
                for fill_col in range(col_idx, col_idx + colspan):
                    grid[fill_row][fill_col] = cell_ref

            col_idx += colspan

    width = max((len(row) for row in grid), default=0)
    for row in grid:
        row.extend([None] * (width - len(row)))
    return grid


def _flatten_html_table_headers(header_rows: List[List[dict]], total_cols: int) -> List[str]:
    if total_cols <= 0:
        return []
    if not header_rows:
        return [f"Column {idx + 1}" for idx in range(total_cols)]

    grid = _build_html_table_grid(header_rows)
    for row in grid:
        row.extend([None] * (total_cols - len(row)))

    headers: List[str] = []
    for col_idx in range(total_cols):
        parts: List[str] = []
        last_id = None
        for row in grid:
            cell = row[col_idx] if col_idx < len(row) else None
            if cell and cell.get("id") != last_id and cell.get("text"):
                parts.append(str(cell["text"]))
                last_id = cell.get("id")
        headers.append(" - ".join(parts) if parts else f"Column {col_idx + 1}")
    return headers


def _expand_html_table_body_rows(body_rows: List[List[dict]], total_cols: int) -> List[List[str]]:
    expanded_rows: List[List[str]] = []
    pending_rowspans = [0] * total_cols

    for row in body_rows:
        current_row = [""] * total_cols
        new_rowspans = [0] * total_cols
        col_idx = 0

        for cell in row:
            colspan = max(int(cell.get("colspan", 1) or 1), 1)
            rowspan = max(int(cell.get("rowspan", 1) or 1), 1)

            while col_idx < total_cols and pending_rowspans[col_idx] > 0:
                col_idx += 1

            while (
                col_idx < total_cols
                and any(pending_rowspans[c] > 0 for c in range(col_idx, min(col_idx + colspan, total_cols)))
            ):
                col_idx += 1
                while col_idx < total_cols and pending_rowspans[col_idx] > 0:
                    col_idx += 1

            if col_idx >= total_cols:
                break

            current_row[col_idx] = str(cell.get("text", ""))
            end_col = min(col_idx + colspan, total_cols)
            if rowspan > 1:
                for span_col in range(col_idx, end_col):
                    new_rowspans[span_col] = max(new_rowspans[span_col], rowspan - 1)

            col_idx = end_col

        expanded_rows.append(current_row)
        pending_rowspans = [max(value - 1, 0) for value in pending_rowspans]
        for idx in range(total_cols):
            pending_rowspans[idx] = max(pending_rowspans[idx], new_rowspans[idx])

    return expanded_rows


def _html_table_to_markdown(table_html: str) -> str:
    thead_match = re.search(r"(?is)<thead\b[^>]*>(.*?)</thead>", table_html or "")
    tbody_match = re.search(r"(?is)<tbody\b[^>]*>(.*?)</tbody>", table_html or "")

    header_rows = _parse_html_table_rows(thead_match.group(1)) if thead_match else []
    if tbody_match:
        body_rows = _parse_html_table_rows(tbody_match.group(1))
    else:
        all_rows = _parse_html_table_rows(table_html or "")
        if not header_rows:
            split_idx = 0
            for row in all_rows:
                if all(cell.get("tag") == "th" for cell in row):
                    split_idx += 1
                else:
                    break
            header_rows = all_rows[:split_idx]
            body_rows = all_rows[split_idx:]
        else:
            body_rows = all_rows[len(header_rows):]

    header_grid = _build_html_table_grid(header_rows) if header_rows else []
    header_width = max((len(row) for row in header_grid), default=0)
    body_width = max(
        (sum(max(int(cell.get("colspan", 1) or 1), 1) for cell in row) for row in body_rows),
        default=0,
    )
    total_cols = max(header_width, body_width)
    if total_cols == 0:
        return ""

    headers = _flatten_html_table_headers(header_rows, total_cols)
    rows = _expand_html_table_body_rows(body_rows, total_cols)
    return _build_markdown_table(headers, rows)


def _convert_html_tables_to_markdown(answer: str) -> str:
    text = (answer or "").strip()
    if "<table" not in text.lower():
        return text

    table_pattern = re.compile(r"(?is)<table[\s>].*?</table>")
    parts: List[str] = []
    last_end = 0

    for match in table_pattern.finditer(text):
        prefix = text[last_end:match.start()].strip()
        if prefix:
            parts.append(prefix)

        table_md = _html_table_to_markdown(match.group(0))
        if table_md:
            parts.append(table_md)

        last_end = match.end()

    suffix = text[last_end:].strip()
    if suffix:
        parts.append(suffix)

    return "\n\n".join(parts).strip()


def _format_answer_as_table_if_needed(answer: str) -> str:
    """
    Convert repeated 'Key: Value' blocks into a markdown table.
    Keeps original text when the structure is not table-like.
    """
    text = _convert_html_tables_to_markdown((answer or "").strip())
    if not text:
        return text

    lines = [ln.rstrip() for ln in text.splitlines()]

    # Already markdown-table-like; do nothing.
    pipe_lines = [ln for ln in lines if re.match(r"^\s*\|.*\|\s*$", ln)]
    if len(pipe_lines) >= 2:
        return text

    kv_pattern = re.compile(r"^\s*([^:\n：]{1,80})\s*[：:]\s*(.+?)\s*$")
    sep_pattern = re.compile(r"^\s*[-_]{3,}\s*$")

    first_kv_idx = next((i for i, ln in enumerate(lines) if kv_pattern.match(ln)), None)
    if first_kv_idx is None:
        return text

    prefix_lines = lines[:first_kv_idx]
    working_lines = lines[first_kv_idx:]

    blocks: List[dict] = []
    current: dict = {}
    current_order: List[str] = []
    consumed = 0

    for idx, line in enumerate(working_lines):
        stripped = line.strip()
        m = kv_pattern.match(line)
        if m:
            key = m.group(1).strip()
            value = m.group(2).strip()
            if key not in current:
                current_order.append(key)
            current[key] = value
            consumed = idx + 1
            continue

        if sep_pattern.match(stripped):
            if current:
                blocks.append({"order": current_order[:], "values": current.copy()})
                current = {}
                current_order = []
            consumed = idx + 1
            continue

        if not stripped:
            consumed = idx + 1
            continue

        # Non-table content starts here.
        break

    if current:
        blocks.append({"order": current_order[:], "values": current.copy()})

    if len(blocks) < 2:
        return text

    headers = blocks[0]["order"]
    if len(headers) < 2:
        return text

    if any(block["order"] != headers for block in blocks[1:]):
        return text

    rows = [[block["values"].get(h, "") for h in headers] for block in blocks]
    table_md = _build_markdown_table(headers, rows)

    suffix_lines = working_lines[consumed:]

    out_parts: List[str] = []
    if any(ln.strip() for ln in prefix_lines):
        out_parts.append("\n".join(prefix_lines).strip())
    out_parts.append(table_md)
    if any(ln.strip() for ln in suffix_lines):
        out_parts.append("\n".join(suffix_lines).strip())

    return "\n\n".join(out_parts).strip()


def _format_answer_for_chat(answer: str) -> Any:
    text = (answer or "").strip()
    if re.search(r"(?is)<table[\s>].*?</table>", text):
        return _render_html_answer_component(text)
    return _format_answer_as_table_if_needed(text)


def _qa_terms(norm_text: str) -> List[str]:
    return [t for t in (norm_text or "").split(" ") if len(t) > 1]


def _is_specific_qa_text(norm_text: str, terms: List[str]) -> bool:
    return len(terms) >= 2 or len((norm_text or "").strip()) >= MIN_QA_QUERY_CHARS


def _same_qa_context(item: Dict[str, Any], grade: Any, unit: Any) -> bool:
    return bool(normalize_scope(grade) or normalize_scope(unit)) and item_in_scope(
        item, grade, unit
    )


def _qa_answer_key(answer: Any) -> str:
    return re.sub(r"\s+", " ", str(answer or "").strip())


def _qa_context_key(grade: Any, unit: Any, norm_q: str) -> Optional[Tuple[str, str, str]]:
    grades = normalize_scope(grade)
    units = normalize_scope(unit)
    if len(grades) != 1 or len(units) != 1 or not norm_q:
        return None
    return (grades[0], units[0], norm_q)


def _ordered_qa_items(grade: Any = None, unit: Any = None) -> List[Dict[str, Any]]:
    knowledge = list(getattr(qa_store, "ALL_QA_KNOWLEDGE", None) or [])
    if not (normalize_scope(grade) or normalize_scope(unit)):
        return knowledge
    return [item for item in knowledge if item_in_scope(item, grade, unit)]


def _ordered_qa_indices(
    grade: Any = None,
    unit: Any = None,
    knowledge: Optional[List[Dict[str, Any]]] = None,
) -> List[int]:
    knowledge = (
        list(knowledge)
        if knowledge is not None
        else list(getattr(qa_store, "ALL_QA_KNOWLEDGE", None) or [])
    )
    indices = list(range(len(knowledge)))
    if not (normalize_scope(grade) or normalize_scope(unit)):
        return indices
    return [
        idx for idx in indices if item_in_scope(knowledge[idx], grade, unit)
    ]


def _ambiguous_exact_reply(items: List[Dict[str, Any]]) -> str:
    options: List[str] = []
    seen = set()
    for item in items:
        grade = str(item.get("grade") or "").strip()
        unit = str(item.get("unit") or "").strip()
        label = f"{grade}/{unit}".strip("/")
        if not label or label in seen:
            continue
        seen.add(label)
        options.append(label)
        if len(options) >= 8:
            break

    option_text = ", ".join(options)
    if len(seen) < len({(str(i.get("grade") or ""), str(i.get("unit") or "")) for i in items}):
        option_text += ", ..."
    return (
        "ຄຳຖາມນີ້ພົບໃນຫຼາຍບົດຮຽນ ແລະ ມີຄຳຕອບຕ່າງກັນ. "
        f"ກະລຸນາເລືອກຊັ້ນ/ບົດຮຽນກ່ອນ ຫຼື ລະບຸບົດຮຽນໃນຄຳຖາມ. ({option_text})"
    )


def _exact_qa_match_details(
    norm_q: str,
    grade: Any = None,
    unit: Any = None,
) -> Tuple[Optional[str], bool, Optional[Dict[str, Any]]]:
    context_key = _qa_context_key(grade, unit, norm_q)
    if context_key:
        item = getattr(qa_store, "QA_INDEX_BY_CONTEXT", {}).get(context_key)
        if item and item.get("a"):
            return str(item.get("a") or ""), False, item

    items = list(getattr(qa_store, "QA_ITEMS_BY_NORM", {}).get(norm_q, []))
    scoped = bool(normalize_scope(grade) or normalize_scope(unit))
    if scoped:
        items = [item for item in items if item_in_scope(item, grade, unit)]
    if not items:
        if scoped:
            return None, False, None
        legacy_answer = getattr(qa_store, "QA_INDEX", {}).get(norm_q)
        return (
            (str(legacy_answer), False, None)
            if legacy_answer
            else (None, False, None)
        )

    distinct_answers = {_qa_answer_key(item.get("a")) for item in items if item.get("a")}
    if len(distinct_answers) <= 1:
        return str(items[0].get("a") or ""), False, items[0]

    return None, True, None


def _exact_qa_match(
    norm_q: str,
    grade: Optional[str] = None,
    unit: Optional[str] = None,
) -> Tuple[Optional[str], bool]:
    answer, ambiguous, _item = _exact_qa_match_details(norm_q, grade, unit)
    return answer, ambiguous


def answer_from_qa(
    question: str,
    grade: Any = None,
    unit: Any = None,
    *,
    match_trace: Optional[Dict[str, Any]] = None,
) -> Optional[Any]:
    """
    Goal: match BOTH full questions and short "topic/keyword" queries safely.

    Priority:
    1) Exact match (highest precision)
    2) Partial/substring match
    3) Improved fuzzy match (overlap + coverage)
    4) Embedding similarity (semantic match; conservative thresholds)

    Safety:
    - If the user types only 1 keyword, do NOT force Q&A (let Glossary handle).
    """
    if match_trace is not None:
        match_trace.clear()
        match_trace.update(
            {
                "accepted": False,
                "route": "manual_qa",
                "scope": scope_label(grade, unit),
                "requested_grades": list(normalize_scope(grade)),
                "requested_units": list(normalize_scope(unit)),
                "top_candidates": [],
            }
        )
    if not question or not question.strip():
        if match_trace is not None:
            match_trace["rejection_reason"] = "empty_question"
        return None

    raw_norm = qa_store.normalize_question(question)
    norm_q = normalize_lao_text(raw_norm)
    if not norm_q:
        if match_trace is not None:
            match_trace["rejection_reason"] = "empty_normalized_question"
        return None

    q_terms = _qa_terms(norm_q)
    if not _is_specific_qa_text(norm_q, q_terms):
        if match_trace is not None:
            match_trace["rejection_reason"] = "question_not_specific_enough"
        return None

    candidate_items = _ordered_qa_items(grade, unit)
    if match_trace is not None:
        match_trace["scope_candidate_count"] = len(candidate_items)
    if not candidate_items:
        if match_trace is not None:
            match_trace["rejection_reason"] = "no_candidates_in_scope"
        return None

    # 1) Exact match (for sufficiently specific questions only)
    exact_answer, ambiguous_exact, exact_item = _exact_qa_match_details(
        norm_q, grade, unit
    )
    if exact_answer:
        source_type = "manual_qa"
        if str((exact_item or {}).get("source") or "").lower() != "manual":
            source_type = "textbook_auto_qa"
        _record_source_match(
            match_trace,
            exact_item,
            source_type=source_type,
            method="exact",
            score=1.0,
        )
        return _format_answer_for_chat(exact_answer)
    if ambiguous_exact and not (grade and unit):
        if match_trace is not None:
            match_trace["rejection_reason"] = "ambiguous_exact_match"
        return _ambiguous_exact_reply(getattr(qa_store, "QA_ITEMS_BY_NORM", {}).get(norm_q, []))

    # 2) Partial/substring match
    best_sub_score = 0.0
    best_sub_answer: Optional[str] = None
    best_sub_item: Optional[Dict[str, Any]] = None

    for item in candidate_items:
        stored = item.get("norm_q", "") or ""
        if not stored:
            continue
        stored_terms = _qa_terms(stored)
        if not _is_specific_qa_text(stored, stored_terms):
            continue

        if norm_q in stored or stored in norm_q:
            s = min(len(norm_q), len(stored)) / max(len(norm_q), len(stored))
            if s > best_sub_score:
                best_sub_score = s
                best_sub_answer = item.get("a")
                best_sub_item = item

    if best_sub_answer is not None and best_sub_score >= 0.88:
        print(f"[QA SUBSTRING] score={best_sub_score:.2f}")
        _record_source_match(
            match_trace,
            best_sub_item,
            source_type="manual_qa"
            if str((best_sub_item or {}).get("source") or "").lower() == "manual"
            else "textbook_auto_qa",
            method="substring",
            score=best_sub_score,
        )
        return _format_answer_for_chat(best_sub_answer)

    # 3) Fuzzy match.  Lao whitespace segmentation is not a tokenizer, so use
    # symmetric coverage, Jaccard, and whole-string similarity rather than
    # accepting generic query-only overlap.
    fuzzy_candidates: List[Tuple[float, Dict[str, Any], Dict[str, float]]] = []
    best_fuzzy_answer: Optional[str] = None
    best_fuzzy_item: Optional[Dict[str, Any]] = None

    for item in candidate_items:
        stored_norm = item.get("norm_q", "") or ""
        stored_terms = _qa_terms(stored_norm)
        if not _is_specific_qa_text(stored_norm, stored_terms):
            continue

        query_set = set(q_terms)
        stored_set = set(stored_terms)
        overlap = len(query_set & stored_set)
        query_coverage = overlap / max(len(query_set), 1)
        stored_coverage = overlap / max(len(stored_set), 1)
        union_size = len(query_set | stored_set)
        jaccard = overlap / union_size if union_size else 0.0
        sequence_ratio = SequenceMatcher(None, norm_q, stored_norm).ratio()
        score = max(
            sequence_ratio,
            (query_coverage + stored_coverage + jaccard) / 3.0,
        )
        fuzzy_candidates.append(
            (
                score,
                item,
                {
                    "overlap": float(overlap),
                    "query_coverage": query_coverage,
                    "stored_coverage": stored_coverage,
                    "jaccard": jaccard,
                    "sequence_ratio": sequence_ratio,
                },
            )
        )

    fuzzy_candidates.sort(key=lambda value: value[0], reverse=True)
    if fuzzy_candidates:
        best_fuzzy_score, best_fuzzy_item, fuzzy_stats = fuzzy_candidates[0]
        best_fuzzy_answer = best_fuzzy_item.get("a")
        runner_up_score = (
            fuzzy_candidates[1][0] if len(fuzzy_candidates) > 1 else 0.0
        )
        fuzzy_margin = best_fuzzy_score - runner_up_score
        fuzzy_is_strong = bool(
            fuzzy_stats["sequence_ratio"] >= 0.90
            or (
                fuzzy_stats["overlap"] >= 2
                and fuzzy_stats["query_coverage"] >= 0.75
                and fuzzy_stats["stored_coverage"] >= 0.65
                and fuzzy_stats["jaccard"] >= 0.55
            )
        )
        if (
            best_fuzzy_answer is not None
            and fuzzy_is_strong
            and fuzzy_margin >= 0.04
        ):
            print(
                f"[QA FUZZY] score={best_fuzzy_score:.2f}, "
                f"margin={fuzzy_margin:.2f}"
            )
            _record_source_match(
                match_trace,
                best_fuzzy_item,
                source_type="manual_qa"
                if str((best_fuzzy_item or {}).get("source") or "").lower()
                == "manual"
                else "textbook_auto_qa",
                method="fuzzy",
                score=best_fuzzy_score,
            )
            return _format_answer_for_chat(best_fuzzy_answer)

    # 4) Embedding similarity (semantic fallback)
    try:
        knowledge, qa_embeddings = _qa_embedding_snapshot()
        if qa_embeddings is None:
            return None

        q_emb = embed_model.encode(
            [norm_q],
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )[0]

        sims = np.dot(qa_embeddings, q_emb)

        # Ignore weak keyword-only prompts and hard-filter supplied scope.
        valid_indices: List[int] = []
        for i in _ordered_qa_indices(grade, unit, knowledge):
            item = knowledge[i]
            norm_i = item.get("norm_q", "") or ""
            terms_i = _qa_terms(norm_i)
            if _is_specific_qa_text(norm_i, terms_i):
                valid_indices.append(i)

        if not valid_indices:
            if match_trace is not None:
                match_trace["rejection_reason"] = "no_semantic_candidates_in_scope"
            return None

        valid_sims = sims[valid_indices]
        best_local_idx = int(np.argmax(valid_sims))
        best_idx = valid_indices[best_local_idx]
        best_sim = float(valid_sims[best_local_idx])

        scoped = bool(normalize_scope(grade) or normalize_scope(unit))
        sim_threshold = (
            QA_SCOPED_SEMANTIC_THRESHOLD
            if scoped
            else QA_GLOBAL_SEMANTIC_THRESHOLD
        )
        minimum_margin = (
            QA_SCOPED_SEMANTIC_MARGIN
            if scoped
            else QA_GLOBAL_SEMANTIC_MARGIN
        )
        ranked_valid = sorted(
            valid_indices, key=lambda idx: float(sims[idx]), reverse=True
        )
        second_sim = (
            float(sims[ranked_valid[1]]) if len(ranked_valid) > 1 else None
        )
        accepted, acceptance_reason, margin = semantic_acceptance(
            best_sim,
            second_sim,
            threshold=sim_threshold,
            minimum_margin=minimum_margin,
        )
        ranked_global = np.argsort(sims)[::-1].tolist()
        scoped_top = ranked_valid[:5]
        global_top = ranked_global[:5]
        if match_trace is not None:
            def candidate_row(idx: int, rank: int) -> Dict[str, Any]:
                item = knowledge[idx]
                source_type = (
                    "manual_qa"
                    if str(item.get("source") or "").lower() == "manual"
                    else "textbook_auto_qa"
                )
                return {
                    "rank": rank,
                    "source_id": str(item.get("id") or ""),
                    "source_type": source_type,
                    "grade": str(item.get("grade") or ""),
                    "unit": str(item.get("unit") or ""),
                    "similarity_score": float(sims[idx]),
                    "in_scope": item_in_scope(item, grade, unit),
                    "question": str(item.get("q") or ""),
                    "answer": str(item.get("a") or ""),
                    "source_file": str(item.get("_source_file") or ""),
                }
            match_trace.update(
                {
                    "semantic_threshold": sim_threshold,
                    "minimum_margin": minimum_margin,
                    "similarity_margin": margin,
                    "top_candidates": [candidate_row(idx, rank) for rank, idx in enumerate(global_top, start=1)],
                    "global_top_candidates": [candidate_row(idx, rank) for rank, idx in enumerate(global_top, start=1)],
                    "scoped_top_candidates": [candidate_row(idx, rank) for rank, idx in enumerate(scoped_top, start=1)],
                }
            )

        if accepted:
            ans = knowledge[best_idx].get("a")
            if ans:
                print(f"[QA EMBED] sim={best_sim:.2f}, margin={margin:.2f}")
                best_item = knowledge[best_idx]
                _record_source_match(
                    match_trace,
                    best_item,
                    source_type="manual_qa"
                    if str(best_item.get("source") or "").lower() == "manual"
                    else "textbook_auto_qa",
                    method="semantic",
                    score=best_sim,
                )
                return _format_answer_for_chat(ans)
        elif match_trace is not None:
            match_trace["rejection_reason"] = acceptance_reason
            near_miss: List[Dict[str, Any]] = []
            if scoped and best_sim >= QA_SCOPED_NEAR_MATCH_EVIDENCE_MIN:
                raw_near_miss: List[Dict[str, Any]] = []
                for candidate_rank, idx in enumerate(ranked_valid[:5], start=1):
                    item = knowledge[idx]
                    if str(item.get("source") or "").lower() != "manual":
                        continue
                    score = float(sims[idx])
                    if score < QA_SCOPED_NEAR_MATCH_EVIDENCE_MIN:
                        continue
                    raw_near_miss.append(
                        {
                            "rank": candidate_rank,
                            "source_id": str(item.get("id") or ""),
                            "source_type": "manual_qa",
                            "grade": str(item.get("grade") or ""),
                            "unit": str(item.get("unit") or ""),
                            "source_file": str(item.get("_source_file") or ""),
                            "similarity_score": score,
                            "retrieved_text": (
                                f"ຄຳຖາມ: {item.get('q') or ''}\n"
                                f"ຄຳຕອບ: {item.get('a') or ''}"
                            ),
                            "chapter_title": "",
                            "section_title": "",
                            "title": "",
                            "section": "",
                            "retrieval_scope": "manual_qa_near_miss",
                            "evidence_priority": "manual_qa_near_miss",
                            "direct_rejection_reason": acceptance_reason,
                        }
                    )
                near_miss = select_manual_near_miss_candidates(
                    raw_near_miss, max_candidates=3, score_window=0.025
                )
                for near_rank, item in enumerate(near_miss, start=1):
                    item["rank"] = near_rank
            match_trace["near_miss_evidence_candidates"] = near_miss

    except Exception as e:
        print(f"[QA EMBED WARN] {e}")
        if match_trace is not None:
            match_trace["rejection_reason"] = (
                f"semantic_match_error:{type(e).__name__}"
            )

    if match_trace is not None and not match_trace.get("rejection_reason"):
        if best_sub_answer is not None and best_sub_score < 0.88:
            match_trace["rejection_reason"] = "substring_below_threshold"
        elif fuzzy_candidates:
            match_trace["rejection_reason"] = "fuzzy_match_not_distinct_enough"
        else:
            match_trace["rejection_reason"] = "no_direct_match"
    return None


def _rank_manual_evidence(
    query: str,
    *,
    grade: Any = None,
    unit: Any = None,
    limit: int = 3,
) -> List[Dict[str, Any]]:
    """Rank scoped teacher Manual Q&A records as compound-answer evidence."""
    knowledge, embeddings = _qa_embedding_snapshot()
    if embeddings is None:
        return []
    indices = [
        idx
        for idx, item in enumerate(knowledge)
        if str(item.get("source") or "").lower() == "manual"
        and item_in_scope(item, grade, unit)
    ]
    if not indices:
        return []
    query_embedding = embed_model.encode(
        [normalize_lao_text(query)],
        convert_to_numpy=True,
        normalize_embeddings=True,
        show_progress_bar=False,
    )[0]
    similarities = np.dot(embeddings, query_embedding)
    ranked = sorted(
        indices, key=lambda idx: float(similarities[idx]), reverse=True
    )[: max(int(limit or 0), 0)]
    return [
        {
            "rank": rank,
            "source_id": str(knowledge[idx].get("id") or ""),
            "source_type": "manual_qa",
            "grade": str(knowledge[idx].get("grade") or ""),
            "unit": str(knowledge[idx].get("unit") or ""),
            "source_file": str(knowledge[idx].get("_source_file") or ""),
            "similarity_score": float(similarities[idx]),
            "retrieved_text": (
                f"ຄຳຖາມ: {knowledge[idx].get('q') or ''}\n"
                f"ຄຳຕອບ: {knowledge[idx].get('a') or ''}"
            ),
            "chapter_title": "",
            "section_title": "",
            "title": "",
            "section": "",
            "retrieval_scope": "compound_scoped",
        }
        for rank, idx in enumerate(ranked, start=1)
    ]


def manual_qa_ranking_diagnostics(
    query: str,
    *,
    grade: Any = None,
    unit: Any = None,
    expected_source_ids: Optional[Sequence[str]] = None,
) -> Dict[str, Any]:
    """Persist Manual Q&A global/scoped Top-5 and exact expected ranks."""
    knowledge, embeddings = _qa_embedding_snapshot()
    if embeddings is None:
        return {
            "global_top5": [],
            "scoped_top5": [],
            "expected_source_ranks": [],
            "reason": "qa_embeddings_unavailable",
        }
    manual_indices = [
        idx
        for idx, item in enumerate(knowledge)
        if str(item.get("source") or "").lower() == "manual"
    ]
    if not manual_indices:
        return {
            "global_top5": [],
            "scoped_top5": [],
            "expected_source_ranks": [],
            "reason": "no_manual_qa_entries",
        }
    query_embedding = embed_model.encode(
        [normalize_lao_text(query)],
        convert_to_numpy=True,
        normalize_embeddings=True,
        show_progress_bar=False,
    )[0]
    similarities = np.dot(embeddings, query_embedding)
    global_ranked = sorted(
        manual_indices,
        key=lambda idx: float(similarities[idx]),
        reverse=True,
    )
    scoped_ranked = [
        idx
        for idx in global_ranked
        if item_in_scope(knowledge[idx], grade, unit)
    ]

    def row(idx: int, rank: int) -> Dict[str, Any]:
        item = knowledge[idx]
        return {
            "rank": rank,
            "source_id": str(item.get("id") or ""),
            "grade": str(item.get("grade") or ""),
            "unit": str(item.get("unit") or ""),
            "similarity_score": float(similarities[idx]),
            "question": str(item.get("q") or ""),
            "answer": str(item.get("a") or ""),
        }

    expected = {
        str(source_id or "").strip()
        for source_id in (expected_source_ids or [])
        if str(source_id or "").strip()
    }
    global_rank_by_id = {
        str(knowledge[idx].get("id") or ""): rank
        for rank, idx in enumerate(global_ranked, start=1)
    }
    scoped_rank_by_id = {
        str(knowledge[idx].get("id") or ""): rank
        for rank, idx in enumerate(scoped_ranked, start=1)
    }
    score_by_id = {
        str(knowledge[idx].get("id") or ""): float(similarities[idx])
        for idx in manual_indices
    }
    return {
        "global_top5": [row(idx, rank) for rank, idx in enumerate(global_ranked[:5], start=1)],
        "scoped_top5": [row(idx, rank) for rank, idx in enumerate(scoped_ranked[:5], start=1)],
        "expected_source_ranks": [
            {
                "source_id": source_id,
                "global_rank": global_rank_by_id.get(source_id),
                "scoped_rank": scoped_rank_by_id.get(source_id),
                "similarity_score": score_by_id.get(source_id),
            }
            for source_id in sorted(expected)
        ],
        "manual_candidate_count": len(global_ranked),
        "scoped_manual_candidate_count": len(scoped_ranked),
        "reason": "ranked",
    }


def _rank_glossary_evidence(
    query: str,
    *,
    grade: Any = None,
    unit: Any = None,
    limit: int = 2,
) -> List[Dict[str, Any]]:
    """Rank scoped glossary definitions as evidence, not as a direct answer."""
    glossary = list(getattr(qa_store, "GLOSSARY", None) or [])
    embeddings = getattr(qa_store, "GLOSSARY_EMBEDDINGS", None)
    if not glossary or embeddings is None:
        return []
    indices = [
        idx
        for idx, item in enumerate(glossary)
        if item_in_scope(item, grade, unit)
    ]
    if not indices:
        return []
    query_embedding = embed_model.encode(
        [query],
        convert_to_numpy=True,
        normalize_embeddings=True,
    )[0]
    similarities = np.dot(embeddings, query_embedding)
    ranked = sorted(
        indices, key=lambda idx: float(similarities[idx]), reverse=True
    )[: max(int(limit or 0), 0)]
    return [
        {
            "rank": rank,
            "source_id": str(glossary[idx].get("id") or ""),
            "source_type": "glossary",
            "grade": str(glossary[idx].get("grade") or ""),
            "unit": str(glossary[idx].get("unit") or ""),
            "source_file": str(glossary[idx].get("_source_file") or ""),
            "similarity_score": float(similarities[idx]),
            "retrieved_text": (
                f"ຄຳສັບ: {glossary[idx].get('term') or ''}\n"
                f"ນິຍາມ: {glossary[idx].get('definition') or ''}"
            ),
            "chapter_title": "",
            "section_title": "",
            "title": "",
            "section": "",
            "retrieval_scope": "compound_scoped",
        }
        for rank, idx in enumerate(ranked, start=1)
    ]


def retrieve_compound_context_details(
    question: str,
    max_entries: int = MAX_CONTEXT_ENTRIES,
    *,
    grade: Any = None,
    unit: Any = None,
    decision_trace: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
    """Retrieve at least one scoped evidence item per detected component.

    Normal questions remain fixed at Top-4. Compound questions receive a
    dynamic evidence budget equal to the number of explicit components, with a
    small Top-4 floor for compatibility with the existing interface.
    """
    if decision_trace is not None:
        decision_trace.clear()
    components = extract_compound_components(question)
    if not components:
        components = [
            {"id": "component_1", "label": "component 1", "query": question}
        ]

    per_component: Dict[str, List[Dict[str, Any]]] = {}
    all_candidates: Dict[Tuple[str, str, str, str], Dict[str, Any]] = {}
    component_selections: List[Dict[str, Any]] = []

    for component in components:
        component_id = str(component.get("id") or "")
        component_label = str(component.get("label") or component_id)
        query = str(component.get("query") or question)
        candidates: List[Dict[str, Any]] = []
        candidates.extend(
            retrieve_context_details(
                query,
                max_entries=4,
                grade=grade,
                unit=unit,
                allow_global_fallback=False,
            )
        )
        candidates.extend(
            _rank_manual_evidence(query, grade=grade, unit=unit, limit=5)
        )
        candidates.extend(
            _rank_glossary_evidence(query, grade=grade, unit=unit, limit=3)
        )
        candidates = [
            dict(item)
            for item in candidates
            if item.get("similarity_score") is not None
        ]
        candidates.sort(
            key=lambda item: float(item.get("similarity_score") or 0.0),
            reverse=True,
        )
        per_component[component_id] = candidates
        for item in candidates:
            key = (
                str(item.get("source_type") or ""),
                str(item.get("grade") or ""),
                str(item.get("unit") or ""),
                str(item.get("source_id") or ""),
            )
            existing = all_candidates.get(key)
            if existing is None:
                existing = dict(item)
                existing["component_ids"] = []
                existing["component_labels"] = []
                existing["component_scores"] = {}
                existing["component_queries"] = {}
                all_candidates[key] = existing
            score = float(item.get("similarity_score") or 0.0)
            previous_score = float(
                existing["component_scores"].get(component_id) or 0.0
            )
            if score >= previous_score:
                existing["component_scores"][component_id] = score
                existing["component_queries"][component_id] = query
            if component_id not in existing["component_ids"]:
                existing["component_ids"].append(component_id)
                existing["component_labels"].append(component_label)
            existing["similarity_score"] = max(
                float(existing.get("similarity_score") or 0.0), score
            )

    selected_keys: List[Tuple[str, str, str, str]] = []
    component_to_key: Dict[str, Tuple[str, str, str, str]] = {}
    threshold = RAG_SCOPED_MIN_SIMILARITY

    for component in components:
        component_id = str(component.get("id") or "")
        candidates = per_component.get(component_id) or []
        chosen_key: Optional[Tuple[str, str, str, str]] = None
        chosen_item: Optional[Dict[str, Any]] = None
        for candidate in candidates:
            if float(candidate.get("similarity_score") or 0.0) < threshold:
                continue
            if not component_is_supported(component_id, candidate, question):
                continue
            key = (
                str(candidate.get("source_type") or ""),
                str(candidate.get("grade") or ""),
                str(candidate.get("unit") or ""),
                str(candidate.get("source_id") or ""),
            )
            chosen_key = key
            chosen_item = dict(candidate)
            chosen_item["selected_evidence_span"] = evidence_span_for_component(
                component_id, candidate, question
            )
            break
        if chosen_key is not None:
            component_to_key[component_id] = chosen_key
            if chosen_key not in selected_keys:
                selected_keys.append(chosen_key)
        component_selections.append(
            {
                "component_id": component_id,
                "component_label": str(component.get("label") or component_id),
                "query": str(component.get("query") or ""),
                "selected_source_id": str((chosen_item or {}).get("source_id") or ""),
                "selected_source_type": str((chosen_item or {}).get("source_type") or ""),
                "similarity_score": (
                    float((chosen_item or {}).get("similarity_score") or 0.0)
                    if chosen_item
                    else None
                ),
                "selected_evidence_span": str(
                    (chosen_item or {}).get("selected_evidence_span") or ""
                ),
                "covered": chosen_item is not None,
            }
        )

    dynamic_limit = max(MAX_CONTEXT_ENTRIES, len(components), int(max_entries or 0))
    if len(selected_keys) < dynamic_limit:
        for key, item in sorted(
            all_candidates.items(),
            key=lambda pair: float(pair[1].get("similarity_score") or 0.0),
            reverse=True,
        ):
            if key in selected_keys:
                continue
            if float(item.get("similarity_score") or 0.0) < threshold:
                continue
            selected_keys.append(key)
            if len(selected_keys) >= dynamic_limit:
                break

    expected_component_ids = [
        str(component.get("id") or "")
        for component in components
        if str(component.get("id") or "")
    ]
    component_label_by_id = {
        str(component.get("id") or ""): str(
            component.get("label") or component.get("id") or ""
        )
        for component in components
    }
    selected: List[Dict[str, Any]] = []
    component_evidence_map: Dict[str, Dict[str, Any]] = {}
    for key in selected_keys:
        item = dict(all_candidates[key])
        component_scores = dict(item.get("component_scores") or {})
        owned_components = [
            component_id
            for component_id in expected_component_ids
            if component_to_key.get(component_id) == key
        ]
        item["component_ids"] = owned_components
        item["component_labels"] = [
            component_label_by_id.get(component_id, component_id)
            for component_id in owned_components
        ]
        item["component_scores"] = {
            component_id: float(component_scores.get(component_id) or 0.0)
            for component_id in owned_components
        }
        raw_component_queries = dict(item.get("component_queries") or {})
        item["component_queries"] = {
            component_id: str(raw_component_queries.get(component_id) or "")
            for component_id in owned_components
        }
        evidence_spans = {
            component_id: evidence_span_for_component(
                component_id, item, question
            )
            for component_id in owned_components
        }
        evidence_spans = {
            component_id: span
            for component_id, span in evidence_spans.items()
            if span
        }
        item["selected_evidence_spans"] = evidence_spans
        if len(evidence_spans) == 1:
            item["selected_evidence_span"] = next(iter(evidence_spans.values()))
        if owned_components:
            item["similarity_score"] = max(
                float(component_scores.get(component_id) or 0.0)
                for component_id in owned_components
            )
        for component_id in owned_components:
            component_evidence_map[component_id] = {
                "selected_source_id": str(item.get("source_id") or ""),
                "selected_source_type": str(item.get("source_type") or ""),
                "similarity_score": float(
                    component_scores.get(component_id) or 0.0
                ),
                "selected_evidence_span": evidence_spans.get(component_id, ""),
            }
        selected.append(item)

    for rank, item in enumerate(selected, start=1):
        item["rank"] = rank
        item["retrieval_scope"] = "compound_component_scoped"
        item["expected_component_ids"] = expected_component_ids

    if decision_trace is not None:
        decision_trace.update(
            {
                "scope_mode": "compound_component_scoped",
                "requested_grades": list(normalize_scope(grade)),
                "requested_units": list(normalize_scope(unit)),
                "fallback_used": False,
                "compound_components": components,
                "component_selections": component_selections,
                "component_evidence_map": component_evidence_map,
                "component_required_count": len(components),
                "component_covered_count": sum(
                    1 for item in component_selections if item.get("covered")
                ),
                "dynamic_context_budget": dynamic_limit,
                "candidate_count": len(all_candidates),
                "selected_source_types": [
                    str(item.get("source_type") or "") for item in selected
                ],
                "reason": "component_coverage_retrieval"
                if selected
                else "no_component_evidence",
            }
        )
    return selected


def _plain_summary_text(text: str, max_chars: int = 260) -> str:
    clean = html.unescape(text or "")
    clean = re.sub(r"(?is)<br\s*/?>", " ", clean)
    clean = re.sub(r"(?is)</(td|th|p|div|li|tr)>", " ", clean)
    clean = re.sub(r"(?is)<[^>]+>", " ", clean)
    clean = _format_latexish_text(clean)
    clean = re.sub(r"\s+", " ", clean).strip()
    if len(clean) <= max_chars:
        return clean
    return clean[: max_chars - 3].rstrip() + "..."


def _qa_source_label(item: dict) -> str:
    grade = str(item.get("grade") or "").strip()
    unit = str(item.get("unit") or "").strip()
    source = str(item.get("source") or "").strip()
    parts = [p for p in [grade, unit, source] if p]
    return "/".join(parts)


def _is_manual_qa_item(item: dict) -> bool:
    return str(item.get("source") or "").strip().lower() == "manual"


_LATEX_SYMBOL_REPLACEMENTS = {
    r"\mu": "μ",
    r"\times": "×",
    r"\cdot": "·",
    r"\approx": "≈",
    r"\leq": "≤",
    r"\le": "≤",
    r"\geq": "≥",
    r"\ge": "≥",
    r"\neq": "≠",
    r"\pm": "±",
    r"\div": "÷",
    r"\alpha": "α",
    r"\beta": "β",
    r"\gamma": "γ",
    r"\Delta": "Δ",
    r"\delta": "δ",
    r"\theta": "θ",
    r"\lambda": "λ",
    r"\pi": "π",
    r"\rho": "ρ",
    r"\sigma": "σ",
    r"\Omega": "Ω",
    r"\omega": "ω",
}

_SUBSCRIPT_TRANSLATION = str.maketrans(
    {
        "0": "₀",
        "1": "₁",
        "2": "₂",
        "3": "₃",
        "4": "₄",
        "5": "₅",
        "6": "₆",
        "7": "₇",
        "8": "₈",
        "9": "₉",
        "+": "₊",
        "-": "₋",
        "=": "₌",
        "(": "₍",
        ")": "₎",
        "a": "ₐ",
        "e": "ₑ",
        "h": "ₕ",
        "i": "ᵢ",
        "j": "ⱼ",
        "k": "ₖ",
        "l": "ₗ",
        "m": "ₘ",
        "n": "ₙ",
        "o": "ₒ",
        "p": "ₚ",
        "r": "ᵣ",
        "s": "ₛ",
        "t": "ₜ",
        "u": "ᵤ",
        "v": "ᵥ",
        "x": "ₓ",
    }
)

_SUPERSCRIPT_TRANSLATION = str.maketrans(
    {
        "0": "⁰",
        "1": "¹",
        "2": "²",
        "3": "³",
        "4": "⁴",
        "5": "⁵",
        "6": "⁶",
        "7": "⁷",
        "8": "⁸",
        "9": "⁹",
        "+": "⁺",
        "-": "⁻",
        "=": "⁼",
        "(": "⁽",
        ")": "⁾",
        "n": "ⁿ",
    }
)


def _translate_script(text: str, translation: dict) -> str:
    return text.translate(translation)


def _strip_math_delimiters(text: str) -> str:
    clean = re.sub(r"(?s)\$\$(.*?)\$\$", lambda m: m.group(1), text)
    clean = re.sub(r"(?s)\$(.*?)\$", lambda m: m.group(1), clean)
    clean = re.sub(r"(?s)\\\[(.*?)\\\]", lambda m: m.group(1), clean)
    clean = re.sub(r"(?s)\\\((.*?)\\\)", lambda m: m.group(1), clean)
    return clean


def _format_latexish_text(text: str) -> str:
    clean = _strip_math_delimiters(text or "")

    for _ in range(4):
        updated = re.sub(
            r"\\frac\s*\{([^{}]+)\}\s*\{([^{}]+)\}",
            lambda m: f"{m.group(1).strip()}/{m.group(2).strip()}",
            clean,
        )
        if updated == clean:
            break
        clean = updated

    for source, replacement in _LATEX_SYMBOL_REPLACEMENTS.items():
        clean = clean.replace(source, replacement)

    clean = re.sub(
        r"_\{([^{}]+)\}",
        lambda m: _translate_script(m.group(1), _SUBSCRIPT_TRANSLATION),
        clean,
    )
    clean = re.sub(
        r"\^\{([^{}]+)\}",
        lambda m: _translate_script(m.group(1), _SUPERSCRIPT_TRANSLATION),
        clean,
    )
    clean = re.sub(
        r"_([A-Za-z0-9+\-=()])",
        lambda m: _translate_script(m.group(1).lower(), _SUBSCRIPT_TRANSLATION),
        clean,
    )
    clean = re.sub(
        r"\^([A-Za-z0-9+\-=()])",
        lambda m: _translate_script(m.group(1).lower(), _SUPERSCRIPT_TRANSLATION),
        clean,
    )

    clean = re.sub(r"\\([A-Za-z]+)", r"\1", clean)
    clean = clean.replace(r"\{", "{").replace(r"\}", "}")
    clean = clean.replace("{", "").replace("}", "")
    return clean


def _plain_full_text(text: str) -> str:
    clean = html.unescape(text or "")
    clean = re.sub(r"(?is)<br\s*/?>", "\n", clean)
    clean = re.sub(r"(?is)</(p|div|li|tr)>", "\n", clean)
    clean = re.sub(r"(?is)</(td|th)>", " ", clean)
    clean = re.sub(r"(?is)<[^>]+>", "", clean)
    clean = _format_latexish_text(clean)
    clean = clean.replace("\r\n", "\n").replace("\r", "\n")
    clean = re.sub(r"[ \t\f\v]+", " ", clean)
    clean = re.sub(r" *\n *", "\n", clean)
    clean = re.sub(r"\n{3,}", "\n\n", clean)
    if re.search(r"\b1\.\s+", clean) and re.search(r"\s2\.\s+", clean):
        clean = re.sub(r"(?<!\n)\s+(\d{1,2}\.\s+)", r"\n\1", clean)
    return clean.strip()


def _append_broad_qa_lines(lines: List[str], question: str, answer: str) -> None:
    if "\n" not in answer:
        lines.append(f"- {question}: {answer}")
        return

    lines.append(f"- {question}:")
    for answer_line in answer.splitlines():
        if answer_line.strip():
            lines.append(f"  {answer_line.strip()}")
        else:
            lines.append("")


def _rank_related_qa(question: str, limit: int = BROAD_QA_LIMIT) -> List[dict]:
    raw_norm = qa_store.normalize_question(question)
    norm_q = normalize_lao_text(raw_norm)
    knowledge = list(getattr(qa_store, "ALL_QA_KNOWLEDGE", None) or [])
    if not norm_q or not knowledge:
        return []

    q_terms = set(_qa_terms(norm_q))
    scored: List[tuple] = []

    try:
        knowledge, embeddings = _qa_embedding_snapshot()
        if embeddings is not None:
            q_emb = embed_model.encode(
                [norm_q],
                convert_to_numpy=True,
                normalize_embeddings=True,
                show_progress_bar=False,
            )[0]
            sims = np.dot(embeddings, q_emb)
        else:
            sims = None
    except Exception as e:  # noqa: BLE001
        print(f"[BROAD QA EMBED WARN] {e}")
        sims = None

    for idx, item in enumerate(knowledge):
        if not _is_manual_qa_item(item):
            continue

        answer = str(item.get("a") or "").strip()
        question_text = str(item.get("q") or "").strip()
        stored_norm = item.get("norm_q", "") or normalize_lao_text(qa_store.normalize_question(question_text))
        stored_terms = set(_qa_terms(stored_norm))
        if not answer or not question_text:
            continue

        overlap = len(q_terms & stored_terms)
        coverage = overlap / max(len(q_terms), 1)
        sim = float(sims[idx]) if sims is not None and idx < len(sims) else 0.0

        if sim < BROAD_MIN_EMBED_SIM and overlap == 0:
            continue

        score = sim + (0.10 * coverage) + (0.03 * min(overlap, 4))
        if str(item.get("grade") or "").strip() == "LEGACY":
            score -= 0.08
        scored.append((score, sim, overlap, item))

    scored.sort(key=lambda row: (row[0], row[1], row[2]), reverse=True)
    if not scored:
        return []

    primary = next(
        (item for _, _, _, item in scored if str(item.get("grade") or "").strip() != "LEGACY"),
        scored[0][3],
    )
    primary_grade = str(primary.get("grade") or "").strip()
    primary_unit = str(primary.get("unit") or "").strip()

    if primary_grade and primary_unit and primary_grade != "LEGACY":
        same_unit = [
            row for row in scored
            if str(row[3].get("grade") or "").strip() == primary_grade
            and str(row[3].get("unit") or "").strip() == primary_unit
        ]
        other_units = [
            row for row in scored
            if row not in same_unit and str(row[3].get("grade") or "").strip() != "LEGACY"
        ]
        legacy_units = [
            row for row in scored
            if row not in same_unit and str(row[3].get("grade") or "").strip() == "LEGACY"
        ]
        scored = [*same_unit, *other_units, *legacy_units]

    selected: List[dict] = []
    seen_answers = set()
    for _, _, _, item in scored:
        answer_key = qa_store.normalize_question(_plain_summary_text(str(item.get("a") or ""), 240))
        if not answer_key or answer_key in seen_answers:
            continue
        seen_answers.add(answer_key)
        selected.append(item)
        if len(selected) >= limit:
            break

    return selected


def _retrieve_context_blocks(question: str, limit: int = BROAD_CONTEXT_LIMIT) -> List[str]:
    context = retrieve_context(question, max_entries=limit)
    blocks = [block.strip() for block in context.split("\n\n") if block.strip()]
    return [_plain_summary_text(block, 300) for block in blocks[:limit]]


def build_broad_prompt(question: str, history: Optional[List] = None) -> str:
    context = retrieve_context(question, max_entries=max(MAX_CONTEXT_ENTRIES, 5))
    history_block = _format_history(history)
    return f"""{SYSTEM_PROMPT}

{history_block}ຂໍ້ມູນອ້າງອີງ:
{context}

ຄຳຖາມ: {question}

ຈົ່ງຕອບແບບກວ້າງແຕ່ກະຊັບ ໂດຍເລືອກຈຸດທີ່ກ່ຽວຂ້ອງທີ່ສຸດ 4-6 ບັນທັດ. ຢ່າຕອບຍາວເກີນໄປ.

ຄຳຕອບກວ້າງດ້ວຍພາສາລາວ:"""


def generate_broad_answer(question: str, history: Optional[List] = None) -> str:
    prompt = build_broad_prompt(question, history)
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=260,
            do_sample=False,
        )

    generated_ids = outputs[0][inputs["input_ids"].shape[1] :]
    answer = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    return answer or generate_answer(question, history)


def broad_research_bot(message: str, history: List) -> str:
    msg = (message or "").strip()
    if not msg:
        return "ກະລຸນາພິມຄຳຖາມກ່ອນ."

    if _is_yes_no_reply(msg):
        prev_user_q = _extract_last_user_text(history)
        if prev_user_q:
            msg = prev_user_q
        else:
            return "ກະລຸນາພິມຄຳຖາມໃຫ້ຊັດເຈນອີກຄັ້ງ."

    guided = _guided_reply(msg, history)
    if guided:
        return guided

    related_qa = _rank_related_qa(msg)

    if related_qa:
        lines = ["ຄຳຕອບແບບກວ້າງ:"]
        for item in related_qa:
            q = _plain_full_text(str(item.get("q") or ""))
            a = _plain_full_text(str(item.get("a") or ""))
            _append_broad_qa_lines(lines, q, a)

        return "\n".join(lines)

    return "ຄຳຕອບແບບກວ້າງ:\n- ບໍ່ພົບຄຳຕອບທີ່ກ່ຽວຂ້ອງໃນລາຍການ Q&A ທີ່ກຽມໄວ້."


# -----------------------------
# Main chatbot entry (Q&A > Glossary > evidence-gated RAG)
# -----------------------------
def _answer_with_rag(
    question: str,
    history: Optional[List],
    *,
    grade: Any = None,
    unit: Any = None,
    compound: bool = False,
    additional_contexts: Optional[Sequence[Mapping[str, Any]]] = None,
    answerable_hint: Optional[bool] = None,
    trace: Optional[Dict[str, Any]] = None,
) -> str:
    """Run one shared, trace-independent retrieval/generation path."""
    retrieval_trace: Dict[str, Any] = {}
    if compound:
        raw_contexts = retrieve_compound_context_details(
            question,
            max_entries=MAX_CONTEXT_ENTRIES,
            grade=grade,
            unit=unit,
            decision_trace=retrieval_trace,
        )
    else:
        raw_contexts = retrieve_context_details(
            question,
            max_entries=MAX_CONTEXT_ENTRIES,
            grade=grade,
            unit=unit,
            decision_trace=retrieval_trace,
        )

    if additional_contexts and not compound:
        merged_contexts: List[Dict[str, Any]] = []
        seen = set()
        for item in [*additional_contexts, *raw_contexts]:
            saved = dict(item)
            key = (
                str(saved.get("source_type") or ""),
                str(saved.get("source_id") or ""),
            )
            if key in seen:
                continue
            seen.add(key)
            merged_contexts.append(saved)
            if len(merged_contexts) >= MAX_CONTEXT_ENTRIES:
                break
        raw_contexts = merged_contexts
        retrieval_trace["manual_qa_near_miss_evidence"] = [
            dict(item) for item in additional_contexts
        ]

    scoped = bool(normalize_scope(grade) or normalize_scope(unit))
    generation_contexts, confidence = assess_retrieval_confidence(
        question,
        raw_contexts,
        scoped=scoped,
        compound=compound,
    )
    retrieval_trace["confidence"] = confidence

    if trace is not None:
        trace.update(
            {
                "retrieved_contexts": raw_contexts,
                "generation_contexts": generation_contexts,
                "retrieval_decision": retrieval_trace,
                "contexts_used_by_generation": False,
                "seallms_invoked": False,
                "seallms_invocation_count": 0,
            }
        )
    if not generation_contexts:
        if trace is not None:
            trace.update(
                {
                    "route": "insufficient_evidence_refusal",
                    "safe_refusal_reason": confidence.get("reason"),
                }
            )
        return SAFE_LAO_REFUSAL

    generation_validation: Dict[str, Any] = {}
    if trace is not None:
        trace.update(
            {
                "route": (
                    "textbook_rag_compound"
                    if compound
                    else "textbook_rag_fallback"
                ),
                "contexts_used_by_generation": True,
                "seallms_invoked": True,
            }
        )
    try:
        answer = generate_answer(
            question,
            history,
            retrieved_contexts=generation_contexts,
            validation_trace=generation_validation,
            allow_extractive_fallback=(answerable_hint is not False),
            evidence_confidence=confidence,
            compound=compound,
        )
    except Exception as exc:  # noqa: BLE001
        if trace is not None:
            trace.update(
                {
                    "route": "textbook_rag_error",
                    "error_message": f"{type(exc).__name__}: {exc}",
                    "generation_validation": generation_validation,
                    "seallms_invocation_count": max(
                        1, len(generation_validation.get("attempts") or [])
                    ),
                }
            )
        return SAFE_LAO_REFUSAL

    if trace is not None:
        trace.update(
            {
                "generation_validation": generation_validation,
                "seallms_invocation_count": len(
                    generation_validation.get("attempts") or []
                ),
            }
        )
        if generation_validation.get("returned_safe_refusal"):
            trace["route"] = "textbook_rag_output_refusal"
            trace["safe_refusal_reason"] = generation_validation.get(
                "fallback_reason"
            ) or "generation_validation_failed"
        elif generation_validation.get("used_extractive_fallback"):
            trace["route"] = "textbook_rag_extractive_fallback"
            trace["fallback_reason"] = generation_validation.get(
                "fallback_reason"
            ) or "sufficient_evidence_generation_failed"
    return answer


def laos_science_bot(
    message: str,
    history: List,
    grade: Any = None,
    unit: Any = None,
    *,
    evaluation_trace: Optional[Dict[str, Any]] = None,
    answerable_hint: Optional[bool] = None,
) -> Any:
    """
    Main chatbot function for Student tab (Gradio ChatInterface).

    Direct routes are used only for a safe, scoped single-source match.
    Compound questions bypass them and use combined evidence.  All remaining
    questions pass through the same retrieval-confidence and output-safety
    gates whether or not evaluation tracing is enabled.
    """
    local_trace: Dict[str, Any] = (
        evaluation_trace if evaluation_trace is not None else {}
    )
    local_trace.clear()
    local_trace.update(
        {
            "route": "not_started",
            "retrieved_contexts": [],
            "generation_contexts": [],
            "contexts_used_by_generation": False,
            "seallms_invoked": False,
            "seallms_invocation_count": 0,
        }
    )

    msg = (message or "").strip()
    if not msg:
        local_trace["route"] = "empty_input"
        return "ກະລຸນາພິມຄໍາຖາມກ່ອນ."

    guided = _guided_reply(msg, history)
    if guided:
        local_trace["route"] = "guided_navigation"
        return guided

    if _is_yes_no_reply(msg):
        prev_user_q = _extract_last_user_text(history)
        if prev_user_q:
            msg = prev_user_q
            local_trace["linked_yes_no_to_previous_question"] = True
        else:
            local_trace["route"] = "unlinked_yes_no"
            return "ກະລຸນາພິມຄໍາຖາມໃຫ້ຊັດເຈນອີກຄັ້ງ."

    compound = is_compound_question(msg)
    local_trace["question_shape"] = {
        "compound": compound,
        "direct_routes_bypassed": compound,
    }

    manual_near_miss_contexts: List[Dict[str, Any]] = []
    if compound:
        local_trace["qa_match_diagnostics"] = {
            "accepted": False,
            "rejection_reason": "compound_direct_route_bypass",
        }
        local_trace["glossary_match_diagnostics"] = {
            "accepted": False,
            "rejection_reason": "compound_direct_route_bypass",
        }
    else:
        qa_match_trace: Dict[str, Any] = {}
        try:
            direct = answer_from_qa(
                msg,
                grade=grade,
                unit=unit,
                match_trace=qa_match_trace,
            )
            local_trace["qa_match_diagnostics"] = qa_match_trace
            manual_near_miss_contexts = list(
                qa_match_trace.get("near_miss_evidence_candidates") or []
            )
            if direct:
                local_trace.update(
                    {
                        "route": "qa_direct",
                        "direct_source_match": qa_match_trace,
                    }
                )
                return direct
        except Exception as exc:  # noqa: BLE001
            local_trace["qa_route_error"] = f"{type(exc).__name__}: {exc}"
            local_trace["qa_match_diagnostics"] = qa_match_trace

        glossary_match_trace: Dict[str, Any] = {}
        try:
            gloss = answer_from_glossary(
                msg,
                grade=grade,
                unit=unit,
                match_trace=glossary_match_trace,
            )
            local_trace["glossary_match_diagnostics"] = glossary_match_trace
            if gloss:
                local_trace.update(
                    {
                        "route": "glossary_direct",
                        "direct_source_match": glossary_match_trace,
                    }
                )
                return gloss
        except Exception as exc:  # noqa: BLE001
            local_trace["glossary_route_error"] = (
                f"{type(exc).__name__}: {exc}"
            )
            local_trace["glossary_match_diagnostics"] = glossary_match_trace

    return _answer_with_rag(
        msg,
        history,
        grade=grade,
        unit=unit,
        compound=compound,
        additional_contexts=manual_near_miss_contexts,
        answerable_hint=answerable_hint,
        trace=local_trace,
    )