"""Drive the sim-rtl-flow LangGraph workflow from a Claude Code session.

The graph never talks to an agent host directly: subagent work goes through the
``SubagentBackend`` protocol and the two mandatory human checkpoints go through
the ``ReviewGate`` protocol. This driver plugs a *record/replay* implementation
into both, so the agent host can be **this Claude session** and the reviewer can
be **the user**:

* when the graph asks for something that has no recorded answer yet, the driver
  stops the run, writes the pending request(s) to ``pending.json`` and prints
  them;
* Claude performs the subagent task (or relays the review request to the user)
  and calls ``record``;
* ``resume`` re-invokes the graph from scratch -- every already-recorded call is
  replayed from the ledger, so the graph deterministically walks back to where
  it stopped and asks for the next thing.

Replay-from-scratch is why no LangGraph checkpointer is used here: graph nodes
have no side effects of their own (all side effects live in the subagents, which
are replayed from the ledger), so a fresh invoke is both cheaper and immune to
partial-superstep resume semantics.

All run state lives under ``<work_dir>/.sim_rtl_flow/`` -- i.e. inside the
project being designed, never inside the sim-rtl-flow repository (Notes/PR01-2.md).

Commands
--------
  start   --input <json> [--thread-id ID]   begin a run
  resume                                    continue after record
  record  --key K --result '<json>'         record one subagent result
  record  --key K --approve | --reject      answer a human review checkpoint
  record  --all-ok                          default-success for pending SUBAGENTS
  status                                    show ledger / pending / last state
  dry-run --input <json>                    end-to-end ScriptedBackend rehearsal
  mermaid                                   print the graph as Mermaid
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import threading
import traceback
from collections import defaultdict
from pathlib import Path
from typing import Any


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


def _default_project() -> Path:
    """The sim-rtl-flow checkout: SIM_RTL_FLOW_ROOT in ~/skills_env.md."""
    root = _env_lookup("SIM_RTL_FLOW_ROOT")
    if root:
        return Path(root)
    raise SystemExit(
        "sim-rtl-flow 工作流仓库路径未配置：请在 ~/skills_env.md 增加\n"
        "  SIM_RTL_FLOW_ROOT = <sim-rtl-flow 检出目录>\n"
        "或用 --project 显式指定。"
    )

#: Run state lives inside the *calling* project, not in the workflow repo.
STATE_DIRNAME = ".sim_rtl_flow"

CHANGE_SCOPES = {
    "none",
    "simulink_module",
    "non_simulink_rtl",
    "soc",
    "hardware",
    "unknown",
}

RESULT_DEFAULTS: dict[str, Any] = {
    "success": True,
    "gate_passed": True,
    "artifacts": [],
    "notes": "",
    "timing_met": True,
    "hil_passed": True,
    "needs_code_change": False,
    "change_scope": "none",
    "affected_modules": [],
    "metadata": {},
}

REVIEW_DEFAULTS: dict[str, Any] = {"approved": True, "notes": "", "metadata": {}}


def _bootstrap(project: Path):
    src = project / "src"
    if not (src / "sim_rtl_flow").is_dir():
        raise SystemExit(f"sim_rtl_flow package not found under {src}")
    sys.path.insert(0, str(src))
    from sim_rtl_flow.backend import (  # noqa: E402
        ScriptedBackend,
        SubagentResult,
    )
    from sim_rtl_flow.graph import WorkflowContext, build_workflow  # noqa: E402
    from sim_rtl_flow.review import ReviewDecision, ScriptedReviewGate  # noqa: E402

    return {
        "ScriptedBackend": ScriptedBackend,
        "ScriptedReviewGate": ScriptedReviewGate,
        "SubagentResult": SubagentResult,
        "ReviewDecision": ReviewDecision,
        "WorkflowContext": WorkflowContext,
        "build_workflow": build_workflow,
    }


class HandoffRequired(Exception):
    """Raised inside a graph node when an answer is not recorded yet."""

    def __init__(self, key: str) -> None:
        super().__init__(f"handoff required: {key}")
        self.key = key


class Handoffs:
    """Shared record/replay ledger for subagent tasks and human reviews."""

    def __init__(self, ledger: dict[str, Any]) -> None:
        self._ledger = ledger
        self._counter: dict[str, int] = defaultdict(int)
        self._lock = threading.Lock()
        self.pending: list[dict[str, Any]] = []
        self.replayed: list[str] = []

    def resolve(self, name: str, describe) -> dict[str, Any]:
        """Return the recorded payload for ``name``'s next occurrence, or halt."""
        with self._lock:
            self._counter[name] += 1
            key = f"{name}#{self._counter[name]}"
            recorded = self._ledger.get(key)
            if recorded is not None:
                self.replayed.append(key)
                return recorded
            self.pending.append({"key": key, **describe()})
        raise HandoffRequired(key)


class ClaudeSessionBackend:
    """Replay recorded subagent results; hand unknown requests to Claude."""

    def __init__(self, handoffs: Handoffs, SubagentResult) -> None:
        self._handoffs = handoffs
        self._result_cls = SubagentResult

    def run(self, request):  # noqa: ANN001 - duck-typed SubagentRequest
        payload = self._handoffs.resolve(
            request.task,
            lambda: {
                "kind": "subagent",
                "task": request.task,
                "role": request.role,
                "skill": request.skill,
                "module": request.module,
                "allow_code_changes": request.allow_code_changes,
                "work_dir": request.work_dir,
                "prompt": request.prompt,
                "context": request.context,
            },
        )
        values = {**RESULT_DEFAULTS, **payload}
        values["artifacts"] = tuple(values["artifacts"])
        values["affected_modules"] = tuple(values["affected_modules"])
        return self._result_cls(**values)


class ClaudeSessionReviewGate:
    """Relay the two mandatory human checkpoints to the user via Claude."""

    def __init__(self, handoffs: Handoffs, ReviewDecision) -> None:
        self._handoffs = handoffs
        self._decision_cls = ReviewDecision

    def review(self, request):  # noqa: ANN001 - duck-typed ReviewRequest
        payload = self._handoffs.resolve(
            f"review:{request.checkpoint}",
            lambda: {
                "kind": "review",
                "checkpoint": request.checkpoint,
                "title": request.title,
                "work_dir": request.work_dir,
                "prompt": request.prompt,
                "context": request.context,
            },
        )
        values = {**REVIEW_DEFAULTS, **payload}
        return self._decision_cls(**values)


# --------------------------------------------------------------------------
# run directory helpers
# --------------------------------------------------------------------------


class RunDir:
    """``<work_dir>/.sim_rtl_flow/runs/<thread-id>/`` -- inside the design project."""

    def __init__(self, work_dir: Path, thread_id: str, *, create: bool = True) -> None:
        self.work_dir = work_dir.expanduser().resolve()
        self.thread_id = thread_id
        self.path = self.work_dir / STATE_DIRNAME / "runs" / thread_id
        if create:
            self.path.mkdir(parents=True, exist_ok=True)

    def _read(self, name: str, default: Any) -> Any:
        target = self.path / name
        if not target.exists():
            return default
        return json.loads(target.read_text(encoding="utf-8"))

    def _write(self, name: str, payload: Any) -> None:
        self.path.mkdir(parents=True, exist_ok=True)
        (self.path / name).write_text(
            json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
        )

    @property
    def workflow_input(self) -> dict[str, Any]:
        data = self._read("input.json", None)
        if data is None:
            raise SystemExit(
                f"no input.json in {self.path}\n"
                "run `start --input <file>` first, or point --work-dir at the "
                "design project whose .sim_rtl_flow/ holds this run"
            )
        return data

    def set_input(self, payload: dict[str, Any]) -> None:
        self._write("input.json", payload)

    @property
    def ledger(self) -> dict[str, Any]:
        return self._read("ledger.json", {})

    def set_ledger(self, payload: dict[str, Any]) -> None:
        self._write("ledger.json", payload)

    @property
    def pending(self) -> list[dict[str, Any]]:
        return self._read("pending.json", [])

    def set_pending(self, payload: list[dict[str, Any]]) -> None:
        self._write("pending.json", payload)

    def set_state(self, payload: Any) -> None:
        self._write("last_state.json", payload)

    @property
    def state(self) -> Any:
        return self._read("last_state.json", None)


def _resolve_work_dir(args, payload: dict[str, Any] | None = None) -> Path:
    """The calling project's directory -- where all run state must live."""
    if payload is not None:
        work_dir = payload.get("work_dir")
        if not work_dir:
            raise SystemExit("workflow input must define work_dir")
        resolved = Path(work_dir)
    elif args.work_dir is not None:
        resolved = Path(args.work_dir)
    else:
        resolved = Path.cwd()

    resolved = resolved.expanduser().resolve()
    project = Path(args.project).expanduser().resolve()
    if resolved == project or project in resolved.parents:
        raise SystemExit(
            f"work_dir 不能位于 sim-rtl-flow 工作流仓库内（{resolved}）。\n"
            "运行状态必须保存在调用这个工作流的那个工程目录下 —— "
            "请把 work_dir 指向真实的设计工程。"
        )
    return resolved


def _jsonable(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_jsonable(v) for v in value]
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value
    return str(value)


def _unwrap(exc: BaseException) -> HandoffRequired | None:
    seen: set[int] = set()
    stack: list[BaseException] = [exc]
    while stack:
        current = stack.pop()
        if current is None or id(current) in seen:
            continue
        seen.add(id(current))
        if isinstance(current, HandoffRequired):
            return current
        if isinstance(current, BaseExceptionGroup):  # noqa: F821 - py3.11+
            stack.extend(current.exceptions)
        stack.append(current.__cause__)
        stack.append(current.__context__)
    return None


# --------------------------------------------------------------------------
# printing
# --------------------------------------------------------------------------


def _print_pending(pending: list[dict[str, Any]]) -> None:
    reviews = [item for item in pending if item.get("kind") == "review"]
    tasks = [item for item in pending if item.get("kind") != "review"]

    print("\n" + "=" * 72)
    print(
        f"HANDOFF — {len(tasks)} subagent task(s), {len(reviews)} human review(s)"
    )
    print("=" * 72)

    for item in reviews:
        print("\n" + "!" * 72)
        print(f"!! 人工审核确认 — key: {item['key']}")
        print(f"!! {item['title']}")
        print("!! 必须把该问题原样转达用户并等待用户决定；")
        print("!! Claude 不得代替用户批准，也不得依据自己的判断放行。")
        print("!" * 72)
        print(f"    checkpoint : {item['checkpoint']}")
        print(f"    work_dir   : {item['work_dir']}")
        print(f"    问题       : {item['prompt']}")
        print("    context    : " + json.dumps(item["context"], ensure_ascii=False))

    for item in tasks:
        print(f"\n--- key: {item['key']}")
        print(f"    role        : {item['role']}")
        print(f"    skill       : {item['skill'] or '(none — do it inline)'}")
        print(f"    module      : {item['module'] or '-'}")
        print(f"    work_dir    : {item['work_dir']}")
        write = (
            "WRITE ALLOWED"
            if item["allow_code_changes"]
            else "READ-ONLY (no code edits!)"
        )
        print(f"    permission  : {write}")
        print(f"    prompt      : {item['prompt']}")
        print("    context     : " + json.dumps(item["context"], ensure_ascii=False))

    print("\n" + "-" * 72)
    print("Answer each, then resume:")
    for item in reviews:
        print(f"  record --key {item['key']} --approve            # 用户批准后")
        print(f"  record --key {item['key']} --reject --notes '…' # 用户驳回后")
    for item in tasks:
        print(
            f"  record --key {item['key']} "
            "--result '{\"gate_passed\": true, \"notes\": \"...\"}'"
        )
    print("  resume")
    print("-" * 72)


def _print_final(state: Any, replayed: list[str]) -> None:
    status = state.get("status") if isinstance(state, dict) else None
    print("\n" + "=" * 72)
    print(f"WORKFLOW FINISHED — status: {status}")
    print("=" * 72)
    if isinstance(state, dict):
        if state.get("failure_reason"):
            print(f"failure_reason : {state['failure_reason']}")
        for label, key in (("人工确认 · 仿真性能", "sim_review"),
                           ("人工确认 · 功能划分", "partition_review")):
            review = state.get(key)
            if review:
                mark = "批准" if review.get("approved") else "驳回"
                print(f"{label} : {mark} {review.get('notes', '')}".rstrip())
        print(f"modules        : {state.get('rtl_modules')}")
        for item in state.get("module_results", []):
            print(
                f"  - {item['module']}: stage={item['stage']} "
                f"success={item['success']} artifacts={item.get('artifacts')}"
            )
        print(f"vivado         : {_jsonable(state.get('vivado_result'))}")
        print(f"hil            : {_jsonable(state.get('hil_result'))}")
        print(f"\nevents ({len(state.get('events', []))}):")
        for event in state.get("events", []):
            print("  " + json.dumps(_jsonable(event), ensure_ascii=False))
    print(f"\nreplayed answers: {len(replayed)}")


# --------------------------------------------------------------------------
# core invoke
# --------------------------------------------------------------------------


def _invoke(run: RunDir, api: dict[str, Any], args) -> int:
    handoffs = Handoffs(run.ledger)
    graph = api["build_workflow"]()
    context = api["WorkflowContext"](
        backend=ClaudeSessionBackend(handoffs, api["SubagentResult"]),
        max_attempts=args.max_attempts,
        review_gate=ClaudeSessionReviewGate(handoffs, api["ReviewDecision"]),
    )
    try:
        state = graph.invoke(
            run.workflow_input,
            {"recursion_limit": args.recursion_limit},
            context=context,
        )
    except BaseException as exc:  # noqa: BLE001 - re-raised unless it is a handoff
        handoff = _unwrap(exc)
        if handoff is None:
            traceback.print_exc()
            return 2
        run.set_pending(handoffs.pending)
        _print_pending(handoffs.pending)
        return 10

    run.set_pending([])
    run.set_state(_jsonable(state))
    _print_final(state, handoffs.replayed)
    return 0 if state.get("status") == "completed" else 1


# --------------------------------------------------------------------------
# commands
# --------------------------------------------------------------------------


def cmd_start(args, api) -> int:
    payload = json.loads(Path(args.input).read_text(encoding="utf-8"))
    payload.pop("outcomes", None)
    run = RunDir(_resolve_work_dir(args, payload), args.thread_id)
    run.set_input(payload)
    if not args.keep_ledger:
        run.set_ledger({})
    run.set_pending([])
    print(f"work dir : {run.work_dir}")
    print(f"run dir  : {run.path}")
    print(f"thread   : {run.thread_id}")
    print("input    : " + json.dumps(payload, ensure_ascii=False))
    return _invoke(run, api, args)


def cmd_resume(args, api) -> int:
    run = RunDir(_resolve_work_dir(args), args.thread_id, create=False)
    return _invoke(run, api, args)


def _validate_result(payload: dict[str, Any]) -> dict[str, Any]:
    unknown = set(payload) - set(RESULT_DEFAULTS)
    if unknown:
        raise SystemExit(f"unknown result fields: {sorted(unknown)}")
    scope = payload.get("change_scope", "none")
    if scope not in CHANGE_SCOPES:
        raise SystemExit(f"change_scope must be one of {sorted(CHANGE_SCOPES)}")
    if payload.get("needs_code_change") and scope == "none":
        raise SystemExit(
            "needs_code_change=true requires an explicit change_scope "
            "(simulink_module / non_simulink_rtl / soc / hardware)"
        )
    return payload


def _validate_review(payload: dict[str, Any]) -> dict[str, Any]:
    unknown = set(payload) - set(REVIEW_DEFAULTS)
    if unknown:
        raise SystemExit(f"unknown review fields: {sorted(unknown)}")
    if "approved" not in payload:
        raise SystemExit("review decision must set approved (true/false)")
    if not payload["approved"] and not payload.get("notes"):
        raise SystemExit("驳回必须写明理由：请提供 notes")
    return payload


def cmd_record(args, api) -> int:
    run = RunDir(_resolve_work_dir(args), args.thread_id, create=False)
    pending = run.pending
    if not pending:
        raise SystemExit("nothing pending — run `start` or `resume` first")
    ledger = run.ledger
    by_key = {item["key"]: item for item in pending}

    if args.all_ok:
        reviews = [k for k, v in by_key.items() if v.get("kind") == "review"]
        if reviews:
            raise SystemExit(
                "--all-ok 不能用于人工审核断点，必须由用户决定：" + ", ".join(reviews)
            )
        for key in by_key:
            ledger[key] = dict(RESULT_DEFAULTS)
        run.set_ledger(ledger)
        print(f"recorded default success for: {sorted(by_key)}")
        return 0

    if not args.key:
        raise SystemExit("--key is required (or use --all-ok)")
    if args.key not in by_key:
        raise SystemExit(f"--key {args.key} is not pending; pending = {sorted(by_key)}")
    item = by_key[args.key]
    is_review = item.get("kind") == "review"

    if args.approve or args.reject:
        if not is_review:
            raise SystemExit("--approve/--reject 只用于人工审核断点")
        payload: dict[str, Any] = {
            "approved": bool(args.approve),
            "notes": args.notes or "",
        }
    else:
        if args.result_file:
            raw = Path(args.result_file).read_text(encoding="utf-8")
        elif args.result is not None:
            raw = args.result
        else:
            raise SystemExit(
                "provide --approve/--reject (review) or --result '<json>' (subagent)"
            )
        payload = json.loads(raw)
        if not isinstance(payload, dict):
            raise SystemExit("result must be a JSON object")

    payload = _validate_review(payload) if is_review else _validate_result(payload)
    ledger[args.key] = payload
    run.set_ledger(ledger)
    remaining = sorted(set(by_key) - set(ledger))
    print(f"recorded {args.key}: " + json.dumps(payload, ensure_ascii=False))
    print(f"still pending: {remaining if remaining else '(none — run `resume`)'}")
    return 0


def cmd_status(args, api) -> int:
    run = RunDir(_resolve_work_dir(args), args.thread_id, create=False)
    print(f"work dir: {run.work_dir}")
    print(f"run dir : {run.path}")
    if not run.path.exists():
        print("(no such run — check --work-dir / --thread-id)")
        return 1
    ledger = run.ledger
    print(f"ledger  : {len(ledger)} recorded answer(s)")
    for key, value in ledger.items():
        print("  " + key + " -> " + json.dumps(value, ensure_ascii=False))
    pending = run.pending
    if pending:
        _print_pending(pending)
    else:
        state = run.state
        print("pending : none")
        if state:
            print(f"status  : {state.get('status')}")
    return 0


def cmd_dry_run(args, api) -> int:
    payload = json.loads(Path(args.input).read_text(encoding="utf-8"))
    outcomes = payload.pop("outcomes", None)
    reviews = payload.pop("reviews", None)
    backend = api["ScriptedBackend"](outcomes) if outcomes else api["ScriptedBackend"]()
    gate = (
        api["ScriptedReviewGate"](reviews) if reviews else api["ScriptedReviewGate"]()
    )
    graph = api["build_workflow"]()
    state = graph.invoke(
        payload,
        {"recursion_limit": args.recursion_limit},
        context=api["WorkflowContext"](
            backend=backend, max_attempts=args.max_attempts, review_gate=gate
        ),
    )
    print("scripted subagent calls:")
    for call in backend.calls:
        print(
            f"  {call.task:<22} skill={call.skill or '-':<18} "
            f"write={call.allow_code_changes}"
        )
    print("human review checkpoints (auto-approved in rehearsal):")
    for request in gate.requests:
        print(f"  {request.checkpoint:<22} {request.title}")
    _print_final(state, [])
    return 0 if state.get("status") == "completed" else 1


def cmd_mermaid(args, api) -> int:
    print(api["build_workflow"]().get_graph().draw_mermaid())
    return 0


def cmd_schema(args, api) -> int:
    print("subagent result:")
    print(json.dumps(RESULT_DEFAULTS, ensure_ascii=False, indent=2))
    print("\nhuman review decision:")
    print(json.dumps(REVIEW_DEFAULTS, ensure_ascii=False, indent=2))
    return 0


def main(argv: list[str] | None = None) -> int:
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:  # noqa: BLE001 - older streams
        pass

    parser = argparse.ArgumentParser(prog="run_sim_rtl_flow")
    parser.add_argument("--project", type=Path, default=None)
    parser.add_argument(
        "--work-dir",
        default=None,
        help=(
            "the design project whose .sim_rtl_flow/ holds run state; "
            "defaults to cwd (on `start` it comes from the input's work_dir)"
        ),
    )
    parser.add_argument("--thread-id", default="default")
    parser.add_argument("--recursion-limit", type=int, default=200)
    parser.add_argument("--max-attempts", type=int, default=10)
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("start")
    p.add_argument("--input", required=True)
    p.add_argument("--keep-ledger", action="store_true")
    p.set_defaults(func=cmd_start)

    p = sub.add_parser("resume")
    p.set_defaults(func=cmd_resume)

    p = sub.add_parser("record")
    p.add_argument("--key")
    p.add_argument("--result")
    p.add_argument("--result-file")
    p.add_argument("--approve", action="store_true")
    p.add_argument("--reject", action="store_true")
    p.add_argument("--notes")
    p.add_argument("--all-ok", action="store_true")
    p.set_defaults(func=cmd_record)

    p = sub.add_parser("status")
    p.set_defaults(func=cmd_status)

    p = sub.add_parser("dry-run")
    p.add_argument("--input", required=True)
    p.set_defaults(func=cmd_dry_run)

    p = sub.add_parser("mermaid")
    p.set_defaults(func=cmd_mermaid)

    p = sub.add_parser("schema")
    p.set_defaults(func=cmd_schema)

    args = parser.parse_args(argv)
    if getattr(args, "approve", False) and getattr(args, "reject", False):
        raise SystemExit("--approve 与 --reject 不能同时使用")
    if args.project is None:
        args.project = _default_project()
    api = _bootstrap(args.project)
    return args.func(args, api)


if __name__ == "__main__":
    raise SystemExit(main())