"""Prepare ethanol through a running vibe-sketch REST server.

Run: python ethanol_rest.py http://127.0.0.1:8766 ethanol-rest
Uses Python's standard library; downloads into a new directory.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from urllib.request import Request, urlopen


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("server", help="address printed by vibe-sketch serve")
    parser.add_argument("directory", type=Path, help="new directory for the example outputs")
    args = parser.parse_args()
    args.directory.mkdir(parents=True, exist_ok=False)
    base = args.server.rstrip("/")

    def post(route: str, data: dict) -> dict:
        request = Request(
            base + route,
            data=json.dumps(data).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urlopen(request, timeout=120) as response:
            return json.load(response)

    state = post("/api/documents", {"format": "smiles", "data": "CCO"})
    route = f"/api/documents/{state['id']}"
    state = post(route + "/embed", {"revision": state["revision"]})
    written = post(
        route + "/inputs",
        {
            "program": "vibeqc",
            "spec": {"task": "sp", "method": "B3LYP", "basis": "def2-SVP"},
        },
    )
    if written["refused"]:
        raise SystemExit(f"Writer refused the request: {written['refused']}")
    for file in written["files"]:
        # Keep filenames returned by a server inside the selected directory.
        name = file["name"]
        if Path(name).name != name or name in {".", ".."} or "\\" in name:
            raise ValueError("The server returned an unsafe filename")
        with (args.directory / name).open("x", encoding="utf-8") as stream:
            stream.write(file["content"])
    for extension in ("vsk", "svg"):
        with (
            urlopen(base + route + "/export/" + extension, timeout=30) as response,
            (args.directory / f"ethanol.{extension}").open("xb") as stream,
        ):
            stream.write(response.read())
    print(f"Open the shared drawing at {base}/?doc={state['id']}")
    print("Downloaded the document, SVG and vibe-qc script. No calculation was submitted.")


if __name__ == "__main__":
    main()