#!/usr/bin/env python3
"""Generate and apply compact knowledge-base patches.

The pipeline is deliberately conservative:
- the agent generates compact `.patch.md` files only;
- this script applies only `更新动作: 新增` patches;
- each applied patch creates three core files in one run directory:
  `*.patch.md`, `*.backup.md.bak`, and `*.applied.md`;
- stable knowledge files are backed up before being modified.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import sys
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

BACKEND_ROOT = Path(__file__).resolve().parents[3]
KNOWLEDGE_PARENT = Path(__file__).resolve().parents[2]
HARNESS_ROOT = BACKEND_ROOT / "packages" / "harness"
for import_root in (BACKEND_ROOT, HARNESS_ROOT):
    import_root_str = str(import_root)
    if import_root_str not in sys.path:
        sys.path.insert(0, import_root_str)

KNOWLEDGE_ROOT = KNOWLEDGE_PARENT / "knowledge_base"
MAINTENANCE_ROOT = KNOWLEDGE_PARENT / "knowledge_generate" / "maintenance"
RUNS_ROOT = MAINTENANCE_ROOT / "runs"

DEFAULT_MODEL = "GLM-5"
DEFAULT_SKILLS = {"policy-knowledge", "update-knowledge"}
DEFAULT_CONFIG_PATH = BACKEND_ROOT / "config.yaml"


class PatchError(RuntimeError):
    pass


@dataclass
class ParsedPatch:
    patch_path: Path
    target_rel: Path
    action: str
    anchor_heading: str
    merge_content: str
    source_basis: str
    merge_judgment: str
    uncertain_items: str


def slug_from_target(target_rel: Path) -> str:
    return "__".join(target_rel.with_suffix("").parts)


def section(text: str, heading: str) -> str:
    pattern = re.compile(
        rf"^## {re.escape(heading)}\s*\n(?P<body>.*?)(?=^## |\Z)",
        re.M | re.S,
    )
    match = pattern.search(text)
    if not match:
        raise PatchError(f"missing section: ## {heading}")
    return match.group("body").strip()


def first_backtick_path(text: str) -> str:
    match = re.search(r"`([^`]+)`", text)
    if not match:
        raise PatchError("target file section does not contain a backtick path")
    return match.group(1).strip()


def normalize_target_path(raw: str) -> Path:
    raw_path = Path(raw)
    parts = raw_path.parts
    if parts and parts[0] in {"entities", "narratives", "events", "primitives", "policies", "playbooks", "reference", "analysis", "data_pointers", "routing"}:
        return raw_path
    if parts and parts[0] == "knowledge_base":
        return Path(*parts[1:])
    if parts and parts[0] == "/mnt" and len(parts) >= 3 and parts[1] == "knowledge":
        return Path(*parts[2:])
    if "knowledge_base" in parts:
        idx = parts.index("knowledge_base")
        return Path(*parts[idx + 1 :])
    raise PatchError(f"target path must point inside knowledge_base: {raw}")


def extract_merge_block(text: str) -> str:
    body = section(text, "可直接合并内容")
    match = re.search(r"```(?:markdown)?\s*\n(?P<content>.*?)\n```", body, re.S)
    if not match:
        raise PatchError("merge content section must contain a fenced markdown block")
    content = match.group("content").strip()
    if not content:
        raise PatchError("merge content block is empty")
    return content


def extract_anchor(text: str) -> str:
    body = section(text, "建议插入位置")
    headings = re.findall(r"`(#{2,6}\s+[^`]+)`", body)
    if not headings:
        raise PatchError("insert position must contain an existing heading in backticks")
    return headings[0].strip()


def parse_patch(patch_path: Path) -> ParsedPatch:
    text = patch_path.read_text(encoding="utf-8")
    target_rel = normalize_target_path(first_backtick_path(section(text, "目标文件")))
    action = section(text, "更新动作").splitlines()[0].strip(" 。.")
    if action != "新增":
        raise PatchError(f"only 新增 patches are supported, got: {action}")
    merge_judgment = section(text, "合并判断")
    if "暂不合并" in merge_judgment:
        raise PatchError("patch merge judgment says 暂不合并")
    return ParsedPatch(
        patch_path=patch_path,
        target_rel=target_rel,
        action=action,
        anchor_heading=extract_anchor(text),
        merge_content=extract_merge_block(text),
        source_basis=section(text, "来源依据"),
        merge_judgment=merge_judgment,
        uncertain_items=section(text, "不确定项"),
    )


def heading_level(heading: str) -> int:
    return len(heading) - len(heading.lstrip("#"))


def insert_content(target_text: str, anchor_heading: str, merge_content: str) -> str:
    if merge_content in target_text:
        raise PatchError("merge content already exists in target")

    anchor = re.search(rf"^{re.escape(anchor_heading)}\s*$", target_text, re.M)
    if not anchor:
        raise PatchError(f"anchor heading not found: {anchor_heading}")

    level = heading_level(anchor_heading)
    next_heading = re.search(rf"^#{{1,{level}}}\s+", target_text[anchor.end() :], re.M)
    insert_at = anchor.end() + next_heading.start() if next_heading else len(target_text)

    before = target_text[:insert_at].rstrip()
    after = target_text[insert_at:].lstrip("\n")
    return f"{before}\n\n{merge_content.strip()}\n\n{after}"


def write_applied_record(parsed: ParsedPatch, target_path: Path, backup_path: Path, record_path: Path) -> None:
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    record_path.write_text(
        "\n".join(
            [
                "# 知识库自动更新说明",
                "",
                "## 更新时间",
                "",
                now,
                "",
                "## Patch 文件",
                "",
                f"`{parsed.patch_path}`",
                "",
                "## 目标文件",
                "",
                f"`{target_path}`",
                "",
                "## 备份文件",
                "",
                f"`{backup_path}`",
                "",
                "## 更新动作",
                "",
                parsed.action,
                "",
                "## 插入位置",
                "",
                f"`{parsed.anchor_heading}` 章节内、下一个同级标题之前。",
                "",
                "## 来源依据",
                "",
                parsed.source_basis,
                "",
                "## 合并判断",
                "",
                parsed.merge_judgment,
                "",
                "## 不确定项",
                "",
                parsed.uncertain_items,
                "",
                "## 状态",
                "",
                "已自动合并，待人工通过 git diff 审核。",
                "",
            ]
        ),
        encoding="utf-8",
    )


def apply_patch_file(patch_path: Path, run_dir: Path, dry_run: bool = False) -> tuple[Path, Path | None, Path | None]:
    parsed = parse_patch(patch_path)
    target_path = KNOWLEDGE_ROOT / parsed.target_rel
    if not target_path.exists():
        raise PatchError(f"target file does not exist: {target_path}")

    original = target_path.read_text(encoding="utf-8")
    updated = insert_content(original, parsed.anchor_heading, parsed.merge_content)
    slug = slug_from_target(parsed.target_rel)
    backup_path = run_dir / f"{slug}.backup.md.bak"
    record_path = run_dir / f"{slug}.applied.md"

    if dry_run:
        return target_path, None, None

    backup_path.write_text(original, encoding="utf-8")
    target_path.write_text(updated, encoding="utf-8")
    write_applied_record(parsed, target_path, backup_path, record_path)
    return target_path, backup_path, record_path


def build_prompt(live_date: str, workspace_rel_dir: str, limit: int, topic_hint: str | None) -> str:
    topic_line = f"\n优先主题/目标：{topic_hint}\n" if topic_hint else ""
    return f"""
你是知识库维护 agent。请使用 update-knowledge skill，并可结合 policy-knowledge skill 的知识库读取规则。

任务：基于 `/mnt/knowledge/live/{live_date}.txt` 判断是否需要生成可自动合并的知识库 patch 文件。{topic_line}
严格要求：
1. 先读取 `/mnt/skills/custom/update-knowledge/SKILL.md`。
2. 读取 `/mnt/knowledge/INDEX.md`、`/mnt/knowledge/live/_MANIFEST.md`、`/mnt/knowledge/live/{live_date}.txt`。
3. 根据 live 条目识别需要更新的稳定知识文件，并读取必要目标文件。
4. 不要修改 `/mnt/knowledge` 下任何稳定知识正文。
5. 使用 `write_file` 将 patch 写入 `/mnt/user-data/workspace/{workspace_rel_dir}/`。
6. 根据当天 live 内容自主判断需要更新多少个稳定知识文件；允许生成 0 到 {limit} 个 `.patch.md` 文件。
7. `{limit}` 是上限，不是必须生成数量；不要为了凑数量而更新。
8. 如果当天内容与稳定知识库无关、只是短期噪声、或信息不稳定，请不要生成 patch。
9. 不要生成 `.update.md` 长篇报告。
10. 每个 patch 总长度控制在 800-1500 字。
11. 每个 patch 必须只服务一个目标文件。
12. patch 文件必须使用以下固定格式：

```markdown
# 知识库更新补丁候选

## 目标文件

`knowledge_base/.../...md`

## 更新动作

新增。

## 建议插入位置

建议插入在 `## 已存在的二级标题` 章节内、下一个同级标题之前。

## 可直接合并内容

```markdown
可直接进入 knowledge_base 正文的内容。这里不要写维护说明，不要写“建议”“我认为”。
```

## 来源依据

- `live/{live_date}.txt`：条目编号

## 合并判断

建议人工审核后合并。

## 不确定项

- 需要继续核验的事项；没有则写“无”。
```

硬性约束：
- `## 可直接合并内容` 是唯一允许写入稳定知识库的正文。
- 如果信息不稳定，不要生成 patch。
- 不要写长篇背景介绍、完整事件复述、主题分类表。
- `建议插入位置` 中第一个反引号标题必须是目标文件中已经存在的标题。
- 如果判断无需更新，最终回复说明“本日未生成 patch”即可，不要调用 write_file。

完成后，在最终回复中列出实际写入的 patch 文件。
"""


def event_to_jsonable(event) -> dict:
    return {"type": event.type, "data": dict(event.data or {})}


def resolve_secret(value: str | None) -> str | None:
    """Resolve direct values or env references.

    Supported forms:
    - ``sk-xxx``: returned as-is
    - ``$ENV_NAME``: read from environment
    - ``env:ENV_NAME``: read from environment
    """
    if not value:
        return None
    value = value.strip()
    if value.startswith("env:"):
        return os.getenv(value[4:])
    if value.startswith("$"):
        return os.getenv(value[1:])
    return value


def secret_config_value(value: str) -> str:
    """Return the value to write into temp config without expanding env refs."""
    value = value.strip()
    if value.startswith("env:"):
        return f"${value[4:]}"
    return value


def parse_header_values(values: list[str] | None) -> dict[str, str]:
    """Parse repeated ``--header Key=Value`` options."""
    headers: dict[str, str] = {}
    for raw in values or []:
        if "=" not in raw:
            raise ValueError(f"--header must use Key=Value format, got: {raw}")
        key, value = raw.split("=", 1)
        key = key.strip()
        if not key:
            raise ValueError(f"--header key cannot be empty: {raw}")
        headers[key] = value.strip()
    return headers


def parse_headers_json(raw: str | None) -> dict[str, str]:
    if not raw:
        return {}
    data = json.loads(raw)
    if not isinstance(data, dict):
        raise ValueError("--headers-json must be a JSON object")
    return {str(k): str(v) for k, v in data.items()}


def build_runtime_config(
    *,
    source_config: Path,
    run_dir: Path,
    config_model_name: str,
    provider_model_name: str | None,
    api_key: str | None,
    base_url: str | None,
    headers: dict[str, str],
) -> Path | None:
    """Create a temporary config.yaml with an injected model when overrides are provided.

    Returns ``None`` when no override is needed, so callers can use the normal
    DeerFlow config resolution.
    """
    if not any([provider_model_name, api_key, base_url, headers]):
        return None

    try:
        import yaml
    except ImportError as exc:
        raise RuntimeError("PyYAML is required for model override config generation") from exc

    config = yaml.safe_load(source_config.read_text(encoding="utf-8")) or {}
    models = list(config.get("models") or [])
    if not models:
        raise RuntimeError(f"source config has no models: {source_config}")

    base_model = dict(models[0])
    for model in models:
        if model.get("name") == config_model_name:
            base_model = dict(model)
            break

    base_model["name"] = config_model_name
    base_model["display_name"] = config_model_name
    base_model["use"] = "langchain_openai:ChatOpenAI"
    if provider_model_name:
        base_model["model"] = provider_model_name
    if api_key:
        resolved_api_key = resolve_secret(api_key)
        if not resolved_api_key:
            raise RuntimeError(f"api key is empty after env resolution: {api_key}")
        base_model["api_key"] = secret_config_value(api_key)
    if base_url:
        base_model["base_url"] = base_url
    if headers:
        # LangChain ChatOpenAI / OpenAI client uses default_headers.
        existing = base_model.get("default_headers") or {}
        if not isinstance(existing, dict):
            existing = {}
        base_model["default_headers"] = {**existing, **headers}
    base_model["supports_thinking"] = False
    base_model["supports_vision"] = bool(base_model.get("supports_vision", False))

    replaced = False
    new_models = []
    for model in models:
        if model.get("name") == config_model_name:
            new_models.append(base_model)
            replaced = True
        else:
            new_models.append(model)
    if not replaced:
        new_models.insert(0, base_model)
    config["models"] = new_models

    runtime_config = run_dir / "_runtime_config.yaml"
    runtime_config.write_text(yaml.safe_dump(config, allow_unicode=True, sort_keys=False), encoding="utf-8")
    return runtime_config


def generate_patches(
    *,
    live_date: str,
    run_dir: Path,
    model: str,
    config_path: Path | None,
    limit: int,
    topic_hint: str | None,
    keep_log: bool,
) -> list[Path]:
    from deerflow.client import DeerFlowClient
    from deerflow.config.paths import get_paths
    from deerflow.runtime.user_context import get_effective_user_id

    thread_id = f"knowledge-update-{live_date}-{int(time.time())}"
    workspace_rel_dir = f"knowledge_update_runs/{run_dir.name}"
    prompt = build_prompt(live_date, workspace_rel_dir, limit, topic_hint)

    client = DeerFlowClient(
        config_path=str(config_path) if config_path else None,
        model_name=model,
        thinking_enabled=False,
        subagent_enabled=False,
        plan_mode=False,
        available_skills=DEFAULT_SKILLS,
    )

    chunks: dict[str, list[str]] = {}
    last_ai_id = ""
    log_file = (run_dir / "_agent_run.jsonl").open("w", encoding="utf-8") if keep_log else None
    started = time.time()
    try:
        for event in client.stream(prompt, thread_id=thread_id, recursion_limit=160):
            if log_file:
                log_file.write(json.dumps(event_to_jsonable(event), ensure_ascii=False) + "\n")
                log_file.flush()

            if event.type == "messages-tuple":
                data = event.data
                if data.get("type") == "ai":
                    msg_id = data.get("id") or ""
                    content = data.get("content") or ""
                    if content:
                        chunks.setdefault(msg_id, []).append(content)
                        last_ai_id = msg_id
                    if data.get("tool_calls"):
                        names = [tc.get("name") for tc in data.get("tool_calls", []) if tc.get("name")]
                        if names:
                            print(f"[tool_call] {names}", flush=True)
                elif data.get("type") == "tool":
                    name = data.get("name")
                    content = (data.get("content") or "").replace("\n", " ")
                    print(f"[tool_result] {name}: {content[:200]}", flush=True)
            elif event.type == "end":
                print(f"[end] usage={event.data.get('usage')} elapsed={time.time() - started:.1f}s", flush=True)
    finally:
        if log_file:
            log_file.close()

    if keep_log:
        (run_dir / "_agent_final.md").write_text("".join(chunks.get(last_ai_id, [])), encoding="utf-8")

    workspace_dir = get_paths().sandbox_work_dir(thread_id, user_id=get_effective_user_id()) / workspace_rel_dir
    if not workspace_dir.exists():
        raise RuntimeError(f"agent workspace output not found: {workspace_dir}")

    copied: list[Path] = []
    for source in sorted(workspace_dir.glob("*.patch.md")):
        dest = run_dir / source.name
        shutil.copy2(source, dest)
        copied.append(dest)
    return copied


def apply_all_patches(run_dir: Path, dry_run: bool) -> list[tuple[Path, Path, Path | None, Path | None]]:
    results = []
    for patch_path in sorted(run_dir.glob("*.patch.md")):
        target, backup, record = apply_patch_file(patch_path, run_dir, dry_run=dry_run)
        results.append((patch_path, target, backup, record))
    return results


def create_run_dir(live_date: str, run_name: str | None) -> Path:
    name = run_name or f"{live_date}_glm5_{datetime.now().strftime('%H%M%S')}"
    run_dir = RUNS_ROOT / name
    run_dir.mkdir(parents=True, exist_ok=False)
    return run_dir


def main() -> int:
    parser = argparse.ArgumentParser(description="Generate and apply knowledge patches.")
    parser.add_argument("--live-date", default="2026-05-18", help="live file date, e.g. 2026-05-18")
    parser.add_argument("--model", default=DEFAULT_MODEL, help="DeerFlow runtime model config name")
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG_PATH, help="source config.yaml used to build runtime config")
    parser.add_argument("--api-key", help="override API key; supports direct value, $ENV_NAME, or env:ENV_NAME")
    parser.add_argument("--base-url", help="override OpenAI-compatible base URL")
    parser.add_argument("--model-name", help="override provider model name, e.g. Pro/zai-org/GLM-5.1")
    parser.add_argument("--header", action="append", help="extra HTTP header, repeatable, format Key=Value")
    parser.add_argument("--headers-json", help='extra HTTP headers as JSON object, e.g. \'{"X-Org":"abc"}\'')
    parser.add_argument("--limit", type=int, default=5, help="maximum number of patches to generate; agent may generate fewer, including 0")
    parser.add_argument("--topic", help="optional topic/target hint for the agent")
    parser.add_argument("--run-name", help="run directory name under maintenance/runs")
    parser.add_argument("--patch", type=Path, help="apply an existing patch instead of generating")
    parser.add_argument("--dry-run", action="store_true", help="validate but do not modify knowledge_base")
    parser.add_argument("--keep-log", action="store_true", help="keep agent stream log/final text in run dir")
    args = parser.parse_args()

    run_dir = create_run_dir(args.live_date, args.run_name)
    try:
        headers = {
            **parse_headers_json(args.headers_json),
            **parse_header_values(args.header),
        }
        runtime_config = build_runtime_config(
            source_config=args.config,
            run_dir=run_dir,
            config_model_name=args.model,
            provider_model_name=args.model_name,
            api_key=args.api_key,
            base_url=args.base_url,
            headers=headers,
        )
        if args.patch:
            patch_dest = run_dir / args.patch.name
            shutil.copy2(args.patch, patch_dest)
            patches = [patch_dest]
        else:
            patches = generate_patches(
                live_date=args.live_date,
                run_dir=run_dir,
                model=args.model,
                config_path=runtime_config,
                limit=args.limit,
                topic_hint=args.topic,
                keep_log=args.keep_log,
            )

        results = apply_all_patches(run_dir, dry_run=args.dry_run)
    except Exception:
        if not any(run_dir.iterdir()):
            run_dir.rmdir()
        raise

    print(f"run_dir: {run_dir}")
    print(f"patches: {len(patches)}")
    if not patches:
        print("status: no patches generated; no knowledge_base files changed")
        return 0
    for patch_path, target, backup, record in results:
        print(f"patch: {patch_path}")
        print(f"target: {target}")
        if args.dry_run:
            print("status: dry-run ok")
        else:
            print(f"backup: {backup}")
            print(f"applied: {record}")
            print("status: applied")
    return 0


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