"""ProteinMPNN sequence design tool for the Amina CLI."""

import typer
from pathlib import Path
from typing import Optional
from rich.console import Console

METADATA = {
    "name": "proteinmpnn",
    "display_name": "ProteinMPNN",
    "category": "design",
    "description": "Design protein sequences compatible with a 3D structure",
    "modal_function_name": "proteinmpnn_worker",
    "modal_app_name": "proteinmpnn-api",
    "status": "available",
    "outputs": {
        "fasta_filepath": "FASTA file with designed sequences",
        "scores_filepath": "CSV file with sequence scores",
    },
}

console = Console()


def register(app: typer.Typer):
    """Register this tool's command with the app."""
    from amina_cli.commands.tools import run_tool_with_progress

    @app.command("proteinmpnn")
    def run_proteinmpnn(
        pdb: Path = typer.Option(
            ...,
            "--pdb",
            "-p",
            help="Path to PDB file with target structure",
            exists=True,
        ),
        chains: Optional[str] = typer.Option(
            None,
            "--chains",
            "-c",
            help="Comma-separated chain IDs to design (default: all)",
        ),
        fixed_residues: Optional[str] = typer.Option(
            None,
            "--fixed",
            "-f",
            help="Residues to keep fixed, format: 'A:10,15,20;B:5,25'",
        ),
        model_type: str = typer.Option(
            "vanilla",
            "--model-type",
            "-m",
            help="Model variant: vanilla, soluble, or ca_only",
        ),
        num_sequences: int = typer.Option(
            10,
            "--num-sequences",
            "-n",
            help="Number of sequences to generate (1-50)",
        ),
        temperature: float = typer.Option(
            1.0,
            "--temperature",
            help="Sampling temperature (0.01-2.0)",
        ),
        seed: int = typer.Option(
            37,
            "--seed",
            help="Random seed for reproducibility (0 = random)",
        ),
        omit_aas: str = typer.Option(
            "X",
            "--omit-aas",
            help="Amino acids to exclude from design (e.g., 'C' to exclude cysteines)",
        ),
        output: Path = typer.Option(
            ...,
            "--output",
            "-o",
            help="Output directory for results",
        ),
        job_name: Optional[str] = typer.Option(
            None,
            "--job-name",
            "-j",
            help="Custom job name for output files (default: random 4-letter code)",
        ),
    ):
        """
        Design protein sequences compatible with a 3D structure using ProteinMPNN.

        ProteinMPNN is a deep learning model for protein sequence design that
        generates sequences predicted to fold into the given backbone structure.

        Examples:
            amina run proteinmpnn --pdb ./structure.pdb -o ./output/
            amina run proteinmpnn --pdb ./structure.pdb --chains A,B -n 10 -o ./output/
            amina run proteinmpnn --pdb ./structure.pdb -j myjob -o ./output/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

        # Read PDB file content
        pdb_content = pdb.read_text()
        console.print(f"Read PDB from {pdb}")

        params = {
            "pdb_content": pdb_content,
            "model_type": model_type,
            "num_sequences": num_sequences,
            "temperature": temperature,
            "seed": seed,
            "omit_AAs": omit_aas,
        }

        if chains:
            params["chains_to_design"] = [c.strip().upper() for c in chains.split(",")]

        if fixed_residues:
            # Parse fixed residues: "A:10,15,20;B:5,25" -> {"A": [10, 15, 20], "B": [5, 25]}
            fixed_dict = {}
            for spec in fixed_residues.split(";"):
                if ":" in spec:
                    chain, positions = spec.split(":", 1)
                    pos_list = [int(p.strip()) for p in positions.split(",") if p.strip().isdigit()]
                    if pos_list:
                        fixed_dict[chain.strip().upper()] = pos_list
            if fixed_dict:
                params["fixed_residues"] = fixed_dict

        if job_name:
            params["job_name"] = job_name

        run_tool_with_progress("proteinmpnn", params, output)
