"""run_workflow.py — submit workflow runs to the local langgraph dev server.

Instead of running the graph directly in-process, every invocation sends a
run to the langgraph dev server (http://127.0.0.1:2024) via the LangGraph
SDK sync client.  This lets LangGraph Studio track execution in real time.

The Codex review command is forwarded in the run's config so the server does
not need to be restarted when the command changes between projects.
"""
from __future__ import annotations

import argparse
import base64
import json
import logging
import sys
import time
from pathlib import Path

SERVER_URL = "http://127.0.0.1:2024"
ASSISTANT_ID = "plan_review"
THREAD_MAP_FILENAME = "workflow_threads.json"


# ── path helpers ──────────────────────────────────────────────────────────────

def _resolve_work_dir(value: str) -> Path:
    return Path(value or ".").resolve()


def _resolve_notes_dir(work_dir: Path, value: str) -> Path:
    notes_dir = Path(value or "Notes")
    if not notes_dir.is_absolute():
        notes_dir = work_dir / notes_dir
    return notes_dir.resolve()


# ── thread map (task_id → server thread_id) ──────────────────────────────────

def _load_thread_map(notes_dir: Path) -> dict:
    path = notes_dir / THREAD_MAP_FILENAME
    if path.exists():
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            return {}
    return {}


def _save_thread_map(notes_dir: Path, mapping: dict) -> None:
    path = notes_dir / THREAD_MAP_FILENAME
    path.write_text(json.dumps(mapping, indent=2, ensure_ascii=False), encoding="utf-8")


# ── interrupt extraction ──────────────────────────────────────────────────────

def _extract_interrupt(state: dict):
    """Return the first interrupt value from a thread-state dict, or None."""
    # Primary: tasks[].interrupts (langgraph-api >= 0.5)
    for task in state.get("tasks", []):
        for intr in task.get("interrupts", []):
            if isinstance(intr, dict):
                return intr.get("value", intr)
            return intr
    # Fallback: __interrupt__ channel value
    interrupt_list = state.get("values", {}).get("__interrupt__", [])
    if interrupt_list:
        intr = interrupt_list[0]
        return intr.value if hasattr(intr, "value") else intr
    return None


def _is_interrupted(state: dict) -> bool:
    """Return True when the thread is paused at an interrupt."""
    return bool(state.get("next")) and _extract_interrupt(state) is not None


# ── core workflow runner ──────────────────────────────────────────────────────

def run_workflow(
    *,
    task_id: str,
    user_prompt: str = "",
    notes_dir: str = "Notes",
    work_dir: str = "",
    codex_review_command: list[str],
    thread_id: str | None = None,
    resume_payload=None,
    timeout: int = 600,
    server_url: str = SERVER_URL,
) -> dict:
    try:
        from langgraph_sdk import get_sync_client
    except ImportError:
        raise ImportError(
            "langgraph_sdk not installed. Run: "
            "conda run -n langgraph pip install langgraph-sdk"
        )

    import httpx

    effective_task_id = thread_id or task_id
    resolved_work_dir = _resolve_work_dir(work_dir)
    resolved_notes_dir = _resolve_notes_dir(resolved_work_dir, notes_dir)
    resolved_notes_dir.mkdir(parents=True, exist_ok=True)

    # ── verify server is up ───────────────────────────────────────────────────
    try:
        httpx.get(f"{server_url}/ok", timeout=5).raise_for_status()
    except Exception as exc:
        manage_script = Path(__file__).resolve().with_name("manage_dev_server.ps1")
        raise RuntimeError(
            f"langgraph dev server not reachable at {server_url}.\n"
            f"Start it first:\n"
            f"  PowerShell -File {manage_script} start\n"
            f"Error: {exc}"
        )

    client = get_sync_client(url=server_url)

    # ── resolve server thread_id ──────────────────────────────────────────────
    thread_map = _load_thread_map(resolved_notes_dir)
    server_thread_id: str | None = thread_map.get(effective_task_id)

    # Verify the stored thread still exists on the server (server may have restarted)
    if server_thread_id:
        try:
            client.threads.get(server_thread_id)
        except Exception:
            logging.warning(
                "Stored thread %s no longer exists on server (server restarted?); "
                "creating a new thread.",
                server_thread_id,
            )
            server_thread_id = None

    if not server_thread_id:
        thread = client.threads.create(
            metadata={"task_id": effective_task_id}
        )
        server_thread_id = thread["thread_id"]
        thread_map[effective_task_id] = server_thread_id
        _save_thread_map(resolved_notes_dir, thread_map)

    # ── submit run ────────────────────────────────────────────────────────────
    run_config = {
        "configurable": {
            "codex_review_command": codex_review_command,
            "codex_timeout": timeout,
        }
    }

    if resume_payload is not None:
        run = client.runs.create(
            server_thread_id,
            ASSISTANT_ID,
            command={"resume": resume_payload},
            config=run_config,
        )
    else:
        run = client.runs.create(
            server_thread_id,
            ASSISTANT_ID,
            input={
                "task_id": effective_task_id,
                "user_prompt": user_prompt,
                "notes_dir": str(resolved_notes_dir),
                "work_dir": str(resolved_work_dir),
            },
            config=run_config,
        )

    run_id = run["run_id"]
    logging.info("Submitted run %s on thread %s", run_id, server_thread_id)

    # ── wait for completion or interrupt ──────────────────────────────────────
    ACTIVE = {"pending", "running"}
    deadline = time.monotonic() + timeout
    run_status = run.get("status", "pending")

    while run_status in ACTIVE and time.monotonic() < deadline:
        time.sleep(1)
        try:
            final_run = client.runs.get(server_thread_id, run_id)
            run_status = final_run.get("status", run_status)
        except Exception as exc:
            logging.debug("runs.get error (will retry): %s", exc)

    if run_status in ACTIVE:
        raise TimeoutError(
            f"Run {run_id} did not complete within {timeout}s (status={run_status})"
        )

    # ── detect interrupt from thread state ────────────────────────────────────
    state = client.threads.get_state(server_thread_id)

    if _is_interrupted(state):
        interrupt_value = _extract_interrupt(state)
        return {
            "status": "waiting",
            "thread_id": effective_task_id,
            "server_thread_id": server_thread_id,
            "action_required": interrupt_value,
        }

    if run_status not in ("success",):
        error = final_run.get("error") or run_status
        raise RuntimeError(f"Run ended with status '{run_status}': {error}")

    # ── extract final state ───────────────────────────────────────────────────
    values = state.get("values", {})

    return {
        "status": "done",
        "thread_id": effective_task_id,
        "server_thread_id": server_thread_id,
        "task_path": values.get("task_path", ""),
        "latest_discuss_path": values.get("latest_discuss_path", ""),
        "work_dir": values.get("work_dir", ""),
        "codex_session_id": values.get("codex_session_id", ""),
        "has_serious_issue": values.get("has_serious_issue", False),
        "should_exit": values.get("should_exit", True),
    }


# ── CLI ───────────────────────────────────────────────────────────────────────

def _parse_command(value: str | None) -> list[str] | None:
    if not value:
        return None
    try:
        parsed = json.loads(value)
    except json.JSONDecodeError:
        import re
        try:
            parsed = json.loads(re.sub(r'\\(?![\\"/bfnrtu])', r'\\\\', value))
        except json.JSONDecodeError:
            stripped = value.strip()
            if stripped.startswith("[") and stripped.endswith("]"):
                parsed = [i.strip().strip('"').strip("'") for i in stripped[1:-1].split(",")]
            else:
                raise
    if not isinstance(parsed, list) or not all(isinstance(i, str) for i in parsed):
        raise argparse.ArgumentTypeError("command must be a JSON string array")
    return parsed


def _parse_b64(value: str | None) -> list[str] | None:
    if not value:
        return None
    return _parse_command(base64.b64decode(value).decode("utf-8"))


def _parse_json_b64(value: str | None):
    if not value:
        return None
    return json.loads(base64.b64decode(value).decode("utf-8"))


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Submit a plan-review workflow run to the local langgraph dev server."
    )
    parser.add_argument("--task-id", required=True)
    parser.add_argument("--prompt", default="")
    parser.add_argument("--notes-dir", default="Notes")
    parser.add_argument("--work-dir", default="")
    parser.add_argument("--codex-review-command", type=_parse_command, default=None)
    parser.add_argument("--codex-review-command-b64", type=_parse_b64, default=None)
    parser.add_argument("--timeout", type=int, default=600)
    parser.add_argument("--thread-id", default="")
    parser.add_argument("--resume", default="")
    parser.add_argument("--resume-b64", default="")
    parser.add_argument("--server-url", default=SERVER_URL)
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(
            level=logging.DEBUG, stream=sys.stderr,
            format="%(asctime)s %(levelname)s %(message)s",
        )

    codex_cmd = args.codex_review_command_b64 or args.codex_review_command
    if not codex_cmd:
        parser.error("--codex-review-command or --codex-review-command-b64 is required")

    resume_payload = (
        _parse_json_b64(args.resume_b64) if args.resume_b64
        else (json.loads(args.resume) if args.resume else None)
    )

    output = run_workflow(
        task_id=args.task_id,
        user_prompt=args.prompt,
        notes_dir=args.notes_dir,
        work_dir=args.work_dir,
        codex_review_command=codex_cmd,
        thread_id=args.thread_id or None,
        resume_payload=resume_payload,
        timeout=args.timeout,
        server_url=args.server_url,
    )

    print(json.dumps(output, ensure_ascii=False, indent=2))
    return 0


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