# -*- coding: utf-8 -*-
"""
kmdb_rag.py — 用 Books KMDB 已清理的 Markdown 建一個最小 RAG，並比較切分策略
CM979B AI 協作研究｜吳相勳｜2026-10

這支程式刻意不依賴 LangChain／LlamaIndex，每一步都攤開來寫，讓同學看見 RAG 的骨架：
  1. 讀檔：cleaned_*.md 以「### PDF 第 N 頁」為頁碼錨點
  2. 前處理：接回 PDF 斷行、統一繁簡字
  3. 切分（chunking）：七種策略
  4. 建索引：BM25（關鍵字）、向量（語意）、混合（RRF）
  5. 檢索與評估：Hit@5、MRR、送進 LLM 的脈絡長度
  6. 生成：把檢索結果組成附頁碼的提示詞，交給 Claude 或任何 LLM 作答
"""
import re, glob, os, math, time, json
from dataclasses import dataclass, field

# ------------------------------------------------------------------
# 1. 讀檔：每本書切成「頁」，保留頁碼，之後每個片段都能回到原書第幾頁
# ------------------------------------------------------------------
PAGE_HEAD = re.compile(r'^### PDF 第 (\d+) 頁\s*$', re.M)

def load_book(path):
    """KMDB 格式：以「### PDF 第 N 頁」切頁。一般 Markdown：沒有頁碼錨點時，改以標題切成「段落區塊」，
    區塊編號放在 page 欄位、標題文字放在 label 欄位，之後一樣能回到原文位置。"""
    raw = open(path, encoding='utf-8-sig').read()
    title = (raw.lstrip().splitlines() or [os.path.basename(path)])[0].lstrip('# ').strip()
    body = raw.split('## 全書正文', 1)[-1]
    pages = []
    if PAGE_HEAD.search(body):
        parts = PAGE_HEAD.split(body)
        for i in range(1, len(parts), 2):
            txt = re.sub(r'<a id="page-\d+"></a>', '', parts[i + 1]).strip()
            if txt.startswith('[本頁未辨識到足夠文字'):
                continue  # 空白頁、純圖頁：不進索引
            pages.append({'page': int(parts[i]), 'label': f'PDF 第 {parts[i]} 頁', 'text': txt})
    else:
        parts = re.split(r'^(#{1,3} .+)$', raw, flags=re.M)
        if parts[0].strip():
            pages.append({'page': 0, 'label': '開頭', 'text': parts[0].strip()})
        for i in range(1, len(parts), 2):
            if parts[i + 1].strip():
                pages.append({'page': len(pages) + 1, 'label': parts[i].lstrip('# ').strip(), 'text': parts[i + 1].strip()})
    slug = os.path.basename(os.path.dirname(path)) if os.path.basename(path).startswith('cleaned_') \
        else os.path.splitext(os.path.basename(path))[0]
    return {'slug': slug, 'title': title, 'pages': pages, 'source': path}

def load_corpus(root, slugs=None):
    """KMDB：讀 root/*/cleaned_*.md；其他資料夾：讀 root 底下所有 .md"""
    files = sorted(glob.glob(os.path.join(root, '*', 'cleaned_*.md'))) or \
        sorted(glob.glob(os.path.join(root, '**', '*.md'), recursive=True))
    books = []
    for f in files:
        b = load_book(f)
        if slugs and b['slug'] not in slugs:
            continue
        if b['pages']:
            books.append(b)
    return books

# ------------------------------------------------------------------
# 2. 前處理
#    (a) PDF 轉出的文字每行硬換行，句子被攔腰截斷 → 接回段落
#    (b) 簡體原文 vs 繁體提問 → 關鍵字檢索會對不上，統一成繁體
# ------------------------------------------------------------------
CJK = r'㐀-鿿豈-﫿'
CJK_PUNCT = '，。、；：？！「」『』（）《》〈〉—…“”'
END_PUNCT = '。！？!?；;：:」』”）)】'

def rejoin_lines(text):
    text = re.sub(r'^#{1,6} ', '', text, flags=re.M)        # 頁內誤判的標題符號
    text = re.sub(r'-\n(?=[a-z])', '', text)                  # 英文斷字 hyphen
    lines = text.split('\n')
    out = ''
    for ln in lines:
        ln = ln.strip()
        if not ln:
            out += '\n\n'; continue
        if not out or out.endswith('\n'):
            out += ln; continue
        prev = out[-1]
        if prev in END_PUNCT and len(ln) > 0 and (ln[0].isupper() or re.match(f'[{CJK}「“(（]', ln)):
            out += '\n' + ln                                  # 前一行已是句末：保留換行
        elif re.match(f'[{CJK}{CJK_PUNCT}]', prev) and re.match(f'[{CJK}{CJK_PUNCT}]', ln):
            out += ln                                         # 中文：直接接起來
        else:
            out += ' ' + ln                                   # 英文：補空白
    return re.sub(r'\n{3,}', '\n\n', out).strip()

try:
    import opencc
    _s2t = opencc.OpenCC('s2tw')   # s2tw：轉成臺灣用字（吃、為），s2t 會轉成「喫」「爲」
    def to_trad(s): return _s2t.convert(s)
except Exception:                                             # 沒裝 opencc 就跳過
    def to_trad(s): return s

def prepare(books, rejoin=True, normalize_script=True):
    """回傳每本書的連續全文 + 每頁在全文中的起訖位置（用來把片段對回頁碼）"""
    docs = []
    for b in books:
        buf, spans, pos, labels = [], [], 0, {}
        for p in b['pages']:
            t = p['text']
            if rejoin: t = rejoin_lines(t)
            if normalize_script: t = to_trad(t)
            buf.append(t); spans.append((pos, pos + len(t), p['page'])); labels[p['page']] = p.get('label', str(p['page']))
            pos += len(t) + 2                                 # 頁與頁之間以空行相接
        docs.append({'slug': b['slug'], 'title': to_trad(b['title']) if normalize_script else b['title'],
                     'text': '\n\n'.join(buf), 'spans': spans, 'labels': labels})
    return docs

def pages_of(doc, start, end):
    return [pg for (s, e, pg) in doc['spans'] if s < end and e > start]

# ------------------------------------------------------------------
# 3. 切分策略
# ------------------------------------------------------------------
@dataclass
class Chunk:
    id: str
    slug: str
    title: str
    pages: list
    text: str                 # 拿去建索引的文字
    context: str = ''         # 檢索到之後真正送給 LLM 的文字（small-to-big 時會比 text 長）
    meta: dict = field(default_factory=dict)

def _mk(doc, i, s, e, text=None, prefix=''):
    t = doc['text'][s:e] if text is None else text
    pg = pages_of(doc, s, e)
    return Chunk(f"{doc['slug']}#{i}", doc['slug'], doc['title'], pg, prefix + t.strip(), t.strip(),
                 {'labels': [doc['labels'].get(p, str(p)) for p in pg], 'start': s, 'end': e})

# 3-1 固定長度：每 N 個字元切一刀，不管句子
def chunk_fixed(doc, size=500, overlap=0):
    step = size - overlap; out = []
    for i, s in enumerate(range(0, len(doc['text']), step)):
        e = min(s + size, len(doc['text']))
        if doc['text'][s:e].strip():
            out.append(_mk(doc, i, s, e))
        if e == len(doc['text']): break
    return out

# 3-2 遞迴切分：先試段落，太長再試句子，再不行才硬切（LangChain RecursiveCharacterTextSplitter 的原理）
SEPS = ['\n\n', '\n', '。', '！', '？', '. ', '? ', '! ', '；', '; ', '，', ', ', ' ', '']

def _split(text, size, seps):
    if len(text) <= size: return [text]
    sep = next((s for s in seps if s and s in text), '')
    if sep == '':
        return [text[i:i + size] for i in range(0, len(text), size)]
    pieces = [p + sep for p in text.split(sep)]
    pieces[-1] = pieces[-1][:-len(sep)]
    rest = seps[seps.index(sep) + 1:]
    out = []
    for p in pieces:
        out.extend(_split(p, size, rest) if len(p) > size else [p])
    return out

def chunk_recursive(doc, size=500):
    pieces = _split(doc['text'], size, SEPS)
    out, cur, cur_s, pos = [], '', 0, 0
    for p in pieces:                                          # 把小碎片合併到接近 size
        if cur and len(cur) + len(p) > size:
            out.append(_mk(doc, len(out), cur_s, cur_s + len(cur)))
            cur, cur_s = '', pos
        if not cur: cur_s = pos
        cur += p; pos += len(p)
    if cur.strip(): out.append(_mk(doc, len(out), cur_s, cur_s + len(cur)))
    return out

# 3-3 依文件結構：這批書的 Markdown 標題就是「PDF 第 N 頁」，所以結構切分＝一頁一塊
def chunk_page(doc):
    return [_mk(doc, i, s, e) for i, (s, e, pg) in enumerate(doc['spans'])]

# 3-4 語意切分：逐句做向量，相鄰兩句相似度掉到谷底就切開（Kamradt 式 semantic chunking）
SENT = re.compile(r'(?<=[。！？!?])|(?<=\.)\s+(?=[A-Z])')

def chunk_semantic(doc, embed_fn, pct=20, min_len=150, max_len=900):
    import numpy as np
    sents, offs, pos = [], [], 0
    for s in SENT.split(doc['text']):
        if s and s.strip():
            i = doc['text'].find(s, pos); sents.append(s); offs.append(i); pos = i + len(s)
    if len(sents) < 3: return chunk_recursive(doc)
    E = embed_fn(sents)
    sims = (E[:-1] * E[1:]).sum(1)
    cut = np.percentile(sims, pct)
    out, start = [], offs[0]
    for k in range(1, len(sents)):
        cur_len = offs[k] - start
        if (sims[k - 1] < cut and cur_len >= min_len) or cur_len >= max_len:
            out.append(_mk(doc, len(out), start, offs[k])); start = offs[k]
    out.append(_mk(doc, len(out), start, len(doc['text'])))
    return out

# 3-5 小塊檢索、大塊作答（small-to-big／parent-child）：用 200 字小塊找，回傳整頁
def chunk_small_to_big(doc, child=200):
    page_text = {pg: doc['text'][s:e] for (s, e, pg) in doc['spans']}
    out = chunk_recursive(doc, child)
    for c in out:
        c.context = '\n'.join(page_text[p] for p in c.pages)
        c.meta['parent'] = c.pages
    return out

# 3-6 加上脈絡標頭：每塊前面補「書名＋頁碼」（Contextual Retrieval 的低成本版；正式版請 LLM 寫 50–100 tokens）
def chunk_with_header(doc, size=500):
    out = chunk_recursive(doc, size)
    for c in out:
        head = f"《{doc['title']}》{c.meta['labels'][0].replace('PDF ', '')}：" if c.pages else f"《{doc['title']}》："
        c.text = head + c.context
    return out

# ------------------------------------------------------------------
# 4. 索引與檢索
# ------------------------------------------------------------------
def tokenize(s):
    """中文用字元二元組（bigram），英文用單字；不用斷詞套件也能跑"""
    s = s.lower(); toks = re.findall(r'[a-z0-9]+', s)
    for run in re.findall(f'[{CJK}]+', s):
        toks += [run[i:i + 2] for i in range(len(run) - 1)] or [run]
    return toks

class BM25Index:
    def __init__(self, chunks):
        from rank_bm25 import BM25Okapi
        self.chunks = chunks
        self.bm = BM25Okapi([tokenize(c.text) for c in chunks])
    def search(self, q, k=10):
        import numpy as np
        sc = self.bm.get_scores(tokenize(q)); idx = np.argsort(-sc)[:k]
        return [(self.chunks[i], float(sc[i])) for i in idx]

class Embedder:
    """multilingual-e5-small：118M 參數、中英文都能用；e5 規定要加 query:/passage: 前綴"""
    def __init__(self, name='intfloat/multilingual-e5-small'):
        from sentence_transformers import SentenceTransformer
        self.m = SentenceTransformer(name, device=os.environ.get('DEVICE', 'cpu'))
        self.tok = self.m.tokenizer
        self.max_tokens = self.m.max_seq_length
    def passages(self, texts, bs=32):
        return self.m.encode(['passage: ' + t for t in texts], batch_size=bs, normalize_embeddings=True, show_progress_bar=False)
    def queries(self, texts):
        return self.m.encode(['query: ' + t for t in texts], normalize_embeddings=True, show_progress_bar=False)
    def sym(self, texts, bs=64):
        return self.m.encode(['query: ' + t for t in texts], batch_size=bs, normalize_embeddings=True, show_progress_bar=False)
    def n_tokens(self, text):
        return len(self.tok('passage: ' + text, add_special_tokens=True)['input_ids'])

class DenseIndex:
    def __init__(self, chunks, emb, vectors=None):
        self.chunks, self.emb = chunks, emb
        self.V = vectors if vectors is not None else emb.passages([c.text for c in chunks])
    def search(self, q, k=10):
        import numpy as np
        sc = self.V @ self.emb.queries([q])[0]; idx = np.argsort(-sc)[:k]
        return [(self.chunks[i], float(sc[i])) for i in idx]

def rrf(*result_lists, k=60, top=10):
    """Reciprocal Rank Fusion：把關鍵字排名與語意排名合併"""
    score, keep = {}, {}
    for res in result_lists:
        for r, (c, _) in enumerate(res):
            score[c.id] = score.get(c.id, 0) + 1 / (k + r + 1); keep[c.id] = c
    ids = sorted(score, key=lambda i: -score[i])[:top]
    return [(keep[i], score[i]) for i in ids]

def dedupe_parents(results, k=5):
    """small-to-big：多個小塊落在同一頁，只回傳一次整頁"""
    seen, out = set(), []
    for c, s in results:
        key = (c.slug, tuple(c.meta.get('parent', c.pages)))
        if key in seen: continue
        seen.add(key); out.append((c, s))
        if len(out) == k: break
    return out

# ------------------------------------------------------------------
# 5. 評估：答案頁是否出現在前 5 名
# ------------------------------------------------------------------
def is_hit(c, gold):
    return any(c.slug == g['book'] and set(c.pages) & set(g['pages']) for g in gold)

def evaluate(search_fn, qs, k=5):
    hits, rr, ctx = 0, 0.0, 0
    per_q = []
    for q in qs:
        res = search_fn(q['q'])
        rank = next((i + 1 for i, (c, _) in enumerate(res[:10]) if is_hit(c, q['gold'])), None)
        hits += 1 if rank and rank <= k else 0
        rr += 1 / rank if rank else 0
        ctx += sum(len(c.context or c.text) for c, _ in res[:k])
        per_q.append({'id': q['id'], 'rank': rank, 'top1': (res[0][0].slug, res[0][0].pages[:3]) if res else None})
    n = len(qs)
    return {'hit@5': hits / n, 'mrr@10': rr / n, 'ctx_chars@5': ctx / n, 'per_q': per_q}

# ------------------------------------------------------------------
# 6. 生成：組出附來源頁碼的提示詞
# ------------------------------------------------------------------
def build_prompt(question, results, k=5):
    blocks = []
    for i, (c, _) in enumerate(results[:k], 1):
        lab = c.meta.get('labels') or [str(p) for p in c.pages]
        loc = lab[0] if len(lab) == 1 else f"{lab[0]}～{lab[-1]}"
        blocks.append(f"[{i}] 《{c.title}》{loc}\n{(c.context or c.text)[:1500]}")
    return (
        "你是研究助理。只能根據下列書籍摘錄回答；每個論點後面用 [編號] 標出依據。"
        "摘錄不足以回答時，直接說明缺什麼，不要自行補充。請用臺灣繁體中文作答。\n\n"
        + "\n\n".join(blocks) + f"\n\n問題：{question}")

def ask_claude(prompt, model=os.environ.get('CLAUDE_MODEL', 'claude-sonnet-5-5')):
    """有 ANTHROPIC_API_KEY 才會呼叫；沒有的話把 prompt 貼到 Claude 對話框即可"""
    import anthropic
    msg = anthropic.Anthropic().messages.create(model=model, max_tokens=1200,
                                               messages=[{'role': 'user', 'content': prompt}])
    return msg.content[0].text