"""ProteinIQ cookbook, Python SDK 0.2.0+. MIT licensed example code.

Quote by default. --submit and --max-credits explicitly authorize one fixed-price
job. Keep the output directory to resume; use a new directory for changed work.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import tempfile
import uuid

from proteiniq import ProteinIQ
from proteiniq.errors import ProteinIQError


def save_json(path: Path, value: object) -> None:
    with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        temporary = Path(handle.name)
        json.dump(value, handle, indent=2)
        handle.write("\n")
    try:
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def structure_input(file: Path, slot: str) -> dict:
    format_name = file.suffix.lower().lstrip(".")
    if format_name not in {"pdb", "ent", "cif", "mmcif"}:
        raise ValueError("Use a PDB or mmCIF structure file.")
    return {"id": slot, "slotId": slot, "kind": "protein", "format": format_name,
            "content": file.read_text(encoding="utf-8"), "label": file.name,
            "source": {"type": "text"}}


def make_request(tool: str, structure: Path, reference: Path | None) -> dict:
    if tool == "molprobity":
        if reference is not None:
            raise ValueError("MolProbity takes one structure.")
        inputs = [structure_input(structure, "protein")]
        settings = {}
    else:
        if reference is None:
            raise ValueError("USAlign needs --reference as well as --structure.")
        inputs = [structure_input(structure, "structure1"), structure_input(reference, "structure2")]
        if any(item["format"] == "mmcif" for item in inputs):
            raise ValueError("For USAlign, give mmCIF files the .cif extension.")
        settings = {"alignment_mode": "monomer", "ter_mode": "2", "mol_type": "auto",
                    "het_mode": "0", "fast_mode": False}
    return {"tool": tool, "input": {"inputs": inputs}, "settings": settings}


def check_quote(quote, ceiling: int | None) -> None:
    if quote.blocking_errors:
        raise ValueError(f"Quote rejected: {quote.blocking_errors}")
    if quote.billing.get("mode") != "fixed":
        raise ValueError("This example only submits fixed-price jobs; inspect metered billing separately.")
    if ceiling is not None and quote.estimated_credits > ceiling:
        raise ValueError(f"Quote {quote.estimated_credits} exceeds your {ceiling}-credit limit.")


def collect(client: ProteinIQ, state: dict, output: Path, timeout: float) -> None:
    job = client.jobs.wait(state["job_id"], timeout=timeout)
    if job.status != "COMPLETED":
        raise RuntimeError(f"Job {job.id}: {job.status}: {job.error}")
    result = client.results.wait(job.id, timeout=120)
    save_json(output / "results.json", result.results)
    manifest_path = output / "files.json"
    manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {}
    for index, file in enumerate(result.files, 1):
        if not file.url:
            raise RuntimeError(f"No download URL for output {file.name}; rerun to refresh links.")
        name = re.sub(r"[^a-zA-Z0-9._-]", "_", file.name or "result-file")[:160]
        destination = output / f"{index:02d}-{name}"
        if destination.exists():
            digest = hashlib.sha256(destination.read_bytes()).hexdigest()
            if manifest.get(destination.name, {}).get("sha256") != digest:
                raise RuntimeError(f"Existing file has no matching checksum: {destination}")
        else:
            client.download_url(file.url, destination)
            digest = hashlib.sha256(destination.read_bytes()).hexdigest()
            manifest[destination.name] = {"source_name": file.name, "sha256": digest}
            save_json(manifest_path, manifest)
        print(destination)
    print(f"Completed job: {job.id}. Structured result: {output / 'results.json'}")


def run(args: argparse.Namespace, client: ProteinIQ) -> None:
    output = args.output
    state_path = output / "request.json"
    workspace_id = client.account.get().workspace.get("id")
    if not workspace_id:
        raise ValueError("Account response did not identify the workspace.")
    if state_path.exists():
        if not args.resume:
            raise ValueError("This output directory already contains a request. Use --resume or a new directory.")
        state = json.loads(state_path.read_text(encoding="utf-8"))
        if state["base_url"] != client.base_url or state["workspace_id"] != workspace_id:
            raise ValueError("Resume with the same ProteinIQ server and workspace as the original request.")
        if state.get("job_id"):
            collect(client, state, output, args.timeout)
            return
        if not args.submit:
            raise ValueError("Retrying an unconfirmed submission also requires --submit.")
        request, ceiling = state["request"], state["max_credits"]
    else:
        if args.resume:
            raise ValueError("No saved request to resume.")
        if args.tool is None or args.structure is None:
            raise ValueError("Supply --tool and --structure for a new analysis.")
        request = make_request(args.tool, args.structure, args.reference)
        ceiling = args.max_credits
        state = {"request": request, "max_credits": ceiling, "idempotency_key": str(uuid.uuid4()),
                 "base_url": client.base_url, "workspace_id": workspace_id}
    quote = client.jobs.quote(**request)
    print(json.dumps(quote.raw, indent=2))
    check_quote(quote, ceiling)
    if not args.submit:
        print("Quote only. No scientific job submitted.")
        return
    if ceiling is None or ceiling < 1:
        raise ValueError("Submission requires a positive integer --max-credits.")
    # Persist the exact request and key BEFORE a network call can start a job.
    output.mkdir(parents=True, exist_ok=True)
    save_json(state_path, state)
    job = client.jobs.submit(**request, name=f"Cookbook: {request['tool']}",
                             idempotency_key=state["idempotency_key"], max_quoted_credits=ceiling)
    state["job_id"] = job.id
    save_json(state_path, state)
    print(f"Job: {job.id}. Resume with --output {output} --resume")
    collect(client, state, output, args.timeout)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tool", choices=["molprobity", "usalign"])
    parser.add_argument("--structure", type=Path)
    parser.add_argument("--reference", type=Path)
    parser.add_argument("--output", type=Path, default=Path("outputs"))
    parser.add_argument("--submit", action="store_true")
    parser.add_argument("--max-credits", type=int)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--timeout", type=float, default=3600)
    args = parser.parse_args()
    # An exclusive lock also prevents two local processes overwriting recovery state.
    args.output.mkdir(parents=True, exist_ok=True)
    lock = args.output / ".cookbook.lock"
    descriptor = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
    try:
        with ProteinIQ() as client:
            run(args, client)
    except (ProteinIQError, ValueError, RuntimeError, TimeoutError) as error:
        parser.exit(1, f"{error}\nKeep this directory and resume the same request. A timeout does not cancel compute.\n")
    finally:
        os.close(descriptor)
        lock.unlink(missing_ok=True)


if __name__ == "__main__":
    main()