"""知识库管理模块 - RAG 检索增强生成

架构：文档 → 精细分块(代码块独立/标题层级切分) → 中文embedding(BGE) + numpy余弦相似度 → 关键词+向量混合检索 → LLM生成回答
"""

import hashlib
import json
import os
import re
from pathlib import Path
from urllib.parse import quote

# ============ 配置 ============

KNOWLEDGE_DIR = Path(__file__).parent
INDEX_DIR = KNOWLEDGE_DIR / "index"
CACHE_FILE = KNOWLEDGE_DIR / "qa_cache.json"
CACHE_VERSION = "v3-detail-grounded"
DOCS_SOURCE = Path("/vol3/1000/Code/ngfchl.github.io/docs")
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen3:8b")
SITE_URL = "https://ptools.fun"

# 切分参数
CHUNK_SIZE = 900
CHUNK_OVERLAP = 120
TOP_K = 10

# 中文 embedding 模型
EMBED_MODEL = "BAAI/bge-small-zh-v1.5"

# 需要跳过的文件
SKIP_FILES = {"agreement.md", "index.md", "password.md", "termnal.md", "ssh.md"}
SKIP_SECTION_KEYWORDS = ["油猴", "tampermonkey", "greasemonkey", "userscript"]


# ============ Embedding ============

_model = None


def get_model():
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer
        cache_dir = str(KNOWLEDGE_DIR / "models")
        os.environ["CUDA_VISIBLE_DEVICES"] = ""
        print(f"📦 加载中文 embedding 模型: {EMBED_MODEL}")
        _model = SentenceTransformer(EMBED_MODEL, cache_folder=cache_dir)
        print("✅ 模型加载完成")
    return _model


def embed_texts(texts: list[str]):
    model = get_model()
    return model.encode(texts, normalize_embeddings=True, show_progress_bar=True,
                        batch_size=32, convert_to_numpy=True)


# ============ URL 生成 ============

def doc_path_to_url(rel_path: str) -> str:
    html_path = rel_path.replace(".md", ".html")
    parts = html_path.split("/")
    encoded = "/".join(quote(p, safe="") for p in parts)
    return f"{SITE_URL}/{encoded}"


# ============ 文档清洗（保留结构） ============

def clean_markdown(text: str) -> str:
    """清洗但保留代码块标记和链接"""
    text = re.sub(r'<[^>]+>', '', text)
    # 图片替换为 alt 文本
    text = re.sub(r'!\[([^\]]*)\]\([^\)]+\)', r'[图片:\1]', text)
    # 链接保留文本+URL
    text = re.sub(r'\[([^\]]+)\]\(([^\)]+)\)', r'\1 (\2)', text)
    # frontmatter
    text = re.sub(r'^---\n.*?---\n', '', text, flags=re.DOTALL)
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


# ============ 代码块检测 ============

def find_code_blocks(text: str) -> list[tuple[int, int]]:
    """找到所有代码块的起止位置"""
    blocks = []
    lines = text.split('\n')
    in_block = False
    start = -1

    for i, line in enumerate(lines):
        stripped = line.strip()
        if re.match(r'^```\w*$', stripped):
            if not in_block:
                in_block = True
                start = i
            else:
                blocks.append((start, i))
                in_block = False
    return blocks


# ============ 精细分块 ============

def split_document(text: str, rel_path: str, doc_title: str, doc_url: str) -> list[dict]:
    """
    精细分块策略：
    1. 代码块作为独立 chunk（不拆分）
    2. 每个标题层级下的内容独立成 chunk
    3. 大段落按 CHUNK_SIZE 拆分
    4. 每个 chunk 带上标题路径上下文
    """
    chunks = []
    lines = text.split('\n')

    # 先找代码块位置
    code_blocks = find_code_blocks(text)
    code_block_lines = set()
    for start, end in code_blocks:
        for i in range(start, end + 1):
            code_block_lines.add(i)

    # 按标题切分
    sections = []  # [(level, heading, line_start, line_end)]
    current_heading_stack = []  # [(level, heading)]
    current_start = 0

    for i, line in enumerate(lines):
        if i in code_block_lines:
            continue
        m = re.match(r'^(#{1,6})\s+(.+)', line.strip())
        if m:
            level = len(m.group(1))
            heading = m.group(2).strip()
            # 保存之前的 section
            if i > current_start:
                sections.append({
                    "level": current_heading_stack[-1][0] if current_heading_stack else 0,
                    "heading": current_heading_stack[-1][1] if current_heading_stack else "",
                    "path": [h for _, h in current_heading_stack],
                    "line_start": current_start,
                    "line_end": i - 1,
                })
            # 更新 heading stack
            while current_heading_stack and current_heading_stack[-1][0] >= level:
                current_heading_stack.pop()
            current_heading_stack.append((level, heading))
            current_start = i

    # 最后一个 section
    sections.append({
        "level": current_heading_stack[-1][0] if current_heading_stack else 0,
        "heading": current_heading_stack[-1][1] if current_heading_stack else "",
        "path": [h for _, h in current_heading_stack],
        "line_start": current_start,
        "line_end": len(lines) - 1,
    })

    chunk_id = 0

    for sec in sections:
        sec_lines = lines[sec["line_start"]:sec["line_end"] + 1]
        sec_text = '\n'.join(sec_lines).strip()

        if len(sec_text) < 15:
            continue

        # 检查是否包含油猴关键词
        if any(kw in sec_text.lower() for kw in SKIP_SECTION_KEYWORDS):
            continue

        # 构建标题路径
        title_path = " > ".join(sec["path"]) if sec["path"] else doc_title
        if title_path and title_path != doc_title:
            title_path = f"{doc_title} > {title_path}"
        else:
            title_path = doc_title

        # 代码块独立成 chunk
        code_text_parts = []
        non_code_parts = []
        in_code = False
        current_block = []

        for line in sec_lines:
            stripped = line.strip()
            if re.match(r'^```\w*$', stripped):
                if not in_code:
                    in_code = True
                    current_block = [line]
                else:
                    current_block.append(line)
                    code_text = '\n'.join(current_block)
                    if len(code_text.strip()) > 20:
                        code_text_parts.append(code_text)
                    current_block = []
                    in_code = False
                continue

            if in_code:
                current_block.append(line)
                continue

            non_code_parts.append(line)

        # 代码块 chunk
        for code_text in code_text_parts:
            chunks.append({
                "id": f"{rel_path}#{chunk_id}",
                "text": f"[代码块] {title_path}\n{code_text}",
                "title": title_path,
                "doc_title": doc_title,
                "source": rel_path,
                "url": doc_url,
                "heading": sec["heading"],
            })
            chunk_id += 1

        # 非代码部分 chunk
        non_code_text = '\n'.join(non_code_parts).strip()
        if not non_code_text or len(non_code_text) < 15:
            continue

        if len(non_code_text) <= CHUNK_SIZE:
            chunks.append({
                "id": f"{rel_path}#{chunk_id}",
                "text": f"{title_path}\n{non_code_text}",
                "title": title_path,
                "doc_title": doc_title,
                "source": rel_path,
                "url": doc_url,
                "heading": sec["heading"],
            })
            chunk_id += 1
        else:
            # 按段落拆分大文本
            paragraphs = non_code_text.split('\n\n')
            current = ""
            for para in paragraphs:
                para = para.strip()
                if not para:
                    continue
                if len(current) + len(para) + 2 <= CHUNK_SIZE:
                    current += ("\n\n" if current else "") + para
                else:
                    if current:
                        chunks.append({
                            "id": f"{rel_path}#{chunk_id}",
                            "text": f"{title_path}\n{current}",
                            "title": title_path,
                            "doc_title": doc_title,
                            "source": rel_path,
                            "url": doc_url,
                            "heading": sec["heading"],
                        })
                        chunk_id += 1
                    if len(para) > CHUNK_SIZE:
                        # 按行拆
                        sub_lines = para.split('\n')
                        current = ""
                        for sl in sub_lines:
                            if len(current) + len(sl) + 1 <= CHUNK_SIZE:
                                current += ("\n" if current else "") + sl
                            else:
                                if current:
                                    chunks.append({
                                        "id": f"{rel_path}#{chunk_id}",
                                        "text": f"{title_path}\n{current}",
                                        "title": title_path,
                                        "doc_title": doc_title,
                                        "source": rel_path,
                                        "url": doc_url,
                                        "heading": sec["heading"],
                                    })
                                    chunk_id += 1
                                current = sl
                    else:
                        current = para
            if current.strip():
                chunks.append({
                    "id": f"{rel_path}#{chunk_id}",
                    "text": f"{title_path}\n{current}",
                    "title": title_path,
                    "doc_title": doc_title,
                    "source": rel_path,
                    "url": doc_url,
                    "heading": sec["heading"],
                })
                chunk_id += 1

    return chunks


# ============ 文档加载 ============

def load_all_docs() -> list[dict]:
    docs = []
    target_dirs = ("收割机", "通用教程")

    for md_file in sorted(DOCS_SOURCE.rglob("*.md")):
        rel_path = str(md_file.relative_to(DOCS_SOURCE))
        parts = md_file.relative_to(DOCS_SOURCE).parts

        if parts[0] not in target_dirs:
            continue
        if md_file.name in SKIP_FILES:
            continue

        try:
            content = md_file.read_text(encoding="utf-8")
            doc_title = ""
            title_match = re.search(r'^title:\s*(.+)$', content, re.MULTILINE)
            if title_match:
                doc_title = title_match.group(1).strip().strip('"').strip("'")
                doc_title = re.sub(r'^\d+\.\s*', '', doc_title)

            body = clean_markdown(content)
            if len(body) < 30:
                continue
            if not doc_title:
                doc_title = md_file.stem

            doc_url = doc_path_to_url(rel_path)
            chunks = split_document(body, rel_path, doc_title, doc_url)
            docs.extend(chunks)

            code_count = sum(1 for c in chunks if "[代码块]" in c["text"])
            print(f"  📄 {doc_title}: {len(chunks)} chunks ({code_count} 代码块)")

        except Exception as e:
            print(f"⚠️ 跳过 {md_file}: {e}")

    return docs


# ============ 索引管理 ============

def rebuild_index():
    INDEX_DIR.mkdir(parents=True, exist_ok=True)
    docs = load_all_docs()
    if not docs:
        print("❌ 没有找到文档")
        return

    print(f"\n📚 加载了 {len(docs)} 个文档块")

    texts = [d["text"] for d in docs]
    print("🔄 正在生成中文向量...")
    embeddings = embed_texts(texts)

    import numpy as np
    np.save(INDEX_DIR / "embeddings.npy", embeddings)
    (INDEX_DIR / "metadata.json").write_text(
        json.dumps(docs, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    save_qa_cache({})
    print(f"\n🎉 知识库索引完成，共 {len(docs)} 个文档块")


def load_metadata() -> list[dict]:
    """只加载轻量 metadata，运行时 lite 检索使用。"""
    meta_path = INDEX_DIR / "metadata.json"
    if not meta_path.exists():
        return []
    return json.loads(meta_path.read_text(encoding="utf-8"))


def load_index():
    """加载向量索引。仅 vector 模式/离线构建环境需要 numpy。"""
    import numpy as np

    emb_path = INDEX_DIR / "embeddings.npy"
    meta_path = INDEX_DIR / "metadata.json"
    if not emb_path.exists() or not meta_path.exists():
        return np.array([]), []
    embeddings = np.load(emb_path)
    metadata = json.loads(meta_path.read_text(encoding="utf-8"))
    return embeddings, metadata


# ============ 问答缓存 ============

def _cache_key(question: str) -> str:
    normalized = re.sub(r'\s+', '', question.strip().lower())
    return hashlib.md5(f"{CACHE_VERSION}:{normalized}".encode()).hexdigest()


def load_qa_cache() -> dict:
    if CACHE_FILE.exists():
        try:
            content = CACHE_FILE.read_text(encoding="utf-8")
            if content.strip():
                return json.loads(content)
        except Exception:
            return {}
    return {}


def save_qa_cache(cache: dict):
    CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)
    CACHE_FILE.write_text(json.dumps(cache, ensure_ascii=False, indent=2), encoding="utf-8")


def get_cached_answer(question: str) -> str | None:
    cache = load_qa_cache()
    entry = cache.get(_cache_key(question))
    if entry:
        return entry.get("answer")
    return None


def set_cached_answer(question: str, answer: str):
    cache = load_qa_cache()
    cache[_cache_key(question)] = {
        "question": question,
        "answer": answer,
    }
    if len(cache) > 500:
        oldest = min(cache.keys(), key=lambda k: cache[k].get("question", ""))
        del cache[oldest]
    save_qa_cache(cache)


# ============ 关键词→文档映射 ============

KEYWORD_DOC_MAP = {
    "安装": ["收割机/install.md"],
    "部署": ["收割机/install.md", "收割机/使用指南.md"],
    "docker": ["收割机/install.md", "收割机/使用指南.md"],
    "compose": ["收割机/install.md", "收割机/compose.md"],
    "群晖": ["收割机/install.md"],
    "飞牛": ["收割机/install.md"],
    "cookie": ["收割机/import.md", "通用教程/CookieCloud.md"],
    "导入": ["收割机/import.md"],
    "站点导入": ["收割机/import.md"],
    "签到": ["收割机/common.md"],
    "通知": ["收割机/common.md", "收割机/home.md"],
    "下载器": ["收割机/common.md", "收割机/home.md"],
    "app": ["收割机/app.md"],
    "辅种": ["收割机/common.md"],
    "主题": ["收割机/custom-theme.md"],
    "自定义": ["收割机/custom-add.md", "收割机/custom-theme.md"],
    "常见问题": ["收割机/common.md"],
    "微信": ["收割机/wechat.md"],
    "baidu": ["通用教程/baidu-ocr.md"],
    "ocr": ["通用教程/baidu-ocr.md"],
    "frp": ["通用教程/tunnel.md"],
    "内网穿透": ["通用教程/tunnel.md"],
    "tailscale": ["通用教程/tailscale-nat.md"],
    "telegram": ["通用教程/telegram-bot.md"],
    "阿里云": ["通用教程/aliyun.md"],
    "emby": ["收割机/common.md"],
    "做种": ["收割机/common.md"],
    "日志": ["收割机/common.md"],
    "代理": ["收割机/common.md"],
    "ua": ["收割机/common.md"],
    "登录": ["收割机/使用指南.md"],
    "添加站点": ["收割机/import.md", "收割机/custom-add.md"],
    "计划任务": ["收割机/home.md", "收割机/使用指南.md"],
    "辅种助手": ["收割机/import.md"],
    "种子迁移": ["收割机/使用指南.md"],
    "数据库": ["收割机/install.md"],
    "postgresql": ["收割机/install.md"],
    "环境变量": ["收割机/install.md"],
    "webhook": ["收割机/common.md"],
    "授权": ["收割机/common.md", "收割机/install.md"],
    "密码": ["收割机/使用指南.md"],
    "token": ["收割机/使用指南.md"],
}


def keyword_boost(query: str) -> list[str]:
    query_lower = query.lower()
    matched = []
    for keyword, docs in KEYWORD_DOC_MAP.items():
        if keyword in query_lower:
            matched.extend(docs)
    return list(dict.fromkeys(matched))


# ============ 混合检索 ============

def tokenize_query(query: str) -> list[str]:
    """轻量中文/英文分词：无需 jieba，运行镜像零重依赖。"""
    q = query.lower().strip()
    tokens: list[str] = []

    # 英文、数字、版本号等
    tokens.extend(re.findall(r"[a-z0-9][a-z0-9_.:-]{1,}", q))

    # 连续中文词 + 2~4 字滑窗，兼顾“怎么安装/反代/环境变量”等短问法
    for block in re.findall(r"[\u4e00-\u9fff]+", q):
        if len(block) <= 4:
            tokens.append(block)
        else:
            tokens.append(block)
            for n in (2, 3, 4):
                tokens.extend(block[i:i + n] for i in range(0, len(block) - n + 1))

    # 去重并过滤太泛的单字
    return list(dict.fromkeys(t for t in tokens if len(t) >= 2))


def lite_search(query: str, top_k: int = TOP_K) -> list[dict]:
    """轻量关键词/BM25-like 检索。

    运行时不加载 sentence-transformers/torch，大幅降低 Docker 镜像体积。
    对当前 200~500 个文档块规模，关键词+领域词表加权的性价比最高。
    """
    metadata = load_metadata()
    if not metadata:
        return []

    query_lower = query.lower().strip()
    tokens = tokenize_query(query)
    boosted_sources = set(keyword_boost(query))
    scored: list[tuple[float, int]] = []

    for i, doc in enumerate(metadata):
        title = (doc.get("title") or "").lower()
        source = (doc.get("source") or "").lower()
        text = (doc.get("text") or "").lower()
        hay = f"{title}\n{source}\n{text}"

        score = 0.0

        # 领域词表强加权
        if doc.get("source") in boosted_sources:
            score += 20.0

        # 完整问题/短语命中
        if query_lower and query_lower in hay:
            score += 6.0
        if query_lower and query_lower in title:
            score += 5.0

        # token 频次 + 字段权重
        for token in tokens:
            if token in title:
                score += 4.0
            if token in source:
                score += 3.0
            count = text.count(token)
            if count:
                # 类 BM25：频次递减，避免长文刷分
                score += min(3.0, 1.0 + count ** 0.5)

        # 配置/安装类问题：代码块和关键配置变量优先，避免模型拿不到原文示例后自行编造
        config_query = any(k in query_lower for k in ["compose", "docker", "postgres", "postgresql", "安装", "部署", "配置", "环境变量"])
        if config_query and "[代码块]" in text:
            score += 10.0
            config_hits = sum(1 for k in [
                "image:", "services:", "environment:", "volumes:", "ports:",
                "postgres_db", "postgres_user", "postgres_password",
                "email", "token", "go_web_port", "database", "db_",
            ] if k in text)
            score += min(20.0, config_hits * 3.0)

        # 查询中英文/数字 token 对安装、docker、postgresql 等问题很关键
        exact_tokens = re.findall(r"[a-z0-9][a-z0-9_.:-]{1,}", query_lower)
        for token in exact_tokens:
            if token in hay:
                score += 2.5

        if score > 0:
            scored.append((score, i))

    scored.sort(reverse=True, key=lambda x: x[0])

    docs = []
    for score, idx in scored:
        doc = metadata[idx].copy()
        doc["score"] = float(score)
        src_count = sum(1 for d in docs if d["source"] == doc["source"])
        if src_count < 2:
            docs.append(doc)
        if len(docs) >= top_k:
            break

    return docs


def vector_search(query: str, top_k: int = TOP_K) -> list[dict]:
    """向量检索：仅开发/宿主机模式使用，需要 sentence-transformers。"""
    embeddings, metadata = load_index()
    if len(embeddings) == 0:
        return []

    query_vec = get_model().encode([query], normalize_embeddings=True)
    scores = (embeddings @ query_vec.T).flatten()

    boosted_sources = keyword_boost(query)
    if boosted_sources:
        for i, doc in enumerate(metadata):
            if doc["source"] in boosted_sources:
                scores[i] += 0.3

    query_words = set(re.findall(r'[\u4e00-\u9fff]+', query))
    for i, doc in enumerate(metadata):
        text_lower = doc["text"].lower()
        matches = sum(1 for w in query_words if w in text_lower)
        if matches >= 2:
            scores[i] += 0.15 * matches

    top_indices = np.argsort(scores)[::-1][:top_k]

    docs = []
    for idx in top_indices:
        doc = metadata[idx].copy()
        doc["score"] = float(scores[idx])
        src_count = sum(1 for d in docs if d["source"] == doc["source"])
        if src_count < 2:
            docs.append(doc)
    return docs


def expand_context(docs: list[dict], window: int = 1, max_chars: int = 12000, max_sources: int = 5) -> list[dict]:
    """为命中的 chunk 补充同源相邻上下文，减少回答过短/断章取义。"""
    if not docs:
        return []

    metadata = load_metadata()
    by_id = {d.get("id"): d for d in metadata}
    expanded: list[dict] = []
    seen: set[str] = set()
    total_chars = 0

    def parse_id(doc_id: str) -> tuple[str, int] | None:
        if not doc_id or "#" not in doc_id:
            return None
        src, idx = doc_id.rsplit("#", 1)
        if not idx.isdigit():
            return None
        return src, int(idx)

    allowed_sources = list(dict.fromkeys(d.get("source") for d in docs if d.get("source")))[:max_sources]
    allowed_sources_set = set(allowed_sources)

    for doc in docs:
        if doc.get("source") not in allowed_sources_set:
            continue
        parsed = parse_id(doc.get("id", ""))
        candidates = [doc]
        if parsed:
            src, idx = parsed
            neighbors = []
            for offset in range(-window, window + 1):
                item = by_id.get(f"{src}#{idx + offset}")
                if item:
                    neighbors.append(item)
            candidates = neighbors or [doc]

        for item in candidates:
            doc_id = item.get("id", "")
            if doc_id in seen:
                continue
            text_len = len(item.get("text", ""))
            if total_chars + text_len > max_chars and expanded:
                continue
            copy = item.copy()
            copy.setdefault("score", doc.get("score", 0))
            expanded.append(copy)
            seen.add(doc_id)
            total_chars += text_len

    return expanded


def search(query: str, top_k: int = TOP_K, expand: bool = True, context_window: int = 1) -> list[dict]:
    mode = os.environ.get("KNOWLEDGE_SEARCH_MODE", "lite").lower()
    if mode in {"vector", "embedding", "semantic"}:
        docs = vector_search(query, top_k=top_k)
    else:
        docs = lite_search(query, top_k=top_k)
    return expand_context(docs, window=context_window) if expand else docs


def get_stats() -> dict:
    metadata = load_metadata()
    cache = load_qa_cache()
    return {
        "total_chunks": len(metadata),
        "cached_answers": len(cache),
        "site_url": SITE_URL,
        "search_mode": os.environ.get("KNOWLEDGE_SEARCH_MODE", "lite"),
        "embedding_model": EMBED_MODEL,
        "cache_version": CACHE_VERSION,
    }


# ============ CLI ============

if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "rebuild":
        rebuild_index()
    elif len(sys.argv) > 1 and sys.argv[1] == "search":
        query = " ".join(sys.argv[2:])
        results = search(query)
        for r in results:
            print(f"\n📄 {r['title']} ({r['source']}) [分数:{r.get('score', 0):.3f}]")
            print(f"   URL: {r['url']}")
            print(f"   {r['text'][:200]}...")
    elif len(sys.argv) > 1 and sys.argv[1] == "stats":
        print(json.dumps(get_stats(), ensure_ascii=False, indent=2))
    elif len(sys.argv) > 1 and sys.argv[1] == "cache":
        cache = load_qa_cache()
        for k, v in cache.items():
            print(f"\n🔑 {k[:8]}... Q: {v['question'][:50]}")
    else:
        print("用法:")
        print("  python -m knowledge rebuild      - 重建索引")
        print("  python -m knowledge search <查询>  - 搜索")
        print("  python -m knowledge stats        - 统计信息")
        print("  python -m knowledge cache        - 查看缓存")
