"""ESM-IF1 inverse folding tool for the Amina CLI."""

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

METADATA = {
    "name": "esm-if1",
    "display_name": "ESM-IF1",
    "category": "design",
    "description": "Generate protein sequences from structure via inverse folding",
    "modal_function_name": "esm_if1_worker",
    "modal_app_name": "esm-if1-api",
    "status": "available",
    "outputs": {
        "fasta_filepath": "FASTA file with designed sequences (sample mode)",
        "csv_filepath": "CSV file with scores and rankings",
    },
}

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("esm-if1")
    def run_esm_if1(
        pdb: Path = typer.Option(
            ...,
            "--pdb",
            "-p",
            help="Path to PDB/CIF file with target structure",
            exists=True,
        ),
        chain: str = typer.Option(
            ...,
            "--chain",
            "-c",
            help="Target chain ID to design/score (e.g., 'A', 'B', 'C')",
        ),
        mode: str = typer.Option(
            "sample",
            "--mode",
            "-m",
            help="Operation mode: 'sample' (generate sequences) or 'score' (evaluate sequences)",
        ),
        fasta: Optional[Path] = typer.Option(
            None,
            "--fasta",
            "-f",
            help="Path to FASTA file with sequences to score (required for score mode)",
            exists=True,
        ),
        num_samples: int = typer.Option(
            3,
            "--num-samples",
            "-n",
            help="Number of sequence designs to generate (1-300, sample mode only)",
        ),
        temperature: float = typer.Option(
            1.0,
            "--temperature",
            "-t",
            help="Sampling temperature (0-5.0). Lower for native-like, higher for diversity",
        ),
        multichain: bool = typer.Option(
            False,
            "--multichain/--single-chain",
            help="Use entire complex as context for prediction",
        ),
        padding_length: int = typer.Option(
            10,
            "--padding",
            help="Padding between chains for scoring mode",
        ),
        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)",
        ),
    ):
        """
        Generate protein sequences from structure using ESM-IF1 inverse folding.

        ESM-IF1 (Inverse Folding 1) is a deep learning model that generates
        amino acid sequences predicted to fold into a given protein backbone
        structure. It can operate in two modes:

        SAMPLE MODE (default): Generate new sequences compatible with the structure.
        The model samples sequences ranked by log-likelihood (confidence that the
        sequence would fold into the given structure).

        SCORE MODE: Evaluate how well existing sequences fit the structure.
        Requires a FASTA file with sequences to score.

        Examples:
            # Generate 10 sequences at temperature 0.5
            amina run esm-if1 --pdb ./protein.pdb --chain A -n 10 -t 0.5 -o ./output/

            # Score existing sequences
            amina run esm-if1 --pdb ./protein.pdb --chain A --mode score --fasta ./seqs.fasta -o ./output/

            # Design with multichain context
            amina run esm-if1 --pdb ./complex.pdb --chain B --multichain -o ./output/

            # Custom job name
            amina run esm-if1 --pdb ./protein.pdb --chain A -j myjob -o ./output/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

        # Validate mode
        mode = mode.lower().strip()
        if mode not in ["sample", "score"]:
            console.print(f"[red]Error:[/red] Invalid mode '{mode}'. Must be 'sample' or 'score'")
            raise typer.Exit(1)

        # Validate chain ID
        chain = chain.strip().upper()
        if len(chain) != 1 or not chain.isalpha():
            console.print(f"[red]Error:[/red] Chain ID must be a single letter (e.g., 'A'), got '{chain}'")
            raise typer.Exit(1)

        # Validate score mode requires FASTA
        if mode == "score" and fasta is None:
            console.print("[red]Error:[/red] Score mode requires --fasta / -f with sequences to evaluate")
            raise typer.Exit(1)

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

        # Build params dict
        params = {
            "pdb_content": pdb_content,
            "chain_id": chain,
            "mode": mode,
            "temperature": temperature,
            "num_samples": num_samples,
            "multichain_backbone": multichain,
            "padding_length": padding_length,
        }

        # Add FASTA content for score mode
        if fasta:
            fasta_content = fasta.read_text()
            params["fasta_content"] = fasta_content
            console.print(f"Read FASTA from {fasta}")

        # Add job name if provided
        if job_name:
            params["job_name"] = job_name

        run_tool_with_progress("esm-if1", params, output)
