"""Standalone GreenFold API example; requires only requests.

    python -m pip install requests
    python greenfold_api_example.py --uniprot P63313 --model af3
    python greenfold_api_example.py --uniprot P69905 --uniprot P68871 --model boltz

Also runs as python -m api.tools.msa_api_client_example.
Public examples require no API key, GPU or prediction software.
"""

from __future__ import annotations

import argparse
import json
import time
import zipfile
from pathlib import Path
from typing import Iterable
from urllib.parse import quote

import requests


class MsaApiClient:
    def __init__(self, base_url: str, timeout_s: float = 120.0):
        self.base_url = base_url.rstrip("/")
        self.timeout_s = timeout_s

    def request(self, method: str, path: str, **kwargs):
        headers = kwargs.pop("headers", {})
        timeout = kwargs.pop("timeout", self.timeout_s)
        # Bounded retries make rate limits recoverable without indefinite waits.
        for attempt in range(3):
            response = requests.request(
                method, f"{self.base_url}{path}", headers=headers,
                timeout=timeout, **kwargs,
            )
            if response.status_code != 429 or attempt == 2:
                response.raise_for_status()
                return response
            try:
                delay = min(10.0, max(1.0, float(response.headers.get("Retry-After", "2"))))
            except ValueError:
                delay = 2.0
            response.close()
            print(f"Rate limited; retrying in {delay:g}s.", flush=True)
            time.sleep(delay)
        raise RuntimeError("Request retries exhausted.")

    def get_json(self, path: str, params=None):
        with self.request("GET", path, params=params) as response:
            return response.json()

    def download_file(self, path: str, out_path: str | Path, params=None, headers=None) -> Path:
        out_path = Path(out_path)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        partial = out_path.with_name(out_path.name + ".part")
        try:
            with self.request("GET", path, params=params, headers=headers or {}, stream=True) as response:
                with partial.open("wb") as handle:
                    for chunk in response.iter_content(chunk_size=1024 * 1024):
                        if chunk:
                            handle.write(chunk)
            partial.replace(out_path)
        finally:
            partial.unlink(missing_ok=True)
        return out_path.resolve()

    def health(self):
        return self.get_json("/v1/health")

    def resolve(self, uniprot: str):
        return self.get_json("/v1/resolve", params={"uniprot": uniprot})

    def list_msas(self, uniprot: str):
        return self.get_json(f"/v1/msas/{quote(uniprot, safe='')}")

    def download_af3_bundle(self, uniprot_ids: Iterable[str], out_path: str | Path) -> Path:
        params = [("uniprot_ids", uniprot_id) for uniprot_id in uniprot_ids]
        return self.download_file("/v1/download_af3_bundle", out_path, params=params)


def parse_args(argv=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Call GreenFold services and parse their results.")
    parser.add_argument("--base-url", default="https://greenfold.dsdd.one")
    parser.add_argument("--uniprot", action="append", help="Repeat for multiple chains; default: P63313.")
    parser.add_argument("--model", choices=["af2", "af3", "boltz", "openfold", "protenix"], default="af3")
    parser.add_argument("--output-dir", type=Path, default=Path("greenfold_example"))
    parser.add_argument("--timeout", type=float, default=120.0, help="HTTP timeout in seconds.")
    parser.add_argument("--download-af3-bundle", metavar="PATH", help="Also download a direct AF3 bundle.")
    parser.add_argument("--sto-job-db", help="Optionally submit, poll and download an STO job, e.g. mgnify.")
    parser.add_argument("--job-timeout", type=float, default=120.0, help="Maximum job polling time.")
    args = parser.parse_args(argv)
    # Explicit targets replace the default, rather than being appended after it.
    args.uniprot = args.uniprot or ["P63313"]
    if args.timeout <= 0 or args.job_timeout <= 0:
        parser.error("--timeout and --job-timeout must be positive.")
    return args


def main() -> None:
    args = parse_args()
    client = MsaApiClient(args.base_url, timeout_s=args.timeout)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    # 1. Resolve requested identifiers and retrieve full query sequences.
    health = client.health()
    if not health.get("ok"):
        raise RuntimeError(f"Unexpected health response: {health}")
    print("Service ready.", flush=True)
    chains = []
    for chain_index, identifier in enumerate(args.uniprot):
        if chain_index >= 26:
            raise ValueError("This template supports 26 chains; adapt IDs for larger inputs.")
        reference_id = identifier.split("[", 1)[0]
        resolved = client.resolve(reference_id)
        info = client.get_json("/v1/sequence_info", params={"uniprot": identifier})
        sequence = info["sequence"]
        if info.get("sequence_truncated") or len(sequence) != info["length"]:
            raise RuntimeError(f"Incomplete sequence for {identifier}: expected {info['length']} aa.")
        chain_id = chr(ord("A") + chain_index)
        chains.append({"id": chain_id, "sequence": sequence, "uniprot": resolved["uniprot"]})
        print(f"Chain {chain_id}: {resolved['uniprot']}, {len(sequence)} aa", flush=True)
    first = chains[0]
    accession = first["uniprot"]
    encoded = quote(accession, safe="")

    # 2. These metadata counts describe canonical evidence, before transforms.
    metadata = client.get_json(f"/v1/msa_metadata/{encoded}")
    (args.output_dir / "msa_metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
    for kind in ("paired", "unpaired"):
        item = metadata.get("greenfold_msa", {}).get(kind) or {}
        print(f"{kind}: {item.get('n_sequences')} records, depth {item.get('depth')} (excludes query)")
    print("MSA snapshot:", metadata.get("greenfold_msa", {}).get("source_snapshot"))

    # 3. Submit a sequence and parse candidate fields.
    with client.request("POST", "/v1/reverse_sequence_lookup",
                        json={"sequence": first["sequence"], "max_candidates": 2}) as response:
        lookup = response.json()
    (args.output_dir / "reverse_lookup.json").write_text(json.dumps(lookup, indent=2) + "\n")
    for candidate in lookup.get("candidates", []):
        print("Candidate:", candidate["id"], "identity=", candidate["identity"],
              "variation=", candidate.get("modification_pattern"), "ambiguous=", candidate.get("ambiguous"))
    if not lookup.get("candidates"):
        raise RuntimeError(f"No reference candidate for {args.uniprot[0]}.")

    # 4. Download A3M and check the query, preserving a requested variation.
    pattern = args.uniprot[0][args.uniprot[0].find("["):] if "[" in args.uniprot[0] else None
    variation_headers = {"mutation-pattern": pattern} if pattern else {}
    preview = client.get_json(f"/v1/preview_text/{encoded}",
                             params={"format": "a3m", "kind": "unpaired", "max_chars": 2000})
    print("Preview:", preview.get("preview_chars"), "characters; not the full alignment depth.")
    a3m_path = client.download_file(
        f"/v1/download_a3m/{encoded}", args.output_dir / "query_unpaired.a3m",
        params={"kind": "unpaired"}, headers=variation_headers,
    )
    query_lines = []
    record_count = 0
    with a3m_path.open() as handle:
        for line in handle:
            if line.startswith(">"):
                record_count += 1
            elif record_count == 1:
                query_lines.append(line.strip())
    query = "".join(query_lines).replace("-", "").upper()
    if query != first["sequence"]:
        raise RuntimeError(f"A3M query mismatch for {args.uniprot[0]}: {len(query)} vs {len(first['sequence'])} aa.")
    print(f"Downloaded {record_count} A3M records; query verified.", flush=True)

    # 5. Native formats are explicit. Replace this block with your own raw file
    # contents to include ligands, templates or other model-specific settings.
    name = "greenfold_api_example"
    if args.model == "af2":
        filename = "target.fasta"
        native_text = "".join(f">{c['id']}\n{c['sequence']}\n" for c in chains)
    elif args.model == "af3":
        filename = "input_af3.json"
        native_text = json.dumps({
            "name": name, "modelSeeds": [1], "dialect": "alphafold3", "version": 2,
            "sequences": [{"protein": {"id": c["id"], "sequence": c["sequence"]}} for c in chains],
        }, indent=2) + "\n"
    elif args.model == "boltz":
        filename = "input_boltz.yaml"
        native_text = "version: 1\nsequences:\n" + "".join(
            f"  - protein:\n      id: {c['id']}\n      sequence: {c['sequence']}\n" for c in chains
        )
    elif args.model == "openfold":
        filename = "query_openfold.json"
        native_text = json.dumps({"queries": {name: {"chains": [
            {"molecule_type": "protein", "chain_ids": [c["id"]], "sequence": c["sequence"]}
            for c in chains
        ]}}}, indent=2) + "\n"
    else:
        filename = "input_protenix.json"
        native_text = json.dumps([{"name": name, "sequences": [
            {"proteinChain": {"id": [c["id"]], "sequence": c["sequence"], "count": 1}} for c in chains
        ]}], indent=2) + "\n"
    (args.output_dir / filename).write_text(native_text)
    params = {"input_filename": filename}
    if args.model in {"boltz", "openfold", "protenix"}:
        params["fallback_msa_strategy"] = "greenfold_colabfold"
    zip_path = args.output_dir / "prepared_input.zip"
    with client.request("POST", f"/v1/prepare_input/{args.model}", params=params,
                        data=native_text.encode(), headers={"Content-Type": "text/plain; charset=utf-8"}) as response:
        zip_path.write_bytes(response.content)

    # 6. Parse the package without executing structure prediction or ZIP code.
    with zipfile.ZipFile(zip_path) as archive:
        report = json.loads(archive.read("greenfold_input_preparation_report.json"))
        manifest = json.loads(archive.read("bundle_manifest.json"))
        if report["model"] != args.model or report["prepared_input"] not in archive.namelist():
            raise RuntimeError(f"Unexpected prepared input contract: {report}")
        (args.output_dir / "preparation_report.json").write_text(json.dumps(report, indent=2) + "\n")
        (args.output_dir / "bundle_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
        print("Prepared input:", report["prepared_input"])
        for chain in report["chains"]:
            print("Prepared chain:", chain.get("label"), chain.get("uniprot"),
                  "identity=", chain.get("identity"), "source=", chain.get("msa_source"),
                  "pairing=", chain.get("pairing_mode"), "native=", chain.get("fallback_native_msa"))
        for warning in report.get("warnings", []):
            print("WARNING:", warning)
        print("Generated files:", len(manifest.get("generated_files", [])))
    print("Package:", zip_path.resolve())
    print("Next: unpack it and follow RUN_COMMANDS.md in the chosen model environment.", flush=True)

    if args.download_af3_bundle:
        print("Direct AF3 bundle:", client.download_af3_bundle(args.uniprot, args.download_af3_bundle))

    # Optional asynchronous service family: submit, poll with a deadline, download.
    if args.sto_job_db:
        with client.request("POST", "/v1/jobs/sto",
                            params={"uniprot": accession, "db_name": args.sto_job_db}) as response:
            job = response.json()
        job_id = quote(job["job_id"], safe="")
        deadline = time.monotonic() + args.job_timeout
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError(f"STO job {job_id} exceeded {args.job_timeout:g}s; poll its status later.")
            with client.request("GET", f"/v1/jobs/{job_id}",
                                timeout=min(args.timeout, remaining)) as response:
                status = response.json()
            print("STO job:", status["status"], flush=True)
            if status["status"] == "done":
                break
            if status["status"] == "failed":
                raise RuntimeError(f"STO job failed: {status}")
            if time.monotonic() >= deadline:
                raise TimeoutError(f"STO job {job_id} exceeded {args.job_timeout:g}s; poll its status later.")
            time.sleep(min(2, max(0, deadline - time.monotonic())))
        print("STO artifact:", client.download_file(
            f"/v1/jobs/{job_id}/download", args.output_dir / "source.sto.zst",
        ))


if __name__ == "__main__":
    main()