"""kf_flow.py — 驱动 kicad-flow 的 LangGraph 图（进程内，无 langgraph dev）。

图本体在 ~/skills_env.md 的 KICAD_FLOW_ROOT（app/main_graph.py），本脚本只做三件事：
  1) 用 SqliteSaver 打开持久化 checkpoint，按 thread_id 续跑；
  2) 把当前挂起的 interrupt 以结构化 JSON 打印给 Claude 会话；
  3) 把 Claude 执行完真实动作后的结论 Command(resume=...) 回填、推进到下一个 interrupt。

**数据库落点**：与 altium-kicad 不同，kicad-flow 的 checkpoint 库落在
<被调用的 KiCad 工程>/.kicad-flow/checkpoint.sqlite，不在图工程根下。所以每个板子
一个库，`list` 也是按工程列线程。落点由图自己的 app/config.py 闸门强制校验。

为什么不直接用工程里的 run.py：run.py 的 resume 只接受一个位置参数 JSON 字符串，
在 PowerShell 里转义引号极易出错。本脚本用 --set key=value 免转义回填。

退出码：0 = 图已结束(done)；2 = 停在 interrupt 等回填(waiting)；1 = 出错。
"""
from __future__ import annotations

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

# Windows 控制台默认 GBK，interrupt 里的中文 hint 会变乱码看不清。自己转 UTF-8，
# 别指望调用方每次都记得带 PYTHONIOENCODING。
for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, ValueError):
        pass


def _env_lookup(key: str, default: str | None = None) -> str:
    """环境变量优先，其次 ~/skills_env.md 的 `KEY = value` 行（机器相关参数的真源）。"""
    v = os.environ.get(key)
    if v:
        return v
    path = os.path.join(os.path.expanduser("~"), "skills_env.md")
    try:
        with open(path, encoding="utf-8") as f:
            for line in f:
                m = re.match(rf"\s*{key}\s*=\s*(.+?)(?:\s+#.*)?$", line)
                if m:
                    return m.group(1).strip()
    except OSError:
        pass
    if default is None:
        sys.exit(f"missing {key}: add it to ~/skills_env.md or set it as an env var")
    return default


DEFAULT_GRAPH_ROOT = _env_lookup("KICAD_FLOW_ROOT", r"E:\LangGraph_Design\kicad-flow")

#: 与 app/state.py 的 max_*_retry 一一对应
RETRY_LIMITS = {
    "partition": 3, "material": 5, "layout": 5, "vivado": 3, "mcu": 3,
    "footprint": 5, "drawio-sync": 3, "recheck": 3, "requirement": 3,
    "pcb-rules": 3,
    # adjust 超限不是 safety_end，而是 d144「放弃调整原理图」，按原原理图继续布线
    "adjust": 3,
}


# --------------------------------------------------------------------------- utils
def _out(*parts: str) -> None:
    print(*parts, flush=True)


def _fatal(msg: str) -> None:
    _out(f"[ERROR] {msg}")
    sys.exit(1)


def slug_for(prj_dir: str) -> str:
    """E:\\Kicad_design\\MyBoard_V1p0 -> kf-myboard-v1p0"""
    name = Path(prj_dir.rstrip("\\/")).name or "run"
    name = re.sub(r"[^0-9A-Za-z]+", "-", name).strip("-").lower()
    return f"kf-{name or 'run'}"


def coerce(raw: str):
    """--set 的值做最小类型推断：true/false→bool，纯数字→int，[ 或 { 开头→JSON。"""
    s = raw.strip()
    low = s.lower()
    if low in ("true", "yes"):
        return True
    if low in ("false", "no"):
        return False
    if re.fullmatch(r"-?\d+", s):
        return int(s)
    if s[:1] in ("[", "{"):
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            return raw
    return raw


def pending_interrupt(snap):
    for task in snap.tasks:
        if task.interrupts:
            return task.interrupts[0].value
    return None


def load_graph_root(graph_root: Path):
    """把图工程根加进 sys.path，返回 (build_graph, config 模块)。"""
    if not (graph_root / "app" / "main_graph.py").is_file():
        _fatal(f"图工程根不对：{graph_root} 下没有 app\\main_graph.py。"
               f"用 --graph-root 指定，设环境变量 KICAD_FLOW_ROOT，"
               f"或改 ~/skills_env.md 的 KICAD_FLOW_ROOT。")
    sys.path.insert(0, str(graph_root))
    from app.main_graph import build_graph          # noqa: E402
    from app import config                          # noqa: E402
    return build_graph, config


def init_state(layout: dict, thread_id: str, limits: dict[str, int]) -> dict:
    state = {**layout, "thread_id": thread_id,
             "references": [], "sim_results": [], "history": []}
    for name, default in RETRY_LIMITS.items():
        key = name.replace("-", "_")
        state[f"{key}_retry"] = 0
        state[f"max_{key}_retry"] = limits.get(name) or default
    return state


# --------------------------------------------------------------------------- report
def report(graph, thread: dict, before_len: int) -> int:
    """打印当前状态；返回进程退出码。"""
    snap = graph.get_state(thread)
    values = snap.values or {}
    history = values.get("history", [])
    walked = [h.get("node") for h in history[before_len:]]
    if walked:
        _out("[NODES] " + " -> ".join(str(n) for n in walked))

    _out(f"[THREAD] {thread['configurable']['thread_id']}")
    if values.get("prj_dir"):
        _out(f"[PATHS]  prj_dir  = {values['prj_dir']}")
        _out(f"[PATHS]  db       = {values.get('flow_home', '')}\\checkpoint.sqlite")
        _out(f"[PATHS]  evidence = {values.get('runs_dir', '')}")
    for key, label in (("drawio_path", "drawio"), ("schematic_path", "sch"),
                       ("pcb_path", "pcb")):
        if values.get(key):
            _out(f"[FILES]  {label:8s} = {values[key]}")

    retries = {k: v for k, v in values.items()
               if k.endswith("_retry") and not k.startswith("max_") and v}
    if retries:
        _out("[RETRY]  " + json.dumps(retries))

    if snap.next:
        payload = pending_interrupt(snap)
        _out(f"[WAITING] next node: {', '.join(snap.next)}")
        if payload is not None:
            _out("[INTERRUPT] " + json.dumps(payload, ensure_ascii=False, indent=2))
        else:
            _out("[INTERRUPT] (无 payload — 直接 resume 推进)")
        return 2

    if values.get("halted"):
        _out(f"[SAFETY-END] {values.get('halted_reason', '')}")
        _out(f"[REPORT] {values.get('report_path', '')}")
        _out("先解决这个卡点，再用 `start --restart` 重跑，或用 --max-*-retry 调高上限。")
        return 0

    _out("[DONE] 图已跑完：原理图完成、PCB 已布局、关键线（阻抗/差分/等长）已布完。")
    _out(f"[REPORT] {values.get('report_path', '(未生成)')}")
    _out(f"[PCB]    {values.get('pcb_path', '(未回填)')}")
    return 0


# --------------------------------------------------------------------------- commands
def _open(args):
    """解析工程路径 + 落点，返回 (build_graph, db_path, thread, layout)。"""
    build_graph, config = load_graph_root(Path(args.graph_root))
    prj_raw = getattr(args, "prj_dir", None)
    if not prj_raw:
        _fatal("缺少 --prj-dir（目标 KiCad 工程根目录）。kicad-flow 的数据库就建在它下面，"
               "没有默认值，猜错会把库写到别处。")
    try:
        prj = config.resolve_project_dir(prj_raw)
        thread_id = args.thread or slug_for(str(prj))
        layout = config.prepare_project_layout(prj, thread_id)
    except config.ProjectDirError as exc:
        _fatal(str(exc))
    return build_graph, Path(layout["db_path"]), \
        {"configurable": {"thread_id": thread_id}}, layout


def cmd_start(args):
    from langgraph.checkpoint.sqlite import SqliteSaver
    build_graph, db_path, thread, layout = _open(args)

    if args.restart:
        n = wipe_thread(db_path, thread["configurable"]["thread_id"])
        _out(f"[RESTART] 清掉 {n} 条 checkpoint 记录")

    limits = {name: getattr(args, f"max_{name.replace('-', '_')}_retry", None)
              for name in RETRY_LIMITS}
    with SqliteSaver.from_conn_string(str(db_path)) as cp:
        graph = build_graph(checkpointer=cp)
        snap = graph.get_state(thread)
        before = len((snap.values or {}).get("history", []))
        if snap.next or snap.values:
            _out("[RESUME] 该工程已有进行中的线程，直接续跑（不重置）")
            graph.invoke(None, thread)
        else:
            graph.invoke(init_state(layout, thread["configurable"]["thread_id"], limits),
                         thread)
        return report(graph, thread, before)


def cmd_status(args):
    from langgraph.checkpoint.sqlite import SqliteSaver
    build_graph, db_path, thread, _ = _open(args)
    with SqliteSaver.from_conn_string(str(db_path)) as cp:
        graph = build_graph(checkpointer=cp)
        snap = graph.get_state(thread)
        return report(graph, thread, len((snap.values or {}).get("history", [])))


def cmd_resume(args):
    from langgraph.checkpoint.sqlite import SqliteSaver
    from langgraph.types import Command
    build_graph, db_path, thread, _ = _open(args)

    payload: dict = {}
    if args.json:
        try:
            parsed = json.loads(args.json)
        except json.JSONDecodeError as exc:
            _fatal(f"--json 不是合法 JSON：{exc}")
        if not isinstance(parsed, dict):
            _fatal("--json 必须是一个对象")
        payload.update(parsed)
    for item in args.set or []:
        if "=" not in item:
            _fatal(f"--set 要写成 KEY=VALUE，收到：{item}")
        k, v = item.split("=", 1)
        payload[k.strip()] = coerce(v)

    if not payload:
        _fatal("回填内容不能为空 —— langgraph 把 falsy 的 resume 当作『没提供』，"
               "会在同一个节点再中断一次，看起来像卡住。至少给一个键。")

    with SqliteSaver.from_conn_string(str(db_path)) as cp:
        graph = build_graph(checkpointer=cp)
        snap = graph.get_state(thread)
        if not snap.next:
            _fatal("当前没有挂起的 interrupt，无需回填。先跑 status 看状态。")
        # 判定键缺失时节点会把它当 False 处理、悄悄走另一条边（2026-09-23 DVBS2_dmod 的 d20
        # 就这样被跳到了 d22）。所以 interrupt 声明的第一个键（判定/主结论键）必须给。
        want = None
        for task in snap.tasks:
            if task.interrupts and isinstance(task.interrupts[0].value, dict):
                want = task.interrupts[0].value.get("resume")
        if isinstance(want, dict) and want and not args.allow_partial:
            key0 = next(iter(want))
            if key0 not in payload:
                _fatal(f"回填缺少本节点的主键 {key0}（{want[key0]}）。"
                       f"本节点声明的键：{list(want)}。确实只想填部分键就加 --allow-partial")
        before = len((snap.values or {}).get("history", []))
        _out(f"[RESUME] {json.dumps(payload, ensure_ascii=False)}")
        graph.invoke(Command(resume=payload), thread)
        return report(graph, thread, before)


def cmd_list(args):
    """列出某个工程库里的所有线程（DB 按工程隔离，所以要给 --prj-dir）。"""
    _, config = load_graph_root(Path(args.graph_root))
    try:
        prj = config.resolve_project_dir(args.prj_dir)
    except config.ProjectDirError as exc:
        _fatal(str(exc))
    db = config.db_path(prj)
    if not db.is_file():
        _out(f"[EMPTY] {prj} 下还没有 kicad-flow 库（{db} 不存在）")
        return 0
    con = sqlite3.connect(str(db))
    try:
        rows = con.execute(
            "SELECT thread_id, COUNT(*) FROM checkpoints GROUP BY thread_id"
        ).fetchall()
    except sqlite3.DatabaseError as exc:
        _fatal(f"读 {db} 失败：{exc}")
    finally:
        con.close()
    if not rows:
        _out("[EMPTY] 库里没有线程")
        return 0
    _out(f"[DB] {db}")
    for tid, n in rows:
        _out(f"  {tid:<28} {n} 条 checkpoint")
    return 0


def wipe_thread(db_path: Path, thread_id: str) -> int:
    if not db_path.is_file():
        return 0
    con = sqlite3.connect(str(db_path))
    total = 0
    try:
        for table in ("checkpoints", "writes"):
            try:
                cur = con.execute(f"DELETE FROM {table} WHERE thread_id = ?", (thread_id,))
                total += cur.rowcount or 0
            except sqlite3.DatabaseError:
                pass
        con.commit()
    finally:
        con.close()
    return total


def cmd_reset(args):
    _, config = load_graph_root(Path(args.graph_root))
    try:
        prj = config.resolve_project_dir(args.prj_dir)
    except config.ProjectDirError as exc:
        _fatal(str(exc))
    thread_id = args.thread or slug_for(str(prj))
    n = wipe_thread(config.db_path(prj), thread_id)
    _out(f"[RESET] {thread_id}：清掉 {n} 条记录")
    return 0


# --------------------------------------------------------------------------- main
def main() -> int:
    ap = argparse.ArgumentParser(prog="kf_flow", description="驱动 kicad-flow 的 LangGraph 图")
    ap.add_argument("--graph-root", default=os.environ.get("KICAD_FLOW_ROOT", DEFAULT_GRAPH_ROOT),
                    help="图工程根（默认取 ~/skills_env.md 的 KICAD_FLOW_ROOT）")
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("start", help="新建或续跑某个板子的原理图设计流程")
    p.add_argument("--prj-dir", help="目标 KiCad 工程根目录（数据库就建在它下面）")
    p.add_argument("--thread", help="线程 id（默认由 prj-dir 目录名推导）")
    p.add_argument("--restart", action="store_true", help="清空该线程重新开始")
    for name, default in RETRY_LIMITS.items():
        p.add_argument(f"--max-{name}-retry", type=int,
                       help=f"{name} 回环重试上限（默认 {default}）")
    p.set_defaults(func=cmd_start)

    p = sub.add_parser("status", help="看当前停在哪个 interrupt")
    p.add_argument("--prj-dir")
    p.add_argument("--thread")
    p.set_defaults(func=cmd_status)

    p = sub.add_parser("resume", help="回填结论并推进到下一个 interrupt")
    p.add_argument("--prj-dir")
    p.add_argument("--thread")
    p.add_argument("--set", action="append", metavar="KEY=VALUE",
                   help="免转义回填；true/false→bool，纯数字→int，[ 或 { 开头→JSON")
    p.add_argument("--json", help="直接给 JSON 对象（与 --set 可混用）")
    p.add_argument("--allow-partial", action="store_true",
                   help="允许不填 interrupt 声明的主键（默认缺主键直接报错）")
    p.set_defaults(func=cmd_resume)

    p = sub.add_parser("list", help="列出该工程库里的所有线程")
    p.add_argument("--prj-dir")
    p.set_defaults(func=cmd_list)

    p = sub.add_parser("reset", help="删除某个线程的 checkpoint")
    p.add_argument("--prj-dir")
    p.add_argument("--thread")
    p.set_defaults(func=cmd_reset)

    args = ap.parse_args()
    return args.func(args)


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