"""ak_flow.py — 驱动 Altium→KiCad 迁移的 LangGraph 图（进程内，无 langgraph dev）。

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

为什么不直接用工程里的 run.py：run.py 的 resume 只接受一个位置参数 JSON 字符串，
在 PowerShell 里转义引号极易出错；且 thread 固定为 AK_THREAD 单值，多工程并发会串。
本脚本用 --set key=value 免转义回填，thread 由 target_dir 推导，DB 沿用
runs/skill_checkpoint.sqlite（已跑完的 ak-xmod-v1p0 / ak-gq-exct7g-v1 仍可续）。

退出码：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

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)
    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


# 图工程根是机器相关的：真源是 ~/skills_env.md 的 ALTIUM_KICAD_GRAPH_ROOT。
DEFAULT_PROJECT_ROOT = _env_lookup("ALTIUM_KICAD_GRAPH_ROOT", r"E:\LangGraph_Design\Altium-Kicad")
DB_RELATIVE = ("runs", "skill_checkpoint.sqlite")


# --------------------------------------------------------------------------- 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(target_dir: str) -> str:
    """E:\\Kicad_design\\XMOD_V1p0 -> ak-xmod-v1p0  (与既有线程命名一致)。"""
    name = Path(target_dir.rstrip("\\/")).name or "run"
    name = re.sub(r"[^0-9A-Za-z]+", "-", name).strip("-").lower()
    return f"ak-{name or 'run'}"


def coerce(raw: str):
    """--set 的值做最小类型推断：true/false→bool，纯数字→int，其余原样字符串。"""
    low = raw.strip().lower()
    if low in ("true", "yes", "1"):
        return True
    if low in ("false", "no", "0"):
        return False
    if re.fullmatch(r"-?\d+", raw.strip()):
        return int(raw.strip())
    return raw


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


def load_graph(project_root: Path):
    if not (project_root / "app" / "main_graph.py").is_file():
        _fatal(f"工程根不对：{project_root} 下没有 app\\main_graph.py。"
               f"用 --project-root 指定，设环境变量 AK_PROJECT_ROOT，或改 ~/skills_env.md 的 ALTIUM_KICAD_GRAPH_ROOT。")
    sys.path.insert(0, str(project_root))
    from app.main_graph import build_graph  # noqa: E402
    return build_graph


RETRY_LIMITS = {"bridge": 3, "import": 3, "netlist": 5, "gerber": 5}


def init_state(source_prj: str, target_dir: str, limits: dict[str, int]) -> dict:
    state = {
        "source_prj": source_prj,
        "target_dir": target_dir,
        "migration_dir": str(Path(target_dir) / "_migration"),
        "history": [],
    }
    for name, default in RETRY_LIMITS.items():
        state[f"{name}_retry"] = 0
        state[f"max_{name}_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("source_prj"):
        _out(f"[PATHS]  source_prj = {values['source_prj']}")
        _out(f"[PATHS]  target_dir = {values['target_dir']}")
        _out(f"[PATHS]  migration  = {values.get('migration_dir', '')}")

    retries = {k: values.get(k) for k in
               ("bridge_retry", "import_retry", "netlist_retry", "gerber_retry")
               if values.get(k)}
    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 any(h.get("node") == "safety_exit" for h in history):
        _out("[SAFETY-END] 某个判定回环超过重试上限，图已优雅退出。"
             "先修好卡住的那一步，再用 restart 重跑，或调高对应 max_*_retry。")
        return 0

    _out("[DONE] 图已跑完。")
    _out(f"[REPORT] {values.get('report_path', '(未回填 report_path)')}")
    return 0


# --------------------------------------------------------------------------- commands
def cmd_start(args, graph_builder, db_path: Path):
    from langgraph.checkpoint.sqlite import SqliteSaver

    source = args.source_prj
    target = args.target_dir
    if not source or not target:
        _fatal("start 需要 --source-prj 与 --target-dir（均为绝对路径，无默认值）。")
    if not Path(source).is_file():
        _fatal(f"源工程不存在：{source}（应指向 .PrjPCB 文件本身）")
    Path(target).mkdir(parents=True, exist_ok=True)

    thread_id = args.thread or slug_for(target)
    thread = {"configurable": {"thread_id": thread_id}}

    with SqliteSaver.from_conn_string(str(db_path)) as cp:
        graph = graph_builder(checkpointer=cp)
        existing = graph.get_state(thread)
        if (existing.next or existing.values) and not args.restart:
            _out(f"[INFO] 线程 {thread_id} 已存在，续跑（不重置）。要从头开始请加 --restart。")
            before = len(existing.values.get("history", []))
            if existing.next and pending_interrupt(existing) is not None:
                # 已经停在 interrupt 上，什么都不推进，直接把 payload 亮出来
                return report(graph, thread, before)
            graph.invoke(None, thread)
            return report(graph, thread, before)

        if args.restart:
            _out(f"[INFO] --restart：清空线程 {thread_id} 的 checkpoint。")
            wipe_thread(db_path, thread_id)

        limits = {name: getattr(args, f"max_{name}_retry", None) for name in RETRY_LIMITS}
        _out(f"[INFO] 新建线程 {thread_id}")
        graph.invoke(init_state(source, target, limits), thread)
        return report(graph, thread, 0)


def cmd_status(args, graph_builder, db_path: Path):
    from langgraph.checkpoint.sqlite import SqliteSaver

    thread_id = resolve_thread(args)
    thread = {"configurable": {"thread_id": thread_id}}
    with SqliteSaver.from_conn_string(str(db_path)) as cp:
        graph = graph_builder(checkpointer=cp)
        snap = graph.get_state(thread)
        if not snap.values and not snap.next:
            _fatal(f"线程 {thread_id} 没有任何 checkpoint。先跑 start。")
        return report(graph, thread, len(snap.values.get("history", [])))


def cmd_resume(args, graph_builder, db_path: Path):
    from langgraph.checkpoint.sqlite import SqliteSaver
    from langgraph.types import Command

    payload: dict = {}
    if args.json:
        try:
            payload = json.loads(args.json)
        except json.JSONDecodeError as exc:
            _fatal(f"--json 不是合法 JSON：{exc}")
    for item in args.set or []:
        if "=" not in item:
            _fatal(f"--set 要写成 key=value，收到：{item}")
        key, raw = item.split("=", 1)
        payload[key.strip()] = coerce(raw)

    if not payload:
        _fatal("resume 的回填不能为空——langgraph 会把 falsy 的 resume 当作『没提供』"
               "从而在同一个节点再次中断。至少给一个键，例如 --set acknowledged=true。")

    thread_id = resolve_thread(args)
    thread = {"configurable": {"thread_id": thread_id}}
    with SqliteSaver.from_conn_string(str(db_path)) as cp:
        graph = graph_builder(checkpointer=cp)
        snap = graph.get_state(thread)
        if not snap.next:
            _fatal(f"线程 {thread_id} 没有挂起的 interrupt（图已结束或未启动）。先跑 status。")
        before = len(snap.values.get("history", []))
        _out("[RESUME] " + json.dumps(payload, ensure_ascii=False))
        graph.invoke(Command(resume=payload), thread)
        return report(graph, thread, before)


def cmd_list(args, graph_builder, db_path: Path):
    from langgraph.checkpoint.sqlite import SqliteSaver

    if not db_path.is_file():
        _out("[INFO] 还没有任何 checkpoint 数据库。")
        return 0
    conn = sqlite3.connect(str(db_path))
    try:
        threads = sorted({r[0] for r in conn.execute("select distinct thread_id from checkpoints")})
    finally:
        conn.close()
    if not threads:
        _out("[INFO] 数据库里没有线程。")
        return 0
    with SqliteSaver.from_conn_string(str(db_path)) as cp:
        graph = graph_builder(checkpointer=cp)
        for tid in threads:
            snap = graph.get_state({"configurable": {"thread_id": tid}})
            v = snap.values or {}
            state = f"waiting@{','.join(snap.next)}" if snap.next else "done"
            _out(f"[THREAD] {tid:<24} {state:<28} target={v.get('target_dir', '?')}")
    return 0


def cmd_reset(args, graph_builder, db_path: Path):
    thread_id = resolve_thread(args)
    n = wipe_thread(db_path, thread_id)
    _out(f"[INFO] 已删除线程 {thread_id} 的 {n} 条 checkpoint 记录。")
    return 0


def wipe_thread(db_path: Path, thread_id: str) -> int:
    if not db_path.is_file():
        return 0
    conn = sqlite3.connect(str(db_path))
    try:
        n = conn.execute("delete from checkpoints where thread_id=?", (thread_id,)).rowcount
        conn.execute("delete from writes where thread_id=?", (thread_id,))
        conn.commit()
    finally:
        conn.close()
    return n


def resolve_thread(args) -> str:
    if getattr(args, "thread", None):
        return args.thread
    if getattr(args, "target_dir", None):
        return slug_for(args.target_dir)
    _fatal("需要 --thread <id> 或 --target-dir <目标 KiCad 目录> 来定位线程。用 list 看所有线程。")


# --------------------------------------------------------------------------- main
def main() -> int:
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass

    ap = argparse.ArgumentParser(description="Drive the Altium->KiCad migration LangGraph flow.")
    ap.add_argument("--project-root", default=os.environ.get("AK_PROJECT_ROOT", DEFAULT_PROJECT_ROOT))
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("start", help="新建或续跑一次迁移")
    p.add_argument("--source-prj", help="源 Altium .PrjPCB 绝对路径")
    p.add_argument("--target-dir", help="目标 KiCad 工程目录绝对路径")
    p.add_argument("--thread", help="线程 id（默认由 target-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}，超限走 safety_end）")
    p.set_defaults(func=cmd_start)

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

    p = sub.add_parser("resume", help="回填结论并推进到下一个 interrupt")
    p.add_argument("--thread")
    p.add_argument("--target-dir")
    p.add_argument("--set", action="append", metavar="KEY=VALUE",
                   help="免转义回填，可重复；true/false→bool，纯数字→int")
    p.add_argument("--json", help="直接给 JSON 对象（与 --set 可混用）")
    p.set_defaults(func=cmd_resume)

    p = sub.add_parser("list", help="列出所有迁移线程")
    p.set_defaults(func=cmd_list)

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

    args = ap.parse_args()
    root = Path(args.project_root)
    db_path = root.joinpath(*DB_RELATIVE)
    db_path.parent.mkdir(parents=True, exist_ok=True)
    graph_builder = load_graph(root)
    return args.func(args, graph_builder, db_path)


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