# requires: jieba, rank_bm25, numpy
"""最小 RAG 流水线：切分 → 嵌入 → 混合检索（向量 + BM25，RRF 融合）→ 权限过滤 → 拒答阈值 → 带引用的提示词 → 检索评测。

为了不依赖外部服务，这里的“嵌入”用字符 n-gram 哈希向量代替真实的嵌入模型（只能捕捉字面相似，不懂语义），
目的是把流程跑通、把评测方法讲清楚。换成真实嵌入模型时，只需要替换 embed() 函数。
运行：python mini_rag.py
"""
from __future__ import annotations

import hashlib
import logging
import re
from dataclasses import dataclass, field

import jieba
import numpy as np
from rank_bm25 import BM25Okapi

jieba.setLogLevel(logging.WARNING)

DIM = 512


@dataclass
class Doc:
    id: str
    title: str
    text: str
    acl: set[str]  # 允许访问的部门；"all" 表示全员


@dataclass
class Chunk:
    id: str
    doc_id: str
    title: str
    text: str
    acl: set[str]
    vec: np.ndarray = field(repr=False, default=None)


DOCS = [
    Doc("hr-leave", "员工手册-休假制度", """# 年假
入职满一年的正式员工享有带薪年假 5 天，满十年 10 天，满二十年 15 天。年假需提前 3 个工作日在 OA 中申请。
# 病假
病假需提供医院证明，每年累计 10 天以内按全额工资发放。
# 调休
加班可在 3 个月内申请调休，过期作废。""", {"all"}),
    Doc("it-vpn", "IT 指南-远程办公", """# VPN
在家办公需先安装公司 VPN 客户端，使用工号和动态口令登录。连接失败请检查系统时间是否准确。
# 设备
笔记本电脑丢失需在 24 小时内报告 IT 部门并修改所有密码。""", {"all"}),
    Doc("fin-reimburse", "财务制度-报销", """# 差旅报销
出差住宿标准：一线城市每晚 500 元以内，其他城市 350 元以内。发票须在出差结束后 30 天内提交。
# 餐费
出差餐补每天 100 元，无需发票。""", {"all"}),
    Doc("hr-salary", "薪酬制度（仅 HR）", """# 调薪
年度调薪在每年 3 月进行，幅度由绩效等级决定，A 级 10% 至 15%。
# 保密
薪酬信息属于保密信息，不得在公开渠道讨论。""", {"hr"}),
]


def split(doc: Doc, max_len: int = 80, overlap: int = 15) -> list[Chunk]:
    """按 Markdown 标题分节；节内超过 max_len 再按句子切，相邻块保留 overlap 个字符的重叠。
    每块前面加上“文档标题 / 小节标题”，让孤立的块也有上下文（例如“满十年 10 天”属于“年假”）。"""
    out: list[Chunk] = []
    for sec in re.split(r"\n(?=# )", doc.text.strip()):
        heading, _, body = sec.partition("\n")
        heading = heading.lstrip("# ").strip()
        sentences = [s for s in re.split(r"(?<=[。！？])", body.strip()) if s]
        buf = ""
        pieces: list[str] = []
        for s in sentences:
            if buf and len(buf) + len(s) > max_len:
                pieces.append(buf)
                buf = buf[-overlap:]
            buf += s
        if buf:
            pieces.append(buf)
        for p in pieces:
            out.append(Chunk(f"{doc.id}#{len(out)}", doc.id, doc.title, f"【{doc.title} / {heading}】{p}", doc.acl))
    return out


def embed(text: str) -> np.ndarray:
    """字符 1~2 gram 哈希到 DIM 维并 L2 归一化（真实项目替换为嵌入模型 API）。"""
    v = np.zeros(DIM, dtype=np.float32)
    t = re.sub(r"\s+", "", text)
    for n in (1, 2):
        for i in range(len(t) - n + 1):
            h = int.from_bytes(hashlib.md5(t[i:i + n].encode()).digest()[:4], "little")
            v[h % DIM] += 1.0
    norm = np.linalg.norm(v)
    return v / norm if norm else v


def tokenize(text: str) -> list[str]:
    return [w for w in jieba.lcut(text) if w.strip() and not re.fullmatch(r"[，。、：；！？【】/\s]+", w)]


class Index:
    def __init__(self, docs: list[Doc]) -> None:
        self.chunks = [c for d in docs for c in split(d)]
        for c in self.chunks:
            c.vec = embed(c.text)
        self.matrix = np.stack([c.vec for c in self.chunks])
        self.bm25 = BM25Okapi([tokenize(c.text) for c in self.chunks])

    def allowed(self, user_depts: set[str]) -> np.ndarray:
        return np.array([("all" in c.acl) or bool(c.acl & user_depts) for c in self.chunks])

    def vector_search(self, q: str, mask: np.ndarray, k: int) -> list[tuple[int, float]]:
        scores = self.matrix @ embed(q)  # 已归一化，点积 = 余弦相似度
        scores[~mask] = -np.inf          # 权限过滤：检索阶段就排除，而不是生成后再删
        order = np.argsort(-scores)[:k]
        return [(int(i), float(scores[i])) for i in order if np.isfinite(scores[i])]

    def bm25_search(self, q: str, mask: np.ndarray, k: int) -> list[tuple[int, float]]:
        scores = np.array(self.bm25.get_scores(tokenize(q)))
        scores[~mask] = -np.inf
        order = np.argsort(-scores)[:k]
        return [(int(i), float(scores[i])) for i in order if np.isfinite(scores[i]) and scores[i] > 0]

    def hybrid_search(self, q: str, mask: np.ndarray, k: int, rrf_k: int = 60) -> list[tuple[int, float]]:
        """RRF：score = Σ 1 / (rrf_k + 排名)。只用排名，不用担心两种分数量纲不同。"""
        fused: dict[int, float] = {}
        for results in (self.vector_search(q, mask, 20), self.bm25_search(q, mask, 20)):
            for rank, (i, _) in enumerate(results, start=1):
                fused[i] = fused.get(i, 0.0) + 1.0 / (rrf_k + rank)
        return sorted(fused.items(), key=lambda kv: -kv[1])[:k]


def answer_context(index: Index, q: str, user_depts: set[str], k: int = 3, min_sim: float = 0.25) -> dict:
    mask = index.allowed(user_depts)
    hits = index.hybrid_search(q, mask, k)
    # 拒答：最相关块的向量相似度太低，说明知识库里很可能没有答案
    best_sim = max((float(index.matrix[i] @ embed(q)) for i, _ in hits), default=0.0)
    if not hits or best_sim < min_sim:
        return {"refuse": True, "prompt": None, "sources": []}
    sources = [index.chunks[i] for i, _ in hits]
    context = "\n".join(f"[{n}] {c.text}" for n, c in enumerate(sources, start=1))
    prompt = (
        "你是公司制度问答助手。只根据下面的资料回答，资料里没有的内容回答“资料中没有相关规定”。"
        "回答时用 [编号] 标注依据。\n\n"
        f"资料：\n{context}\n\n问题：{q}"
    )
    return {"refuse": False, "prompt": prompt, "sources": [c.id for c in sources], "best_sim": round(best_sim, 3)}


# ---------------- 检索评测 ----------------
EVAL = [  # (问题, 期望命中的文档)
    ("年假有几天", "hr-leave"),
    ("工作满十年有多少天带薪假", "hr-leave"),
    ("病假要什么材料", "hr-leave"),
    ("加班多久内可以调休", "hr-leave"),
    ("VPN 连不上怎么办", "it-vpn"),
    ("电脑丢了怎么处理", "it-vpn"),
    ("出差住酒店能报多少钱", "fin-reimburse"),
    ("发票什么时候交", "fin-reimburse"),
    ("出差吃饭有补贴吗", "fin-reimburse"),
]


def evaluate(index: Index, method: str, k: int = 3) -> tuple[float, float]:
    mask = index.allowed({"all"})
    hit, rr = 0, 0.0
    for q, expected in EVAL:
        results = getattr(index, f"{method}_search")(q, mask, k)
        docs = [index.chunks[i].doc_id for i, _ in results]
        if expected in docs:
            hit += 1
            rr += 1 / (docs.index(expected) + 1)
    return hit / len(EVAL), rr / len(EVAL)


def main() -> None:
    index = Index(DOCS)
    print(f"共 {len(DOCS)} 篇文档，切成 {len(index.chunks)} 块")

    print("\n检索评测（Recall@3 / MRR）：")
    scores = {}
    for m in ("vector", "bm25", "hybrid"):
        recall, mrr = evaluate(index, m)
        scores[m] = recall
        print(f"  {m:7s} recall@3={recall:.2f}  MRR={mrr:.2f}")

    r = answer_context(index, "满十年的年假几天？", {"dev"})
    print("\n问：满十年的年假几天？ 引用：", r["sources"])
    print(r["prompt"][:120], "……")

    r_hr = answer_context(index, "A 级调薪幅度是多少", {"hr"})
    r_dev = answer_context(index, "A 级调薪幅度是多少", {"dev"})
    print("\nHR 提问调薪：", r_hr["sources"][:1], "| 研发提问调薪：", r_dev["sources"])
    assert r_hr["sources"][0].startswith("hr-salary")
    assert not any(s.startswith("hr-salary") for s in r_dev["sources"])  # 无权限的内容不会进入上下文

    r_none = answer_context(index, "公司股票期权怎么行权", {"dev"})
    print("问：公司股票期权怎么行权 → 拒答" if r_none["refuse"] else f"未拒答：{r_none}")
    assert r_none["refuse"]
    assert scores["hybrid"] >= max(scores["vector"], scores["bm25"]) - 1e-9


if __name__ == "__main__":
    main()