# -*- coding: utf-8 -*-
"""
RAG 知识库：文档分块 + jieba 分词 + BM25 检索（本地零依赖）
升级路径：接入商用 embedding + pgvector/Qdrant 做向量检索（见 README）
"""
import math
import re
import jieba
import db

CHUNK_SIZE = 400          # 每块目标字数
CHUNK_OVERLAP = 80        # 块间重叠


def segment(text):
    """jieba 分词，去停用词/单字/空白"""
    text = re.sub(r'[\s，。！？、；：""''（）【】《》·…—0-9a-zA-Z-]+', ' ', text)
    words = [w for w in jieba.cut(text) if len(w.strip()) > 1 and not w.isspace()]
    return words

def split_chunks(content, size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
    """按段落聚合切块，避免从句子中间切断"""
    paras = [p.strip() for p in re.split(r'\n+', content) if p.strip()]
    chunks, buf, buf_len = [], '', 0
    for p in paras:
        if buf_len + len(p) > size and buf:
            chunks.append(buf)
            tail = buf[-overlap:] if overlap else ''
            buf, buf_len = tail + p, len(tail) + len(p)
        else:
            buf += ('\n' if buf else '') + p
            buf_len += len(p)
    if buf:
        chunks.append(buf)
    return chunks or [content[:size]]


def rebuild_document(doc_id):
    """重新分块索引文档"""
    doc = db.q('SELECT * FROM documents WHERE id=?', (doc_id,), one=True)
    if not doc:
        return 0
    db.w('DELETE FROM doc_chunks WHERE document_id=?', (doc_id,))
    chunks = split_chunks(doc['content'])
    for i, c in enumerate(chunks):
        db.w('INSERT INTO doc_chunks (document_id, idx, content, tokens) VALUES (?,?,?,?)',
             (doc_id, i, c, len(segment(c))))
    db.w('UPDATE documents SET chunk_size=?, updated_at=? WHERE id=?', (len(chunks), db.now(), doc_id))
    return len(chunks)


class BM25Index:
    """轻量 BM25：按查询词 IDF 加权打分"""

    def __init__(self, chunks):
        # chunks: [{id, document_id, content, tokens}]
        self.chunks = chunks
        self.doc_len = [max(c['tokens'], 1) for c in chunks]
        self.avg_len = sum(self.doc_len) / max(len(self.doc_len), 1)
        self.N = len(chunks)
        self.k1, self.b = 1.5, 0.75
        # 倒排：term -> set of chunk idx
        self.inv = {}
        self.df = {}
        for i, c in enumerate(chunks):
            seen = set()
            for w in segment(c['content']):
                if w in seen:
                    continue
                seen.add(w)
                self.inv.setdefault(w, []).append(i)
        for w, lst in self.inv.items():
            self.df[w] = len(lst)

    def search(self, query, top_k=5):
        q_terms = [w for w in segment(query)]
        if not q_terms or not self.N:
            return []
        scores = {}
        for w in q_terms:
            postings = self.inv.get(w, [])
            if not postings:
                continue
            idf = math.log(1 + (self.N - self.df[w] + 0.5) / (self.df[w] + 0.5))
            for idx in postings:
                tf = sum(1 for x in segment(self.chunks[idx]['content']) if x == w)
                denom = tf + self.k1 * (1 - self.b + self.b * self.doc_len[idx] / self.avg_len)
                scores[idx] = scores.get(idx, 0) + idf * (tf * (self.k1 + 1)) / denom
        ranked = sorted(scores.items(), key=lambda x: -x[1])[:top_k]
        return [{'chunk_id': self.chunks[i]['id'], 'document_id': self.chunks[i]['document_id'],
                 'content': self.chunks[i]['content'], 'score': round(s, 4)}
                for i, s in ranked]


def search_project(project_id, query, top_k=5):
    """在项目知识库中检索，返回 (片段列表, 是否命中)"""
    chunks = db.q(
        'SELECT c.id, c.document_id, c.content, c.tokens FROM doc_chunks c '
        'JOIN documents d ON d.id=c.document_id WHERE d.project_id=?',
        (project_id,))
    if not chunks:
        return [], False
    idx = BM25Index(chunks)
    return idx.search(query, top_k), True


def build_context(project_id, query, top_k=4):
    """生成注入提示词的检索上下文（含来源标注）"""
    hits, hit = search_project(project_id, query, top_k)
    if not hit or not hits:
        return '', []
    parts, refs = [], []
    for h in hits:
        doc = db.q('SELECT name FROM documents WHERE id=?', (h['document_id'],), one=True)
        name = doc['name'] if doc else f'文档#{h["document_id"]}'
        parts.append(f'【来源：{name}】\n{h["content"]}')
        refs.append(f'{name}#块{h["chunk_id"]}')
    ctx = '以下是项目知识库中的相关资料，回答时请优先参考：\n\n' + '\n\n---\n\n'.join(parts)
    return ctx, refs