# -*- coding: utf-8 -*-
"""
Lõi của demo NamBoQA: truy hồi + đọc hiểu, bám sát thực nghiệm trong bài.

  Truy hồi   BM25 (pyvi) · Bi-encoder · Hybrid  (alpha * norm(BM25) + (1-alpha) * norm(dense))
  Đọc hiểu   Extractive (XLM-R SQuAD2, ngưỡng từ chối) · Generative (Gemini, có grounding)

Code tokenize, chuẩn hoá điểm, prompt Gemini, nhận diện câu từ chối được lấy
nguyên từ nambo_run/code/run_experiments.py và run_reading.py (bộ thực nghiệm
25/10/2026, dữ liệu vòng 3) để demo cho ra đúng hành vi đã báo cáo trong bài.
Cách chia train/dev/test dùng chung chia_tap.py: theo địa điểm, phân tầng theo tỉnh.
"""

from __future__ import annotations

import json
import logging
import os
import re
import threading
import time
import unicodedata
from pathlib import Path

import numpy as np

log = logging.getLogger("namboqa")

NO_ANSWER = "Tài liệu không có thông tin để trả lời."
# chỉ dùng để chọn từ tô sáng trên giao diện, không ảnh hưởng truy hồi
STOPWORDS = {"người", "sau", "khi", "nơi", "được", "của", "những", "các", "này", "đã", "và",
             "là", "có", "cho", "với", "trong", "vào", "nào", "gì", "đến", "năm", "bao", "nhiêu",
             "đâu", "không", "từ", "theo", "tại", "một", "như", "thì", "mà", "nên", "hay", "hoặc",
             "đó", "ở", "để", "về", "ra", "lại", "còn", "rằng", "thời", "gian", "bởi", "do"}
DATA_FILES = ("passages.jsonl", "places_clean.json", "qa_final.jsonl")


# ------------------------------------------------------------------ tiện ích
def load_jsonl(p: Path):
    with open(p, encoding="utf-8") as f:
        return [json.loads(l) for l in f if l.strip()]


def vi_tokens(s: str):
    from pyvi import ViTokenizer
    s = unicodedata.normalize("NFC", str(s).lower())
    s = re.sub(r"[^\w\s]", " ", s)
    return ViTokenizer.tokenize(s).split()


def norm(x):
    x = np.asarray(x, dtype=float)
    return (x - x.min()) / (x.max() - x.min() + 1e-9)


REFUSAL_RE = re.compile(
    r"(không|chưa)\s+(\w+\s+){0,2}"
    r"(có|nêu|đề\s*cập|nhắc|nói|cung\s*cấp|chứa|tìm\s*thấy|xác\s*định|ghi|thể)"
    r"|no[_\s]?answer", re.I)


def is_refusal(pred: str) -> bool:
    p = (pred or "").strip().lower()
    return (not p) or bool(REFUSAL_RE.search(p))


PROMPT_GEN = """Bạn trả lời câu hỏi CHỈ dựa vào tài liệu dưới đây, không dùng kiến thức bên ngoài.

Tài liệu:
<<<
{context}
>>>

Câu hỏi: {question}

Nếu tài liệu không chứa thông tin để trả lời, hãy trả lời đúng một câu: "{no_answer}"
Nếu có, trả lời NGẮN GỌN (tên riêng, năm, con số, cụm từ), không giải thích thêm.

Đáp án:"""


def cpu_threads() -> int:
    """Số luồng CPU thật sự được cấp. os.cpu_count() trong container HF trả về số nhân
    của cả máy chủ (vài chục) trong khi Space CPU basic chỉ có 2 vCPU: đặt theo đó làm
    torch tranh chấp luồng, mỗi câu hỏi chậm 5–10 giây. Đọc hạn mức cgroup thay vào."""
    if os.environ.get("TORCH_NUM_THREADS"):
        return max(1, int(os.environ["TORCH_NUM_THREADS"]))
    try:
        quota, period = open("/sys/fs/cgroup/cpu.max").read().split()[:2]
        if quota != "max":
            return max(1, int(int(quota) / int(period)))
    except Exception:
        pass
    try:
        return max(1, len(os.sched_getaffinity(0)))
    except Exception:
        return max(1, os.cpu_count() or 1)


# ------------------------------------------------------------------ dữ liệu
def ensure_data(data_dir: Path):
    """Dữ liệu KHÔNG nằm trong repo Space (passages.jsonl không được phát hành
    nguyên văn). Nếu thiếu thì tải từ một HF Dataset private bằng HF_TOKEN."""
    missing = [f for f in DATA_FILES if not (data_dir / f).exists()]
    if not missing:
        return
    repo = os.environ.get("DATA_REPO")
    if not repo:
        raise RuntimeError(f"Thiếu {missing} trong {data_dir} và chưa đặt DATA_REPO")
    from huggingface_hub import snapshot_download
    log.info(f"Tải dữ liệu từ dataset private {repo} ...")
    snapshot_download(repo_id=repo, repo_type="dataset", local_dir=str(data_dir),
                      token=os.environ.get("HF_TOKEN") or None)


class Engine:
    def __init__(self, data_dir: str | Path):
        self.data_dir = Path(data_dir)
        self.retriever_name = os.environ.get(
            "RETRIEVER_MODEL", "bkai-foundation-models/vietnamese-bi-encoder")
        self.reader_name = os.environ.get("READER_MODEL", "deepset/xlm-roberta-base-squad2")
        self.gemini_model = os.environ.get("GEMINI_MODEL", "gemini-3.5-flash")
        self.alpha = float(os.environ.get("HYBRID_ALPHA", "0.5"))       # chọn trên dev trong bài
        self.reject_th = float(os.environ.get("REJECT_THRESHOLD", "0.0"))  # chọn trên dev trong bài
        self.max_ctx = int(os.environ.get("MAX_CONTEXT_CHARS", "3000"))
        # Số đoạn truy hồi ghép làm ngữ cảnh cho từng bộ đọc. Bài (Bảng 11): bộ trích
        # xuất đạt F1 cao nhất với 1–3 đoạn, ghép 5 đoạn thì F1 và tỉ lệ từ chối đúng
        # đều giảm; mô hình sinh chịu nhiễu tốt nên giữ 5 như thực nghiệm chính.
        self.ctx_k = {"extractive": int(os.environ.get("EXTRACTIVE_CONTEXT_K", "3")),
                      "generative": int(os.environ.get("GENERATIVE_CONTEXT_K", "5"))}
        self.status = {"data": "pending", "bm25": "pending", "dense": "pending",
                       "extractive": "pending", "generative": "pending"}
        self.errors: dict[str, str] = {}
        self._lock = threading.Lock()

    # ------------------------------------------------------------- khởi động
    def load_fast(self):
        """Phần nhẹ, chạy đồng bộ: dữ liệu + BM25 (~5 giây)."""
        ensure_data(self.data_dir)
        d = self.data_dir
        self.passages = load_jsonl(d / "passages.jsonl")
        self.places = {p["place_id"]: p for p in
                       json.loads((d / "places_clean.json").read_text(encoding="utf-8"))}
        self.qa = load_jsonl(d / "qa_final.jsonl")
        self.place_split = self._split()
        self.ids = [p["id"] for p in self.passages]
        self.txts = [p["text"] for p in self.passages]
        self.pid2idx = {pid: i for i, pid in enumerate(self.ids)}
        self.q_index = {self._qkey(q["question"]): q for q in self.qa}
        self.status["data"] = "ready"

        from rank_bm25 import BM25Okapi
        t0 = time.time()
        self.bm25 = BM25Okapi([vi_tokens(t) for t in self.txts])
        self.status["bm25"] = "ready"
        log.info(f"BM25 xong: {len(self.txts)} passage, {time.time() - t0:.1f}s")

        has_key = bool(os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY"))
        self.status["generative"] = "ready" if has_key else "disabled"

    def _split(self):
        """Tái lập đúng lần chia trong bài: chia_tap.chia_theo_tinh, seed 42 (1.416 / 321 / 708 câu)."""
        import chia_tap
        pids = {p["id"] for p in self.passages}
        qa = [{**q, "gold": [g for g in q["gold_passage_ids"] if g in pids]} for q in self.qa]
        qa = [q for q in qa if q["gold"]]
        prov = {pid: p.get("province") or "?" for pid, p in self.places.items()}
        _, _, _, tk = chia_tap.chia_theo_tinh(qa, prov)
        out = {}
        for name in ("train", "dev", "test"):
            for pid in tk[f"{name}_p"]:
                out[pid] = name
        return out

    def _sources(self, p):
        """Nguồn tư liệu để người xem tự kiểm: link cổng du lịch (theo đoạn) hoặc bài Wikipedia (theo địa điểm)."""
        m = p.get("meta") or {}
        out = []
        for u in (m.get("portal_urls") or ([m["source_url"]] if m.get("source_url") else []))[:2]:
            out.append({"label": u.split("/")[2] if "//" in u else u, "url": u})
        for t in (m.get("wiki_titles") or [])[:2]:
            out.append({"label": f"Wikipedia: {t}",
                        "url": "https://vi.wikipedia.org/wiki/" + t.replace(" ", "_")})
        return out

    def load_heavy(self):
        """Phần nặng, chạy nền: bi-encoder + mô hình extractive."""
        self._load_dense()
        if os.environ.get("DISABLE_EXTRACTIVE", "0") == "1":
            self.status["extractive"] = "disabled"
        else:
            self._load_extractive()

    def _load_dense(self):
        try:
            self.status["dense"] = "loading"
            import torch
            from sentence_transformers import SentenceTransformer
            torch.set_num_threads(cpu_threads())
            self.encoder = SentenceTransformer(self.retriever_name, device="cpu")
            safe = re.sub(r"[^\w.-]", "_", self.retriever_name)
            cache = self.data_dir / f"emb_{safe}.npy"
            if cache.exists():
                emb = np.load(cache)
                if emb.shape[0] != len(self.txts):
                    emb = None
            else:
                emb = None
            if emb is None:
                t0 = time.time()
                log.info("Mã hoá corpus bằng bi-encoder (lần đầu, vài phút trên CPU)...")
                emb = self.encoder.encode(self.txts, batch_size=32, normalize_embeddings=True,
                                          convert_to_numpy=True, show_progress_bar=False)
                try:
                    np.save(cache, emb.astype(np.float32))
                except Exception:
                    pass
                log.info(f"Mã hoá xong trong {time.time() - t0:.0f}s")
            self.emb = emb.astype(np.float32)
            self.status["dense"] = "ready"
        except Exception as e:
            log.exception("không nạp được bi-encoder")
            self.status["dense"] = "error"
            self.errors["dense"] = str(e)[:300]

    def _load_extractive(self):
        try:
            self.status["extractive"] = "loading"
            import torch
            from transformers import AutoModelForQuestionAnswering, AutoTokenizer
            self.qa_tok = AutoTokenizer.from_pretrained(self.reader_name)
            self.qa_model = AutoModelForQuestionAnswering.from_pretrained(self.reader_name).eval()
            self._torch = torch
            self.status["extractive"] = "ready"
            log.info(f"Mô hình extractive: {self.reader_name}")
        except Exception as e:
            log.exception("không nạp được mô hình extractive")
            self.status["extractive"] = "error"
            self.errors["extractive"] = str(e)[:300]

    # ------------------------------------------------------------- truy hồi
    @staticmethod
    def _qkey(q: str) -> str:
        return re.sub(r"\s+", " ", unicodedata.normalize("NFC", q or "").strip().lower())

    def retrieve(self, question: str, method: str = "hybrid", k: int = 5, alpha: float | None = None):
        """Trả về (danh sách (idx, điểm), phương án thực sự dùng)."""
        alpha = self.alpha if alpha is None else alpha
        if method in ("dense", "hybrid") and self.status["dense"] != "ready":
            method = "bm25"                      # bi-encoder chưa nạp xong -> lùi về BM25
        bm = self.bm25.get_scores(vi_tokens(question)) if method != "dense" else None
        if method != "bm25":
            qv = self.encoder.encode([question], normalize_embeddings=True, convert_to_numpy=True)[0]
            dn = self.emb @ qv.astype(np.float32)
        if method == "bm25":
            s = bm
        elif method == "dense":
            s = dn
        else:
            s = alpha * norm(bm) + (1 - alpha) * norm(dn)
        top = np.argsort(-s)[:k]
        return [(int(i), float(s[i])) for i in top], method

    # ------------------------------------------------------------- đọc hiểu
    def extract(self, question: str, context: str, max_len=384, stride=128, max_answer_len=40):
        """Rút span kiểu SQuAD v2 (so với xác suất token CLS). Giống make_extractor
        trong run_reading.py nhưng vector hoá phép tìm cặp (start, end) cho nhanh trên CPU."""
        torch = self._torch
        tok = self.qa_tok
        enc = tok(question, context, truncation="only_second", max_length=max_len, stride=stride,
                  return_overflowing_tokens=True, return_offsets_mapping=True,
                  padding="max_length", return_tensors="pt")
        offsets = enc.pop("offset_mapping")
        enc.pop("overflow_to_sample_mapping", None)
        dev = next(self.qa_model.parameters()).device      # CPU trên Space, GPU khi đánh giá
        with torch.no_grad():
            out = self.qa_model(**{k: v.to(dev) for k, v in enc.items()})
        best = {"answer": "", "score": 0.0, "start": -1, "end": -1}
        for w in range(out.start_logits.size(0)):
            ps = torch.softmax(out.start_logits[w], -1).float().cpu().numpy()
            pe = torch.softmax(out.end_logits[w], -1).float().cpu().numpy()
            null = float(ps[0] * pe[0])
            sid = enc.sequence_ids(w)
            valid = np.array([i for i in range(len(sid)) if sid[i] == 1])
            if not len(valid):
                continue
            lo, hi = valid[0], valid[-1]
            m = np.outer(ps[lo:hi + 1], pe[lo:hi + 1])
            m = np.triu(m) - np.triu(m, max_answer_len)     # j >= i và j < i + max_answer_len
            i, j = np.unravel_index(np.argmax(m), m.shape)
            sc = float(m[i, j])
            if sc > best["score"] and sc > null:
                a, b = offsets[w][lo + i][0].item(), offsets[w][lo + j][1].item()
                best = {"answer": context[a:b], "score": sc, "start": a, "end": b}
        return best

    def generate(self, question: str, context: str) -> str:
        from google import genai
        from google.genai import types
        client = genai.Client()
        r = client.models.generate_content(
            model=self.gemini_model,
            contents=PROMPT_GEN.format(context=context, question=question, no_answer=NO_ANSWER),
            config=types.GenerateContentConfig(temperature=0.0, max_output_tokens=2048),
        )
        txt = ""
        try:
            for part in r.candidates[0].content.parts:
                if getattr(part, "thought", False):
                    continue
                txt += getattr(part, "text", "") or ""
        except Exception:
            txt = r.text or ""
        txt = (txt or "").strip()
        txt = re.sub(r"^(đáp án|answer)\s*[:：]\s*", "", txt, flags=re.I)
        return txt.split("\n")[0].strip().strip('"').rstrip(".")

    # ------------------------------------------------------------- pipeline
    def ask(self, question: str, method="hybrid", reader="generative", k=5, alpha=None):
        t0 = time.time()
        hits, used = self.retrieve(question, method, k, alpha)
        t_ret = time.time() - t0

        # ghép ngữ cảnh đúng như thực nghiệm: nối top-k rồi cắt max_ctx ký tự
        n_ctx = min(len(hits), self.ctx_k.get(reader, len(hits)))
        bounds, parts, pos = [], [], 0
        for idx, _ in hits[:n_ctx]:
            t = self.txts[idx]
            bounds.append((pos, pos + len(t), idx))
            parts.append(t)
            pos += len(t) + 1
        context = " ".join(parts)[: self.max_ctx]

        answer, score, src, span = None, None, None, None
        reader_used = reader
        t1 = time.time()
        if reader == "extractive":
            if self.status["extractive"] != "ready":
                raise RuntimeError(f"Mô hình extractive chưa sẵn sàng ({self.status['extractive']})")
            r = self.extract(question, context)
            score = r["score"]
            if r["answer"] and score >= self.reject_th:
                answer = r["answer"].strip()
                for a, b, idx in bounds:
                    if a <= r["start"] < b:
                        src, span = self.ids[idx], [r["start"] - a, r["end"] - a]
                        break
            else:
                answer = NO_ANSWER
        elif reader == "generative":
            if self.status["generative"] != "ready":
                raise RuntimeError("Chưa cấu hình GEMINI_API_KEY trên back-end")
            try:
                answer = self.generate(question, context)
            except Exception as e:
                raise RuntimeError(f"Gemini không trả lời được (kiểm tra GEMINI_API_KEY, quota, "
                                   f"tên mô hình GEMINI_MODEL): {str(e)[:160]}")
            if not is_refusal(answer):
                for _, _, idx in bounds:          # tìm đoạn chứa nguyên văn đáp án để tô sáng
                    p = self.txts[idx].lower().find(answer.lower())
                    if p >= 0 and len(answer) >= 2:
                        src, span = self.ids[idx], [p, p + len(answer)]
                        break
        else:
            reader_used = "none"
        t_read = time.time() - t1

        qwords = sorted({w.replace("_", " ") for w in vi_tokens(question)
                         if len(w) > 2 and w not in STOPWORDS},
                        key=len, reverse=True)
        max_chars = int(os.environ.get("MAX_PASSAGE_CHARS", "0"))
        passages = []
        for rank, (idx, sc) in enumerate(hits, 1):
            p = self.passages[idx]
            pl = self.places.get(p["place_id"], {})
            text = p["text"] if not max_chars else p["text"][:max_chars]
            passages.append({
                "rank": rank, "id": p["id"], "place_id": p["place_id"],
                "title": p.get("title") or pl.get("name"), "province": pl.get("province"),
                "heritage_level": pl.get("heritage_level"), "score": round(sc, 4),
                "text": text, "sources": self._sources(p),
                "in_context": reader != "none" and rank <= n_ctx and bounds[rank - 1][0] < self.max_ctx,
            })

        gold = None
        q = self.q_index.get(self._qkey(question))
        if q:
            gold = {"id": q["id"], "type": q["type"], "answer": q["answer"],
                    "gold_passage_ids": q["gold_passage_ids"],
                    "split": self.place_split.get(q["place_id"], "?"),
                    "round": "vòng 3" if q.get("vong_duyet") == 3 else "vòng 1–2",
                    "evaluated": q["type"] != "compare"}

        return {
            "question": question, "answer": answer,
            "refused": bool(answer is not None and is_refusal(answer)),
            "reader": reader_used, "reader_model": (self.reader_name if reader_used == "extractive"
                                                    else self.gemini_model if reader_used == "generative"
                                                    else None),
            "score": score, "source_passage_id": src, "answer_span": span,
            "retriever": used, "retriever_requested": method,
            "alpha": (self.alpha if alpha is None else alpha) if used == "hybrid" else None,
            "context_k": n_ctx if reader_used != "none" else 0,
            "passages": passages, "query_terms": qwords, "gold": gold,
            "timing_ms": {"retrieval": round(t_ret * 1000), "reading": round(t_read * 1000)},
        }

    # ------------------------------------------------------------- tiện ích UI
    def samples(self, n=6, qtype=None, split="test", seed=None):
        import random
        # mặc định bỏ "compare": bài không chấm loại này (13 câu test, quá ít)
        pool = [q for q in self.qa
                if (q["type"] == qtype if qtype else q["type"] != "compare")
                and (split == "all" or self.place_split.get(q["place_id"]) == split)]
        rnd = random.Random(seed)
        pick = rnd.sample(pool, min(n, len(pool)))
        return [{"id": q["id"], "type": q["type"], "question": q["question"],
                 "round": "vòng 3" if q.get("vong_duyet") == 3 else "vòng 1–2",
                 "place": self.places.get(q["place_id"], {}).get("name"),
                 "province": self.places.get(q["place_id"], {}).get("province")} for q in pick]

    def stats(self):
        from collections import Counter
        return {
            "n_passages": len(self.passages), "n_places": len(self.places), "n_qa": len(self.qa),
            "qa_by_type": dict(Counter(q["type"] for q in self.qa)),
            "qa_by_split": dict(Counter(self.place_split.get(q["place_id"], "?") for q in self.qa)),
            "places_by_province": dict(Counter(p.get("province") for p in self.places.values())),
            "models": {"retriever": self.retriever_name, "reader": self.reader_name,
                       "generative": self.gemini_model, "alpha": self.alpha,
                       "context_k": self.ctx_k},
        }