"""Run or cancel one real vibe-qc job through REST or MCP.

Download reaction_rest.py and reaction_mcp.py beside this file. Start an
isolated, configured local queue and a vibe-sketch server in that environment.
Run: python queue_lifecycle.py SERVER NEW_DIRECTORY --transport rest --submit
Add --cancel to submit then immediately cancel the new run. --view explicitly
opens the generated structure in vibe-view. No unrelated run is changed.
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

from reaction_mcp import MCP
from reaction_rest import Client


def run(args) -> None:
    output = args.directory
    output.mkdir(parents=True, exist_ok=False)
    rest = Client(args.server)
    readiness = rest.call("GET", "/api/queue")
    if not readiness["daemon"] or "vibeqc" not in readiness["programs"]:
        raise SystemExit(f"Queue is not ready: {readiness}")
    mcp = MCP(args.server) if args.transport == "mcp" else None
    try:
        # Hydrogen is small enough for a first complete run. Use cyclohexene
        # for cancellation so completion is unlikely to race the cancellation.
        smiles = "C1=CCCCC1" if args.cancel else "[H][H]"
        if mcp:
            created = mcp.tool("new_document", format="smiles", data=smiles)
            identifier = created["id"]
        else:
            created = rest.call(
                "POST", "/api/documents", {"format": "smiles", "data": smiles}, expected=201
            )
            identifier = created["id"]
        route = f"/api/documents/{identifier}"
        spec = {"task": "sp", "method": "B3LYP", "basis": "def2-SVP", "resources": {"nprocs": 1}}
        if mcp:
            submitted = mcp.tool("submit_jobs", document=identifier, program="vibeqc", spec=spec)
        else:
            submitted = rest.call(
                "POST",
                route + "/submit",
                {"revision": created["revision"], "program": "vibeqc", "spec": spec},
            )
        if submitted["refused"] or len(submitted["submitted"]) != 1:
            raise RuntimeError(submitted)
        run_id = submitted["submitted"][0]["run"]
        print(
            f"Submitted {run_id}; inspect {args.server.rstrip('/')}/?doc={identifier}", flush=True
        )
        if args.cancel:
            if mcp:
                mcp.tool("cancel_run", run=run_id)
            else:
                rest.call("POST", f"/api/queue/runs/{run_id}/cancel")
        history = []
        deadline = time.monotonic() + 600
        while True:
            rows = (
                mcp.tool("show_runs", document=identifier)
                if mcp
                else rest.call("GET", route + "/runs")
            )
            row = next(item for item in rows if item["run"] == run_id)
            if not history or history[-1] != row["state"]:
                history.append(row["state"])
                print(row["state"], flush=True)
            if row["finished"]:
                break
            if time.monotonic() > deadline:
                raise RuntimeError(f"Timed out; inspect run {run_id} with vq before retrying")
            time.sleep(1)
        if args.cancel:
            if row["state"] not in ("killed", "aborted_by_queue"):
                raise RuntimeError(f"Cancellation did not finish as expected: {row}")
        elif row["state"] != "completed" or row["exit_code"] != 0:
            raise RuntimeError(f"Calculation failed: {row}")
        if mcp:
            mcp.tool("save_document", document=identifier, path=str(output / "calculation.vsk"))
            reopened = mcp.tool("open_document", path=str(output / "calculation.vsk"))
            reopened_rows = mcp.tool("show_runs", document=reopened["id"])
        else:
            data = rest.call("GET", route + "/export/vsk")
            (output / "calculation.vsk").write_bytes(data)
            reopened = rest.call(
                "POST", "/api/documents", {"document": json.loads(data)}, expected=201
            )
            reopened_rows = rest.call("GET", f"/api/documents/{reopened['id']}/runs")
        assert any(item["run"] == run_id for item in reopened_rows)
        if args.view:
            state = rest.call("GET", route)
            sid = state["document"]["species"][0]["id"]
            if mcp:
                mcp.tool("show_in_vibe_view", document=identifier, species=sid)
            else:
                rest.call("POST", route + "/view", {"species": sid})
        # Run state includes local paths. Keep execution records local; do not publish them.
        (output / "run.json").write_text(
            json.dumps({"document": identifier, "run": row, "states": history}, indent=2)
        )
        print(f"vq fetch localhost {run_id} --workspace -o results")
    finally:
        if mcp:
            mcp.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("server")
    parser.add_argument("directory", type=Path)
    parser.add_argument("--transport", choices=("rest", "mcp"), default="rest")
    parser.add_argument(
        "--submit",
        action="store_true",
        required=True,
        help="explicitly submit a real job to the configured local queue",
    )
    parser.add_argument("--cancel", action="store_true")
    parser.add_argument("--view", action="store_true")
    run(parser.parse_args())