"""sfold — protein complex structure prediction with Boltz-2, on HF ZeroGPU."""

import re
import string
import tempfile
import time
from pathlib import Path

import gradio as gr
import py3Dmol
import spaces
import yaml
from boltz.main import download_boltz2, get_cache_path, predict

NUM_ROWS = 6
VALID_AMINO_ACIDS = set("ACDEFGHIKLMNPQRSTVWYXBZUO")
CHAIN_IDS = string.ascii_uppercase

EXAMPLE_UBIQUITIN = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG"
EXAMPLE_CAP = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"

CACHE_DIR = get_cache_path()
print(f"Preparing Boltz-2 weights/CCD cache at {CACHE_DIR} (this happens once at startup)...")
# download_boltz2 writes straight into this directory without creating it first — the normal
# CLI flow apparently creates it elsewhere before reaching this call; we call it directly, so
# we're responsible for making sure it exists.
Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)
download_boltz2(Path(CACHE_DIR))
print("Cache ready.")


def clean_sequence(raw_sequence: str) -> str:
    """Strip a FASTA header/whitespace, uppercase, drop anything that isn't a letter."""
    lines = [line for line in raw_sequence.strip().splitlines() if not line.startswith(">")]
    sequence = "".join(lines).strip().upper()
    return re.sub(r"[^A-Z]", "", sequence)


def render_structure(pdb_str: str) -> str:
    view = py3Dmol.view(width=600, height=500)
    view.addModel(pdb_str, "pdb")
    view.setStyle({"cartoon": {"colorscheme": {"prop": "b", "gradient": "roygb", "min": 50, "max": 90}}})
    view.zoomTo()
    raw_html = view._make_html()
    escaped = raw_html.replace('"', """)
    return f'<iframe style="width: 100%; height: 500px; border: none;" srcdoc="{escaped}"></iframe>'


@spaces.GPU(duration=300)
def fold_rows(*args, progress=gr.Progress()):
    *seqs, recycling_steps, sampling_steps = args

    entities = []
    chain_cursor = 0
    for raw_seq, copies in zip(seqs[:NUM_ROWS], seqs[NUM_ROWS:]):
        sequence = clean_sequence(raw_seq or "")
        if not sequence:
            continue
        invalid_chars = set(sequence) - VALID_AMINO_ACIDS
        if invalid_chars:
            raise gr.Error(f"Invalid characters in sequence: {', '.join(sorted(invalid_chars))}")

        num_copies = max(1, int(copies or 1))
        chain_ids = list(CHAIN_IDS[chain_cursor : chain_cursor + num_copies])
        chain_cursor += num_copies

        entities.append(
            {
                "protein": {
                    "id": chain_ids if num_copies > 1 else chain_ids[0],
                    "sequence": sequence,
                    # Explicit sentinel: skip MSA search entirely (single-sequence mode).
                    # Without this, Boltz raises unless --use_msa_server is set, since a
                    # bare protein entity defaults to "MSA needs generating".
                    "msa": "empty",
                }
            }
        )

    if not entities:
        raise gr.Error("Please enter at least one protein sequence.")

    total_residues = sum(
        len(e["protein"]["sequence"]) * (len(e["protein"]["id"]) if isinstance(e["protein"]["id"], list) else 1)
        for e in entities
    )

    progress(0.1, desc=f"Preparing input: {len(entities)} entities, {chain_cursor} chains, {total_residues} residues total...")

    work_dir = Path(tempfile.mkdtemp(prefix="sfold_"))
    yaml_path = work_dir / "input.yaml"
    with open(yaml_path, "w") as f:
        yaml.safe_dump({"version": 1, "sequences": entities}, f, sort_keys=False)
    out_dir = work_dir / "out"

    progress(0.2, desc="Running Boltz-2 (this can take a while)...")
    start = time.time()
    # `predict` is a click.Command (decorated with @cli.command()), not a plain function —
    # calling it directly routes through click's CLI arg parsing and crashes on kwargs like
    # `data`. `.callback` is click's own stored reference to the original undecorated function.
    predict.callback(
        data=str(yaml_path),
        out_dir=str(out_dir),
        cache=CACHE_DIR,
        devices=1,
        accelerator="gpu",
        recycling_steps=int(recycling_steps),
        sampling_steps=int(sampling_steps),
        diffusion_samples=1,
        output_format="pdb",
        use_msa_server=False,
        model="boltz2",
        seed=0,
    )
    elapsed = time.time() - start

    progress(0.9, desc="Reading structure file...")
    result_dir = out_dir / f"boltz_results_{yaml_path.stem}"
    pdb_files = sorted(result_dir.rglob("*.pdb"))
    if not pdb_files:
        raise gr.Error("Boltz-2 did not produce a structure file — check the Space logs.")
    pdb_path = pdb_files[0]
    pdb_str = pdb_path.read_text()

    viewer_html = render_structure(pdb_str)
    summary_lines = [
        f"**Entities:** {len(entities)}",
        f"**Chains:** {chain_cursor}",
        f"**Total length:** {total_residues} residues",
        f"**Inference time:** {elapsed:.1f}s",
    ]

    return viewer_html, "  \n".join(summary_lines), str(pdb_path)


def load_example(sequence: str):
    return sequence, 1


with gr.Blocks(title="sfold — Boltz-2 Protein Complex Prediction") as demo:
    gr.Markdown(
        """
        # 🧬 sfold
        Predict 3D structure for protein complexes using
        [Boltz-2](https://github.com/jwohlwend/boltz) (MIT license, NVIDIA/MIT lineage).
        Fill in one row per distinct protein, and set copies > 1 for a homooligomer —
        each row becomes one "entity" that Boltz-2 assembles into a single complex.
        Runs single-sequence (no MSA server), so results are best for well-behaved targets.
        """
    )

    seq_boxes: list[gr.Textbox] = []
    copies_boxes: list[gr.Slider] = []

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                example1_btn = gr.Button("Load: Ubiquitin", size="sm")
                example2_btn = gr.Button("Load: CAP protein", size="sm")

            for row_idx in range(NUM_ROWS):
                with gr.Row():
                    seq = gr.Textbox(
                        label=f"Chain {row_idx + 1} sequence",
                        lines=4,
                        max_lines=4,
                        scale=4,
                        placeholder="Amino acid sequence (leave empty to skip this row)...",
                    )
                    copies = gr.Slider(1, 16, value=1, step=1, label="Copies", scale=1)
                seq_boxes.append(seq)
                copies_boxes.append(copies)

            with gr.Accordion("Advanced settings", open=False):
                recycling_steps = gr.Slider(1, 6, value=3, step=1, label="Recycling steps")
                sampling_steps = gr.Slider(20, 200, value=200, step=10, label="Diffusion sampling steps")
                gr.Markdown("Higher values improve accuracy but increase inference time.")

            fold_btn = gr.Button("Fold", variant="primary")
            summary_output = gr.Markdown()
            pdb_output = gr.File(label="Download PDB", interactive=False)

        with gr.Column(scale=2):
            structure_output = gr.HTML(label="Predicted structure")

    example1_btn.click(lambda: load_example(EXAMPLE_UBIQUITIN), outputs=[seq_boxes[0], copies_boxes[0]])
    example2_btn.click(lambda: load_example(EXAMPLE_CAP), outputs=[seq_boxes[0], copies_boxes[0]])

    fold_btn.click(
        fn=fold_rows,
        inputs=[*seq_boxes, *copies_boxes, recycling_steps, sampling_steps],
        outputs=[structure_output, summary_output, pdb_output],
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(theme=gr.themes.Soft())