#!/usr/bin/env python3
"""
fpga-goldenref-flow workflow runner — IN-PROCESS driver (NO langgraph dev server).

Drives the compiled LangGraph graph directly in this Python process with a
persistent SQLite checkpointer. This avoids the langgraph dev server entirely:
no REST, no in-memory store, no background run-queue / enqueue jams, no
watchfiles reload loop. State persists across start/resume/continue invocations
in <output_dir>/.fpga_flow_local.sqlite, keyed by thread_id.

This is the sole driver for the graph (the langgraph dev server path has been
removed). CLI + exit codes:

  start     submit a new run; exits 2 at an interrupt
  resume    answer an interrupt (--plan-model OR approve/--rejected --feedback)
  continue  continue a thread paused at a normal node boundary (input=None)

Exit codes: 0 complete · 1 error · 2 interrupted (human action needed)
"""
import os
import sys
import json
import uuid
import argparse
import sqlite3

# Drive the graph in-process: import the app package.
FPGA_FLOW_ROOT = r"E:\LangGraph_Design\fpga-goldenref-flow"
if FPGA_FLOW_ROOT not in sys.path:
    sys.path.insert(0, FPGA_FLOW_ROOT)

RECURSION_LIMIT = int(os.getenv("FPGA_FLOW_RECURSION_LIMIT", "300"))


def _session_file(output_dir: str = "") -> str:
    base = (output_dir or "").strip() or FPGA_FLOW_ROOT
    return os.path.join(base, ".fpga_flow_session.json")


def _db_path(output_dir: str) -> str:
    base = (output_dir or "").strip() or FPGA_FLOW_ROOT
    return os.path.join(base, ".fpga_flow_local.sqlite")


def save_session(thread_id, output_dir):
    with open(_session_file(output_dir), "w", encoding="utf-8") as f:
        json.dump({"thread_id": thread_id, "output_dir": output_dir, "driver": "local"}, f)


def load_session(output_dir=""):
    path = _session_file(output_dir)
    if not os.path.exists(path):
        print(f"[FAIL] No session file at {path}. Run 'start' first (or pass --thread-id).", flush=True)
        sys.exit(1)
    with open(path, encoding="utf-8") as f:
        return json.load(f)


# ── graph + checkpointer ──────────────────────────────────────────────────────

def _build(output_dir):
    """Compile the graph with a persistent SQLite checkpointer (keeps the conn)."""
    from langgraph.checkpoint.sqlite import SqliteSaver
    from app.main_graph import build_graph
    conn = sqlite3.connect(_db_path(output_dir), check_same_thread=False)
    saver = SqliteSaver(conn)
    try:
        saver.setup()  # create checkpoint tables if not present
    except Exception:
        pass
    return build_graph(checkpointer=saver), conn


def _config(thread_id):
    return {"configurable": {"thread_id": thread_id}, "recursion_limit": RECURSION_LIMIT}


def _stream(graph, config, payload):
    """Run the graph to its next pause, printing node-level progress."""
    for chunk in graph.stream(payload, config, stream_mode="updates"):
        if isinstance(chunk, dict):
            for node in chunk:
                if node and not str(node).startswith("__"):
                    print(f"  -> {node}", flush=True)


def _pending_interrupts(graph, config):
    """Return list of interrupt payload values currently pending (if any)."""
    snap = graph.get_state(config)
    vals = []
    for task in (snap.tasks or ()):
        for it in (getattr(task, "interrupts", None) or ()):
            v = getattr(it, "value", it)
            vals.append(v)
    return vals, snap


def _print_interrupt(graph, config):
    vals, snap = _pending_interrupts(graph, config)
    for v in vals:
        print(f"\n[INTERRUPT] {json.dumps(v, indent=2, default=str, ensure_ascii=False)}", flush=True)
    # plan/requirement paths from state
    values = snap.values or {}
    for serial in values.get("serial_tasks", []):
        for pt in serial.get("parallel_tasks", []):
            tid = pt.get("task_id", "?")
            if pt.get("plan_path"):
                print(f"[PLAN] {tid}: {pt['plan_path']}", flush=True)
            if pt.get("req_path"):
                print(f"[REQ]  {tid}: {pt['req_path']}", flush=True)
    print("\n[ACTION REQUIRED] Claude will read these plans and ask for your approval.", flush=True)


def _show_summary(output_dir):
    summary = os.path.join(output_dir, "summary.md")
    if os.path.exists(summary):
        with open(summary, encoding="utf-8") as f:
            print(f"\n[SUMMARY]\n{f.read()}", flush=True)
    else:
        print(f"[DONE] Workflow completed. (summary not found at {summary})", flush=True)


def _finish(graph, config, output_dir):
    """After a stream returns, classify: interrupted (2) / done (0)."""
    vals, snap = _pending_interrupts(graph, config)
    if vals:
        print("\n[STATUS] interrupted", flush=True)
        _print_interrupt(graph, config)
        sys.exit(2)
    if snap.next:
        # paused without an interrupt (rare) — treat as needing continue
        print(f"\n[STATUS] paused (next={list(snap.next)})", flush=True)
        sys.exit(2)
    print("\n[STATUS] success", flush=True)
    _show_summary(output_dir)
    sf = _session_file(output_dir)
    if os.path.exists(sf):
        os.remove(sf)
    sys.exit(0)


# ── commands ──────────────────────────────────────────────────────────────────

def cmd_start(args):
    raw_out = (args.output_dir or "").strip()
    output_dir = os.path.abspath(raw_out) if raw_out else os.path.join(
        os.path.dirname(os.path.abspath(args.input_dir)), "runs")
    os.makedirs(output_dir, exist_ok=True)
    thread_id = str(uuid.uuid4())
    print(f"[INFO] Driver : in-process (no dev server)", flush=True)
    print(f"[INFO] Thread : {thread_id}", flush=True)
    print(f"[INFO] Output dir : {output_dir}", flush=True)
    print(f"[INFO] Checkpoints : {_db_path(output_dir)}", flush=True)

    graph, conn = _build(output_dir)
    try:
        inp = {"input_dir": args.input_dir, "output_dir": output_dir}
        if args.prompt:
            inp["user_prompt"] = args.prompt
        print("[INFO] Running fpga_flow graph ...\n", flush=True)
        save_session(thread_id, output_dir)
        _stream(graph, _config(thread_id), inp)
        _finish(graph, _config(thread_id), output_dir)
    finally:
        conn.close()


def _resolve_thread(args):
    output_dir = (args.output_dir or "").strip()
    thread_id = (getattr(args, "thread_id", "") or "").strip()
    if not thread_id:
        sess = load_session(output_dir)
        thread_id = sess["thread_id"]
        output_dir = sess.get("output_dir", output_dir)
    return thread_id, output_dir


def cmd_resume(args):
    from langgraph.types import Command
    thread_id, output_dir = _resolve_thread(args)
    graph, conn = _build(output_dir)
    try:
        plan_model = (getattr(args, "plan_model", "") or "").strip()
        if plan_model:
            resume_val = {"plan_model": plan_model}
            print(f"[INFO] Resuming {thread_id}  plan_model={plan_model}", flush=True)
        else:
            resume_val = {"approved": not args.rejected, "feedback": args.feedback or ""}
            print(f"[INFO] Resuming {thread_id}  approved={not args.rejected}", flush=True)
        print("", flush=True)
        _stream(graph, _config(thread_id), Command(resume=resume_val))
        _finish(graph, _config(thread_id), output_dir)
    finally:
        conn.close()


def cmd_continue(args):
    thread_id, output_dir = _resolve_thread(args)
    graph, conn = _build(output_dir)
    try:
        print(f"[INFO] Continuing {thread_id} from checkpoint ...\n", flush=True)
        _stream(graph, _config(thread_id), None)
        _finish(graph, _config(thread_id), output_dir)
    finally:
        conn.close()


def main():
    p = argparse.ArgumentParser(description="fpga-goldenref-flow IN-PROCESS runner (no dev server)")
    sub = p.add_subparsers(dest="cmd")

    ps = sub.add_parser("start")
    ps.add_argument("--input-dir", default="Notes/")
    ps.add_argument("--output-dir", default="")
    ps.add_argument("--prompt", default="")

    pr = sub.add_parser("resume")
    pr.add_argument("--plan-model", default="")
    pr.add_argument("--rejected", action="store_true", default=False)
    pr.add_argument("--feedback", default="")
    pr.add_argument("--output-dir", default="")
    pr.add_argument("--thread-id", default="")

    pc = sub.add_parser("continue")
    pc.add_argument("--output-dir", default="")
    pc.add_argument("--thread-id", default="")

    args = p.parse_args()
    if not args.cmd:
        p.print_help(); sys.exit(1)

    try:
        from app.nodes_design import CodexUnavailable
    except Exception:
        CodexUnavailable = ()  # type: ignore

    try:
        if args.cmd == "start":
            cmd_start(args)
        elif args.cmd == "resume":
            cmd_resume(args)
        elif args.cmd == "continue":
            cmd_continue(args)
    except SystemExit:
        raise
    except CodexUnavailable as exc:
        out = getattr(args, "output_dir", "") or ""
        bar = "=" * 64
        print(f"\n{bar}", flush=True)
        print("[ALARM] CODEX 不可用 —— 工作流已停止（不降级）", flush=True)
        print(f"原因：{exc}", flush=True)
        print("codex(gpt-5.5) 与你其它工作共享额度。请等 codex 空出来后,从断点续跑：", flush=True)
        print(f'  run_fpga_local.py continue --output-dir "{out}"', flush=True)
        print("状态已保存到 SQLite,不会丢失;不会回退到 Claude 或更弱模型。", flush=True)
        print(bar, flush=True)
        sys.exit(3)


if __name__ == "__main__":
    main()