"""Exercise the REST API with a Diels–Alder sketch and revision conflicts.

Run: python reaction_rest.py http://127.0.0.1:8766 NEW_DIRECTORY
Uses only the standard library. Leaves the final sketch open for inspection.
Does not submit jobs or launch a viewer; see queue_lifecycle.py for submission.
"""

from __future__ import annotations

import argparse
import base64
import io
import json
import urllib.error
import urllib.request
import zipfile
from pathlib import Path


class Client:
    def __init__(self, url: str):
        self.url = url.rstrip("/")
        self.calls: list[dict] = []

    def call(self, method: str, path: str, data=None, *, expected=200, headers=None):
        request = urllib.request.Request(
            self.url + path,
            method=method,
            data=json.dumps(data).encode() if data is not None else None,
            headers={"Content-Type": "application/json", **(headers or {})},
        )
        try:
            response = urllib.request.urlopen(request, timeout=120)
        except urllib.error.HTTPError as error:
            response = error
        with response:
            body = response.read()
            status = response.status
            media = response.headers.get("Content-Type", "")
        self.calls.append({"method": method, "path": path, "status": status})
        if status != expected:
            raise RuntimeError(f"{method} {path}: expected {expected}, got {status}: {body!r}")
        return json.loads(body) if "application/json" in media else body


def run(client: Client, output: Path) -> None:
    for path in (
        "/health",
        "/api/version",
        "/api/session",
        "/api/documents",
        "/api/writers",
        "/api/queue",
    ):
        client.call("GET", path)
    schema = client.call("GET", "/api/schema/document")
    (output / "schema.json").write_text(json.dumps(schema, indent=2))
    created = client.call(
        "POST",
        "/api/documents",
        {"format": "smiles", "data": "C=CC=C.C=C>>C1=CCCCC1"},
        expected=201,
    )
    route = f"/api/documents/{created['id']}"
    state = client.call("GET", route)
    sid = state["document"]["species"][0]["id"]
    old_revision = state["revision"]
    state = client.call(
        "PATCH",
        f"{route}/species/{sid}",
        {"revision": old_revision, "label": "1", "name": "1,3-butadiene"},
    )
    # A stale change must fail. Fetch, review, then make a fresh edit.
    client.call("POST", route + "/clean", {"revision": old_revision}, expected=409)
    client.call(
        "PUT", route, state["document"], expected=412, headers={"If-Match": str(old_revision)}
    )
    state = client.call("GET", route)
    state = client.call(
        "PUT", route, state["document"], headers={"If-Match": str(state["revision"])}
    )
    state = client.call(
        "PUT",
        route + "/canvas",
        {"revision": state["revision"], "ket": state["document"]["canvas"]["data"]},
    )
    # Append and replace happen on a scratch document, keeping the balanced scheme intact.
    scratch = client.call("POST", "/api/documents", {}, expected=201)
    scratch_path = f"/api/documents/{scratch['id']}"
    for mode, smiles in (("append", "CCO"), ("append", "O"), ("replace", "[CH3]")):
        scratch = client.call(
            "POST",
            scratch_path + "/import",
            {"revision": scratch["revision"], "format": "smiles", "data": smiles, "mode": mode},
        )
    assert scratch["document"]["species"][0]["multiplicity"] == 2
    client.call("DELETE", scratch_path, expected=204)
    state = client.call("POST", route + "/clean", {"revision": state["revision"]})
    state = client.call("POST", route + "/embed", {"revision": state["revision"], "count": 2})
    report = client.call("GET", route + "/report?names=false")
    assert [row["formula"] for row in report] == ["C4H6", "C2H4", "C6H10"]
    spec = {"task": "sp", "method": "B3LYP", "basis": "def2-SVP"}
    inputs = client.call("POST", route + "/inputs", {"program": "vibeqc", "spec": spec})
    assert len(inputs["files"]) == 3 and not inputs["refused"]
    archive = client.call(
        "POST", route + "/inputs", {"program": "vibeqc", "spec": spec, "format": "zip"}
    )
    with zipfile.ZipFile(io.BytesIO(archive)) as files:
        assert len([name for name in files.namelist() if name.endswith(".py")]) == 3
    (output / "inputs.zip").write_bytes(archive)
    for fmt in ("vsk", "svg", "ket", "rxn", "sdf"):
        data = client.call("GET", route + f"/export/{fmt}")
        # JSON formats may be decoded by the generic client; re-encode them.
        if isinstance(data, dict):
            data = json.dumps(data, indent=2).encode()
        (output / f"diels-alder.{fmt}").write_bytes(data)
    # A browser-created document has no server-side destination. Export it to the client.
    client.call("POST", route + "/save", expected=409)
    queue = client.call("GET", "/api/queue")
    runs = client.call("GET", route + "/runs", expected=200 if queue["installed"] else 503)
    if queue["installed"]:
        assert runs == []
    doc = state["document"]
    assert client.call("POST", "/api/documents/validate", doc)["valid"]
    assert not client.call("POST", "/api/documents/validate", {"schema": "bad"})["valid"]
    converted = client.call("POST", "/api/convert", {"format": "smiles", "data": "CCO"})
    client.call("POST", "/api/clean", converted["document"])
    client.call("POST", "/api/chem/embed3d", {"smiles": "C[C@H](O)C(=O)O", "count": 2})
    client.call("POST", "/api/report?names=false", doc)
    assert b"<svg" in client.call("POST", "/api/export/svg", doc)
    single = client.call(
        "POST", "/api/inputs", {"document": doc, "species": sid, "program": "vibeqc", "spec": spec}
    )
    assert single["files"]
    session = client.call("GET", "/api/session")
    if session["qvf"]:
        qvf = client.call("GET", route + "/export/qvf")
        encoded = base64.b64encode(qvf).decode("ascii")
        assert client.call("POST", "/api/qvf/overwrite-check", {"data": encoded})["allowed"]
        imported = client.call(
            "POST",
            "/api/documents",
            {"format": "qvf", "encoding": "base64", "data": encoded},
            expected=201,
        )
        assert len(imported["document"]["species"]) == 3
        client.call("DELETE", f"/api/documents/{imported['id']}", expected=204)
        (output / "diels-alder.qvf").write_bytes(qvf)
    (output / "requests.json").write_text(json.dumps(client.calls, indent=2))
    print(f"Inspect the shared document: {client.url}/?doc={created['id']}")
    print("Revision conflicts, validation, binary QVF and three vibe-qc input files checked.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("server")
    parser.add_argument("directory", type=Path)
    args = parser.parse_args()
    args.directory.mkdir(parents=True, exist_ok=False)
    run(Client(args.server), args.directory)