#!/usr/bin/env python3
"""
Wiki Curator — 自动知识策展管线 (P5)
扫描 Obsidian vault / markdown 文档，用启发式方法提取知识点，
通过织忆 API 存入结构性记忆。

Usage:
    python3 wiki_curator.py                 # 正常扫描并写入
    python3 wiki_curator.py --dry-run       # 预览（不写入 API）
    python3 wiki_curator.py --dir /tmp/md   # 指定目录
    python3 wiki_curator.py --force         # 忽略状态文件，全部重新处理
"""

import argparse
import hashlib
import json
import os
import re
import sys
from pathlib import Path

try:
    import requests
except ImportError:
    print("ERROR: 'requests' library is required. Install with: pip install requests")
    sys.exit(1)

try:
    import yaml
except ImportError:
    print("WARNING: 'yaml' library not available; LLM mode will use env var fallback")
    yaml = None

# ---------------------------------------------------------------------------
# 配置
# ---------------------------------------------------------------------------
ZHIYI_API = "http://localhost:7821"
ZHIYI_KEY = "zhiyi-dev-key-2026"
STATE_FILE = os.path.expanduser("~/.hermes/wiki_curator_state.json")

# LLM 配置
LLM_API = "http://127.0.0.1:3000/v1/chat/completions"
LLM_MODEL = "minimaxai/minimax-m2.7"  # m3 sometimes returns empty, use m2.7

# 扫描时排除的目录名称（大小写不敏感）
EXCLUDE_DIRS = {
    "__pycache__", ".git", "node_modules", ".obsidian", ".trash",
    "backups", ".gitlab", ".github", ".vscode", ".idea",
    "venv", ".venv", "env", ".env", "__pycache__",
}

# 最小文件长度（字符数）—— 太短的文件没有足够知识量
MIN_CHARS = 500


# ---------------------------------------------------------------------------
# 工具函数
# ---------------------------------------------------------------------------

def compute_sha256(content: str) -> str:
    """计算字符串的 SHA-256 摘要。"""
    return hashlib.sha256(content.encode("utf-8")).hexdigest()


def load_state() -> dict:
    """加载已处理文件的哈希状态。"""
    if os.path.isfile(STATE_FILE):
        try:
            with open(STATE_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError):
            print("  [WARN] 状态文件损坏，重置为空。")
    return {}


def save_state(state: dict):
    """保存处理状态到文件。"""
    os.makedirs(os.path.dirname(STATE_FILE), exist_ok=True)
    with open(STATE_FILE, "w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=False, indent=2)


def should_exclude_dir(dirname: str) -> bool:
    """检查目录名是否在排除列表中。"""
    return dirname.lower() in EXCLUDE_DIRS


def scan_md_files(scan_dir: str, force: bool, state: dict) -> list:
    """递归扫描 .md 文件，返回需要处理的 (相对路径, 绝对路径, 内容, 文件哈希)。"""
    scan_path = Path(scan_dir).expanduser().resolve()
    if not scan_path.is_dir():
        print(f"  [WARN] 目录不存在: {scan_path}")
        return []

    candidates = []
    for root_str, dirs, files in os.walk(str(scan_path)):
        # 过滤排除目录（原地修改 dirs 避免继续深入）
        dirs[:] = [d for d in dirs if not should_exclude_dir(d)]

        for fn in files:
            if not fn.endswith(".md"):
                continue
            # 跳过以下划线开头的文件（草稿/私有文件）
            if fn.startswith("_"):
                continue

            abs_path = Path(root_str) / fn
            rel_path = abs_path.relative_to(scan_path)

            try:
                content = abs_path.read_text(encoding="utf-8")
            except (OSError, UnicodeDecodeError) as e:
                print(f"  [WARN] 读取失败 {abs_path}: {e}")
                continue

            if len(content) < MIN_CHARS:
                print(f"  [SKIP]  {rel_path} (字符数 {len(content)} < {MIN_CHARS})")
                continue

            file_hash = compute_sha256(content)
            key = str(rel_path)

            if not force and state.get(key) == file_hash:
                # 文件未变更，跳过
                continue

            candidates.append((key, str(abs_path), content, file_hash))

    return candidates


# ---------------------------------------------------------------------------
# 知识提取（启发式 / 基于关键词）
# ---------------------------------------------------------------------------

def _heuristic_extract(content: str) -> dict:
    """
    从 markdown 内容中提取知识点（启发式方法）。
    返回结构：
        {
            "concepts": [{"name": "...", "summary": "..."}],
            "entities": [{"name": "...", "attributes": "..."}],
            "relations": [],
        }
    """
    concepts = []
    entities = []

    # 1. 提取 heading 作为概念名称
    #    # 标题, ## 标题, ### 标题 等
    heading_pattern = re.compile(r"^#{1,6}\s+(.+)$", re.MULTILINE)
    for match in heading_pattern.finditer(content):
        heading_text = match.group(1).strip()
        if not heading_text:
            continue
        # 收集该标题下的文本（直到下一个标题或文件末尾）
        start_pos = match.end()
        next_heading = heading_pattern.search(content, start_pos)
        if next_heading:
            section_content = content[start_pos:next_heading.start()].strip()
        else:
            section_content = content[start_pos:].strip()

        # 用段落第一句作为 summary
        summary = ""
        if section_content:
            # 取第一个非空段落作为摘要
            for line in section_content.split("\n"):
                line = line.strip()
                if line and not line.startswith("#") and not line.startswith("-"):
                    summary = line[:200]  # 截断
                    break

        # 过滤掉纯符号或过短的 heading 名称
        if len(heading_text) < 2:
            continue

        concepts.append({
            "name": heading_text,
            "summary": summary,
        })

    # 2. 提取 **粗体** 关键词作为实体
    bold_pattern = re.compile(r"\*\*(.+?)\*\*")
    seen_bolds = set()
    for match in bold_pattern.finditer(content):
        bold_text = match.group(1).strip()
        if not bold_text or len(bold_text) > 80:
            continue
        if bold_text.lower() in seen_bolds:
            continue
        seen_bolds.add(bold_text.lower())

        # 收集该加粗词所在的上下文（前后各 50 字符）
        start = max(0, match.start() - 50)
        end = min(len(content), match.end() + 50)
        # 在上下文截断到行边界
        context = content[start:end].replace("\n", " ").strip()
        # 清理多余的空白
        context = re.sub(r"\s+", " ", context)

        # 标记这个实体
        entities.append({
            "name": bold_text,
            "attributes": context[:200],  # 上下文作为属性描述
        })

    # 3. 提取列表项中的重要短语（- 或 * 开头的行，但不包含 ** 的内容）
    #    这里只处理包含中文字符或关键术语的项
    list_pattern = re.compile(r"^[\s]*[-*]\s+(.+)$", re.MULTILINE)
    seen_list_items = set()
    for match in list_pattern.finditer(content):
        item_text = match.group(1).strip()
        # 跳过空项、纯链接、纯图片
        if not item_text or item_text.startswith("[") or item_text.startswith("!"):
            continue
        # 跳过以冒号/分号结尾的短项
        if len(item_text) < 4:
            continue
        if item_text.lower() in seen_list_items:
            continue
        seen_list_items.add(item_text.lower())

        # 提取冒号前面的部分作为实体名，后面作为描述
        if "：" in item_text or ":" in item_text:
            parts = re.split(r"[：:]", item_text, maxsplit=1)
            name = parts[0].strip()
            desc = parts[1].strip() if len(parts) > 1 else ""
        else:
            # 如果列表中包含 **加粗**，使用加粗内容作为名称
            bm = re.search(r"\*\*(.+?)\*\*", item_text)
            if bm:
                name = bm.group(1).strip()
                # 从列表中去除加粗标记，作为描述
                desc = re.sub(r"\*\*(.+?)\*\*", r"\1", item_text)
            else:
                name = item_text[:60]
                desc = item_text

        if len(name) < 2:
            continue

        entities.append({
            "name": name,
            "attributes": desc[:200],
        })

    return {"concepts": concepts, "entities": entities}


# ---------------------------------------------------------------------------
# LLM 知识提取
# ---------------------------------------------------------------------------

def _get_llm_key() -> str:
    """从 config.yaml 读取 NewAPI key"""
    if yaml is not None:
        try:
            cfg_path = os.path.expanduser("~/.hermes/config.yaml")
            with open(cfg_path) as f:
                cfg = yaml.safe_load(f)
            raw = cfg.get("providers", {}).get("newapi-local", {}).get("api_key", "")
            # NewAPI 的 key 不需要 sk- 前缀
            if raw.startswith("sk-"):
                raw = raw[3:]
            return raw
        except Exception:
            pass
    # 回退到环境变量
    return os.environ.get("NEWAPI_API_KEY", "")


def _extract_with_llm(content: str, filepath: str) -> dict | None:
    """调用 NewAPI LLM 提取结构化知识"""
    llm_key = _get_llm_key()
    if not llm_key:
        print("  ⚠️ No LLM API key found (check config.yaml or NEWAPI_API_KEY env)")
        return None

    prompt = f'''Extract concepts, entities, and relations from this document.
Return ONLY valid JSON:
{{"concepts":[{{"name":"...","summary":"..."}}],"entities":[{{"name":"...","attributes":{{}}}}],"relations":[{{"source":"...","relation":"uses|contains|depends_on|part_of|implements","target":"..."}}]}}

Document:
{content[:2000]}
'''
    try:
        resp = requests.post(LLM_API,
            headers={"Authorization": f"Bearer {llm_key}", "Content-Type": "application/json"},
            json={
                "model": LLM_MODEL,
                "messages": [
                    {"role": "system", "content": "You are a knowledge extraction assistant. Always respond with valid JSON only."},
                    {"role": "user", "content": prompt}
                ],
                "temperature": 0.1,
                "max_tokens": 1000,
            },
            timeout=30)
        data = resp.json()
        if "error" in data and data["error"].get("message"):
            print(f"  ⚠️ LLM API error: {data['error']['message'][:60]}")
            return None
        choices = data.get("choices", [])
        if not choices:
            print("  ⚠️ LLM returned empty choices (API/model may be unavailable)")
            return None
        msg = choices[0].get("message", {})
        text = msg.get("content", "") or ""
        if not text.strip():
            finish = choices[0].get("finish_reason", "")
            print(f"  ⚠️ LLM returned empty content (finish={finish})")
            return None
        # Parse JSON from response
        json_match = re.search(r'\{[\s\S]*\}', text)
        if json_match:
            return json.loads(json_match.group())
    except Exception as e:
        print(f"  ⚠️ LLM extraction failed: {e}")
    return None


# ---------------------------------------------------------------------------
# 织忆 API 交互
# ---------------------------------------------------------------------------

def commit_memory(content: str, category: str, metadata: dict,
                  dry_run: bool = False) -> bool:
    """写入一条记忆到织忆。返回 True 表示成功。"""
    if dry_run:
        print(f"    [DRY-RUN] 写入记忆: category={category}, "
              f"content='{content[:80]}...'")
        return True

    url = f"{ZHIYI_API}/api/v1/commit"
    headers = {
        "X-API-Key": ZHIYI_KEY,
        "Content-Type": "application/json",
    }
    payload = {
        "agent_id": "wiki-curator",
        "content": content,
        "category": category,
        "metadata": metadata,
    }
    try:
        resp = requests.post(url, json=payload, headers=headers, timeout=30)
        if resp.status_code in (200, 201):
            return True
        else:
            print(f"    [FAIL] HTTP {resp.status_code}: {resp.text[:200]}")
            return False
    except requests.exceptions.RequestException as e:
        print(f"    [FAIL] 请求失败: {e}")
        return False


def commit_graph_edge(from_node: str, to_node: str, relation: str,
                      dry_run: bool = False) -> bool:
    """写入一条关系到织忆图谱。返回 True 表示成功。"""
    if dry_run:
        print(f"    [DRY-RUN] 写入关系: {from_node} --[{relation}]--> {to_node}")
        return True

    url = f"{ZHIYI_API}/api/v1/graph/edge"
    headers = {
        "X-API-Key": ZHIYI_KEY,
        "Content-Type": "application/json",
    }
    payload = {
        "from": from_node,
        "to": to_node,
        "relation": relation,
        "namespace": "wiki",
    }
    try:
        resp = requests.post(url, json=payload, headers=headers, timeout=30)
        if resp.status_code in (200, 201):
            return True
        else:
            print(f"    [FAIL] 关系写入 HTTP {resp.status_code}: {resp.text[:200]}")
            return False
    except requests.exceptions.RequestException as e:
        print(f"    [FAIL] 关系请求失败: {e}")
        return False


# ---------------------------------------------------------------------------
# 主流程
# ---------------------------------------------------------------------------

def process_file(rel_path: str, abs_path: str, content: str,
                 dry_run: bool = False, use_llm: bool = False) -> dict:
    """
    处理单个文件：提取知识点并写入织忆。
    返回统计信息。
    """
    print(f"\n  📄 {rel_path}")
    stats = {"concepts": 0, "entities": 0, "relations": 0}

    # 提取知识
    if use_llm:
        method_label = "LLM"
        result = _extract_with_llm(content, abs_path)
        if result:
            concepts = result.get("concepts", [])
            entities = result.get("entities", [])
            relations = result.get("relations", [])
            print(f"  🤖 LLM extracted {len(concepts)} concepts, {len(entities)} entities, {len(relations)} relations")
        else:
            print(f"  ⚠️ LLM failed for {os.path.basename(abs_path)}, falling back to heuristic")
            knowledge = _heuristic_extract(content)
            concepts = knowledge.get("concepts", [])
            entities = knowledge.get("entities", [])
            relations = knowledge.get("relations", [])
            method_label = "heuristic (fallback)"
    else:
        method_label = "heuristic"
        knowledge = _heuristic_extract(content)
        concepts = knowledge.get("concepts", [])
        entities = knowledge.get("entities", [])
        relations = knowledge.get("relations", [])

    # 添加 source 字段
    for c in concepts:
        c["source"] = abs_path
    for e in entities:
        e["source"] = abs_path

    # 写入概念
    for conc in concepts:
        content_line = f"## {conc['name']}"
        if conc.get("summary"):
            content_line += f"\n{conc['summary']}"
        metadata = {
            "source": conc.get("source", abs_path),
            "concept_type": "concept",
        }
        ok = commit_memory(content_line, "wiki", metadata, dry_run=dry_run)
        if ok:
            stats["concepts"] += 1

    # 写入实体
    for ent in entities:
        content_line = f"### {ent['name']}"
        attrs = ent.get("attributes")
        if attrs:
            if isinstance(attrs, dict):
                attrs_str = json.dumps(attrs, ensure_ascii=False)
            else:
                attrs_str = str(attrs)
            content_line += f"\n{attrs_str}"
        metadata = {
            "source": ent.get("source", abs_path),
            "concept_type": "entity",
        }
        ok = commit_memory(content_line, "wiki", metadata, dry_run=dry_run)
        if ok:
            stats["entities"] += 1

    # 写入关系
    for rel in relations:
        source_node = rel.get("source", "")
        target_node = rel.get("target", "")
        relation_type = rel.get("relation", "RELATED_TO").upper()
        if source_node and target_node:
            ok = commit_graph_edge(source_node, target_node,
                                   relation_type, dry_run=dry_run)
            if ok:
                stats["relations"] += 1

    # 写入简单的概念-实体关系（仅在 heuristic 且无 relations 时作为补充）
    if not relations and concepts and entities:
        primary_concept = concepts[0]["name"]
        for ent in entities:
            ok = commit_graph_edge(primary_concept, ent["name"],
                                   "RELATED_TO", dry_run=dry_run)
            if ok:
                stats["relations"] += 1

    return stats


def main():
    parser = argparse.ArgumentParser(
        description="Wiki Curator — 自动知识策展管线 (P5)"
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="预览模式：显示将要处理的内容但不写入 API"
    )
    parser.add_argument(
        "--dir", type=str, default="~/mc/",
        help="扫描目录 (默认: ~/mc/)"
    )
    parser.add_argument(
        "--force", action="store_true",
        help="强制重新处理所有文件，忽略状态文件"
    )
    parser.add_argument(
        "--llm", action="store_true",
        help="Use LLM for concept/entity extraction (default: heuristic)"
    )
    args = parser.parse_args()

    print("=" * 60)
    print("  织忆 Wiki Curator — 知识策展管线")
    print("=" * 60)

    scan_dir = os.path.expanduser(args.dir)
    print(f"\n扫描目录: {scan_dir}")
    if args.dry_run:
        print("模式: 🔍 DRY RUN (仅预览，不写入)")
    if args.force:
        print("模式: 🔄 FORCE (忽略已有状态)")

    # 加载状态
    state = load_state() if not args.force else {}
    print(f"状态文件: {STATE_FILE}")
    print(f"已处理文件: {len(state)}")

    # 扫描文件
    candidates = scan_md_files(scan_dir, args.force, state)
    print(f"\n待处理文件: {len(candidates)}")

    total_stats = {"concepts": 0, "entities": 0, "relations": 0}
    new_state = dict(state)  # 保留旧状态，更新新处理过的

    for rel_path, abs_path, content, file_hash in candidates:
        stats = process_file(rel_path, abs_path, content, dry_run=args.dry_run, use_llm=args.llm)
        total_stats["concepts"] += stats["concepts"]
        total_stats["entities"] += stats["entities"]
        total_stats["relations"] += stats.get("relations", 0)

        # 更新状态（即使 dry-run 也记录，以便下次不重复扫描）
        if not args.dry_run:
            new_state[rel_path] = file_hash

    # 保存状态
    if not args.dry_run:
        save_state(new_state)
        print(f"\n状态已更新: {len(new_state)} 个文件记录")

    # 总结
    print("\n" + "=" * 60)
    print("  📊 处理总结")
    print(f"  处理文件数:   {len(candidates)}")
    print(f"  概念写入数:   {total_stats['concepts']}")
    print(f"  实体写入数:   {total_stats['entities']}")
    print(f"  关系写入数:   {total_stats['relations']}")
    print("=" * 60)

    print("\nWIKI_CURATOR_OK")
    return 0


if __name__ == "__main__":
    sys.exit(main())