"""
FastAPI service that wraps Protenix v2 for protein structure prediction.

POST /fold
  Form fields:
    sequence    (required) — single-letter amino acid sequence
    msa_file    (optional) — pre-computed MSA in .a3m format

Returns the predicted structure as a CIF (or PDB) file download.

Start with:
    uvicorn protenix_api:app --host 0.0.0.0 --port 8000

Environment variables:
    PROTENIX_MODEL   — model name passed to `protenix pred -n` (default: protenix-v2)
    PROTENIX_TIMEOUT — max seconds for inference subprocess (default: 600)
"""

import json
import os
import shutil
import subprocess
import tempfile
import uuid
from pathlib import Path

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

app = FastAPI(title="Protenix Folding API", version="1.0.0")

MODEL_NAME = os.environ.get("PROTENIX_MODEL", "protenix-v2")
TIMEOUT = int(os.environ.get("PROTENIX_TIMEOUT", "600"))


@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_NAME}


@app.post("/fold")
async def fold(
    sequence: str = Form(..., description="Amino acid sequence in single-letter codes"),
    msa_file: UploadFile | None = File(None, description="Pre-computed MSA in .a3m format"),
):
    """
    Fold a protein sequence using Protenix v2.

    Optionally supply a pre-computed MSA (.a3m) to improve accuracy.
    Returns the predicted structure as a CIF file.
    """
    sequence = sequence.strip().upper()
    if not sequence:
        raise HTTPException(status_code=422, detail="sequence must not be empty")

    valid_aa = set("ACDEFGHIKLMNPQRSTVWYX")
    invalid = set(sequence) - valid_aa
    if invalid:
        raise HTTPException(
            status_code=422,
            detail=f"Invalid amino acid characters: {sorted(invalid)}",
        )

    job_id = uuid.uuid4().hex[:12]
    work_dir = Path(tempfile.mkdtemp(prefix=f"protenix_{job_id}_"))

    try:
        protein_chain: dict = {"sequence": sequence, "count": 1}

        if msa_file is not None:
            msa_path = work_dir / "input.a3m"
            msa_path.write_bytes(await msa_file.read())
            # Protenix expects absolute paths for MSA files
            protein_chain["unpairedMsaPath"] = str(msa_path.resolve())

        input_data = [
            {
                "name": job_id,
                "sequences": [{"proteinChain": protein_chain}],
                "covalent_bonds": [],
            }
        ]

        input_json = work_dir / "input.json"
        input_json.write_text(json.dumps(input_data, indent=2))

        output_dir = work_dir / "output"
        output_dir.mkdir()

        result = subprocess.run(
            [
                "protenix", "pred",
                "-i", str(input_json),
                "-o", str(output_dir),
                "-n", MODEL_NAME,
            ],
            capture_output=True,
            text=True,
            timeout=TIMEOUT,
        )

        if result.returncode != 0:
            raise HTTPException(
                status_code=500,
                detail=f"Protenix inference failed:\n{result.stderr[-3000:]}",
            )

        # Prefer CIF over PDB; take the first hit if multiple seeds/samples exist
        structure_files = sorted(output_dir.rglob("*.cif")) or sorted(output_dir.rglob("*.pdb"))
        if not structure_files:
            raise HTTPException(
                status_code=500,
                detail="Protenix produced no output structure file",
            )

        best = structure_files[0]
        out_path = work_dir / f"{job_id}{best.suffix}"
        shutil.copy(best, out_path)

        media_type = "chemical/x-mmcif" if best.suffix == ".cif" else "chemical/x-pdb"
        return FileResponse(
            path=str(out_path),
            media_type=media_type,
            filename=f"{job_id}{best.suffix}",
            background=BackgroundTask(shutil.rmtree, work_dir, ignore_errors=True),
        )

    except HTTPException:
        shutil.rmtree(work_dir, ignore_errors=True)
        raise
    except subprocess.TimeoutExpired:
        shutil.rmtree(work_dir, ignore_errors=True)
        raise HTTPException(
            status_code=504,
            detail=f"Inference timed out after {TIMEOUT}s. Try a shorter sequence or increase PROTENIX_TIMEOUT.",
        )
    except Exception as exc:
        shutil.rmtree(work_dir, ignore_errors=True)
        raise HTTPException(status_code=500, detail=str(exc)) from exc