"""In-process driver for the fpga-imx6-flow LangGraph workflow.

Builds the compiled graph from the project's app package and drives it with a
persistent SqliteSaver — no `langgraph dev` server. Mirrors the fpga-flow helper
convention: exit code 0 = done, 1 = error, 2 = interrupt fired.

Commands:
    start   --project-dir <dir> [--program-mode bitstream|flash] [--bit-path ..]
            [--flash-path ..] [--thread-id ..] [--real] [--max-iterations N]
    resume  --project-dir <dir> --resume-json '<json>' [--thread-id ..] [--real]
    status  --project-dir <dir> [--thread-id ..]

Output: one JSON line on stdout describing status + the pending interrupt (if any).
Run via the env python directly or `conda run --no-capture-output` to avoid the
Windows GBK console re-encode crash on Chinese interrupt payloads.
"""
from __future__ import annotations

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

try:
    sys.stdout.reconfigure(encoding="utf-8")  # emit UTF-8 even on a GBK console
except Exception:
    pass


def _env_lookup(key: str, default: str | None = None) -> str | None:
    """Read KEY from the OS environment, else from a ``KEY = value`` line in
    ~/skills_env.md. Machine-local paths are never hardcoded in a skill; see
    the same helper in fpga_flash/scripts/vivado_env.py."""
    value = os.environ.get(key)
    if value:
        return value
    path = os.path.join(os.path.expanduser("~"), "skills_env.md")
    try:
        with open(path, encoding="utf-8") as handle:
            for line in handle:
                match = re.match(rf"\s*{re.escape(key)}\s*=\s*(.+?)\s*(?:#.*)?$", line)
                if match:
                    found = match.group(1).strip()
                    if found and not found.startswith("<"):
                        return found
    except OSError:
        pass
    return default


# The fpga-imx6-flow checkout: FPGA_IMX6_FLOW_ROOT in ~/skills_env.md.
DEFAULT_PROJECT = _env_lookup("FPGA_IMX6_FLOW_ROOT", "")


def _emit(obj: dict, code: int) -> None:
    print(json.dumps(obj, ensure_ascii=False))
    sys.exit(code)


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


def main() -> None:
    ap = argparse.ArgumentParser(description="Drive the fpga-imx6-flow graph in-process.")
    ap.add_argument("command", choices=["start", "resume", "status"])
    ap.add_argument("--project-dir", default=DEFAULT_PROJECT)
    ap.add_argument("--thread-id", default="imx6-skill-1")
    ap.add_argument("--resume-json", default="")
    ap.add_argument("--program-mode", default="bitstream", choices=["bitstream", "flash"])
    ap.add_argument("--bit-path", default="")
    ap.add_argument("--flash-path", default="")
    ap.add_argument("--real", action="store_true",
                    help="drive real hardware (workflow tools_* must be wired; default is MOCK dry-run)")
    ap.add_argument("--max-iterations", type=int, default=8)
    ap.add_argument("--max-hw-retry", type=int, default=3)
    args = ap.parse_args()

    if not args.project_dir:
        _emit({"status": "error",
               "error": "fpga-imx6-flow 工作流仓库路径未配置：在 ~/skills_env.md 增加 "
                        "FPGA_IMX6_FLOW_ROOT = <检出目录>，或用 --project-dir 显式指定"}, 1)

    project = Path(args.project_dir)
    if not (project / "app" / "main_graph.py").exists():
        _emit({"status": "error",
               "error": f"app/main_graph.py not found under {project}; check --project-dir"}, 1)

    # MOCK by default so the design loop can be dry-run with no hardware attached.
    os.environ["FPGA_IMX6_MOCK"] = "0" if args.real else "1"
    sys.path.insert(0, str(project))

    try:
        from langgraph.checkpoint.sqlite import SqliteSaver
        from langgraph.types import Command
        from app.main_graph import build_graph
    except Exception as exc:  # noqa: BLE001
        _emit({"status": "error",
               "error": f"import failed: {exc}. Run /langgraph-vivado-install to set up the env."}, 1)

    runs = project / "runs"
    runs.mkdir(exist_ok=True)
    db = str(runs / "skill_checkpoint.sqlite")
    cfg = {"configurable": {"thread_id": args.thread_id}}

    with SqliteSaver.from_conn_string(db) as cp:
        graph = build_graph(checkpointer=cp)
        try:
            if args.command == "start":
                init = {
                    "iteration": 0,
                    "program_mode": args.program_mode,
                    "bit_path": args.bit_path,
                    "flash_path": args.flash_path,
                    "max_hw_retry": args.max_hw_retry,
                    "max_iterations": args.max_iterations,
                    "history": [],
                }
                graph.invoke(init, cfg)
            elif args.command == "resume":
                payload = json.loads(args.resume_json) if args.resume_json else {}
                snap0 = graph.get_state(cfg)
                pend = [intr for task in snap0.tasks for intr in task.interrupts]
                if len(pend) <= 1:
                    graph.invoke(Command(resume=payload), cfg)
                else:
                    # Multiple interrupts pending (parallel fan-out, e.g. modify_fpga +
                    # modify_imx6). payload must be a map keyed by each interrupt's "action"
                    # (or interrupt id); resume them all at once with an id->value map.
                    resume_map = {}
                    for intr in pend:
                        val = intr.value if isinstance(intr.value, dict) else {}
                        action = val.get("action")
                        if action in payload:
                            resume_map[intr.id] = payload[action]
                        elif intr.id in payload:
                            resume_map[intr.id] = payload[intr.id]
                        else:
                            resume_map[intr.id] = {}
                    graph.invoke(Command(resume=resume_map), cfg)
            # status: no invoke, just snapshot
            snap = graph.get_state(cfg)
        except Exception as exc:  # noqa: BLE001
            _emit({"status": "error", "error": str(exc)}, 1)

        if snap.next:
            all_intr = [intr.value for task in snap.tasks for intr in task.interrupts]
            _emit({
                "status": "interrupt",
                "thread_id": args.thread_id,
                "mock": not args.real,
                "next": list(snap.next),
                "iteration": snap.values.get("iteration"),
                "interrupt": _pending_interrupt(snap),
                "pending_interrupts": all_intr,
                "pending_count": len(all_intr),
            }, 2)
        else:
            _emit({
                "status": "done",
                "thread_id": args.thread_id,
                "iteration": snap.values.get("iteration"),
                "steps": len(snap.values.get("history", [])),
            }, 0)


if __name__ == "__main__":
    main()