# -*- coding: utf-8 -*-
"""知识库: 把如东诊断 (报告/裁决/交接/现场单/六层文档/skill/findings) 切块入库; 检索 = BM25 (字符 bigram + ASCII 词), 零依赖."""
import json, math, re
from collections import Counter
from pathlib import Path


def _tokens(s):
    s = str(s).lower()
    asc = re.findall(r'[a-z0-9_#.\-]{2,}', s)
    cjk = re.findall(r'[一-鿿]', s)
    return asc + [a + b for a, b in zip(cjk, cjk[1:])]


def _chunks_md(text, source, maxlen=700):
    out, cur, head = [], [], ''
    for line in text.split('\n'):
        if line.startswith('#'):
            if cur:
                out.append((head, '\n'.join(cur)))
                cur = []
            head = line.strip('# ').strip()
        cur.append(line)
        if sum(len(x) for x in cur) > maxlen:
            out.append((head, '\n'.join(cur)))
            cur = []
    if cur:
        out.append((head, '\n'.join(cur)))
    return [dict(source=f'{source} › {h}' if h else source, text=t.strip()) for h, t in out if t.strip()]


PUBLIC_DOC_KEYS = ('收敛报告', '交接_振动', '现场单')   # 公开模式只收结论类文档; 方法/判据/裁决/skill/findings 原文不入库


def build(cfg):
    from src.windcms.redact import PUBLIC, redact
    docs = []
    for p in cfg['knowledge_docs']:
        p = Path(p)
        if not p.exists():
            continue
        if PUBLIC and not any(k in p.name for k in PUBLIC_DOC_KEYS):
            continue
        docs += _chunks_md(p.read_text(encoding='utf-8'), p.name)
    fj = cfg['findings']
    if fj.exists() and not PUBLIC:
        f = json.load(open(fj, encoding='utf-8'))
        for k, v in f.items():
            if k == 'meta':
                continue
            txt = json.dumps(v, ensure_ascii=False)
            for i in range(0, min(len(txt), 6000), 900):
                docs.append(dict(source=f'findings.json › {k}', text=txt[i:i + 900]))
    for d in docs:
        d['text'] = redact(d['text'])
        d['tokens'] = _tokens(d['text'] + ' ' + d['source'])
    out = cfg['out'] / 'kb.json'
    out.parent.mkdir(parents=True, exist_ok=True)
    # encoding 必须给: ensure_ascii=False 写的是中文, 缺 encoding 会按系统 locale(cp936) 编 → 中文 Windows 崩
    with open(out, 'w', encoding='utf-8') as f:
        json.dump(docs, f, ensure_ascii=False)
    print(f'kb: {len(docs)} chunks → {out}')
    _KB.pop(str(cfg['out']), None)
    return docs


_KB = {}


def load(cfg):
    key = str(cfg['out'])
    if key not in _KB:
        p = cfg['out'] / 'kb.json'
        _KB[key] = json.load(open(p, encoding='utf-8')) if p.exists() else build(cfg)
    return _KB[key]


def search(cfg, query, k=5):
    docs = load(cfg)
    N = len(docs)
    q = _tokens(query)
    df = Counter()
    L = []
    for d in docs:
        L.append(len(d['tokens']))
        for t in set(d['tokens']):
            df[t] += 1
    avg = sum(L) / max(N, 1)
    out = []
    for d, l in zip(docs, L):
        tf = Counter(d['tokens'])
        sc = 0.0
        for t in q:
            if t in tf:
                idf = math.log(1 + (N - df[t] + .5) / (df[t] + .5))
                sc += idf * tf[t] * 2.2 / (tf[t] + 1.2 * (0.25 + 0.75 * l / avg))
        if sc > 0:
            out.append(dict(score=sc, source=d['source'], text=d['text']))
    return sorted(out, key=lambda x: -x['score'])[:k]