"""Monte Carlo sequence optimization tool for the Amina CLI."""

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

METADATA = {
    "name": "protein-mc",
    "display_name": "Protein-MC",
    "category": "design",
    "description": "Monte Carlo sequence optimization using pLM-guided simulation",
    "modal_function_name": "protein_mc_worker",
    "modal_app_name": "protein-mc-api",
    "status": "available",
    "outputs": {
        "csv_filepath": "Optimization results and energy scores",
        "plot_filepath": "Energy optimization plot",
        "fasta_filepath": "Generated optimized sequences",
    },
}

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("protein-mc")
    def run_protein_mc(
        sequence: Optional[str] = typer.Option(
            None,
            "--sequence",
            "-s",
            help="Amino acid sequence (single letter code)",
        ),
        fasta: Optional[Path] = typer.Option(
            None,
            "--fasta",
            "-f",
            help="Path to FASTA file containing the protein sequence",
            exists=True,
        ),
        temperature: float = typer.Option(
            0.01,
            "--temperature",
            "-t",
            help="Temperature for Monte Carlo sampling (higher = more exploration)",
        ),
        num_steps: int = typer.Option(
            100,
            "--num-steps",
            "-n",
            help="Number of Monte Carlo steps to run",
        ),
        protected_positions: Optional[str] = typer.Option(
            None,
            "--protected",
            "-p",
            help="Comma-separated list of positions to protect from mutation (1-indexed)",
        ),
        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 diverse mutant protein sequences using Monte Carlo optimization.

        Uses pLM-guided Monte Carlo simulation to optimize protein sequences
        in embedding space while preserving structural features.

        Provide either a sequence directly or a FASTA file.

        Examples:
            amina run protein-mc --sequence "MKFLILLFNILCLFPVLAADNH" -o ./output/
            amina run protein-mc --fasta ./input.fasta -t 0.01 -n 200 -o ./output/
            amina run protein-mc --sequence "MKFLIL" --protected "1,2,3" -o ./output/
            amina run protein-mc --sequence "MKFLIL" -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 input
        if not sequence and not fasta:
            console.print("[red]Error:[/red] Provide either --sequence or --fasta")
            raise typer.Exit(1)

        if sequence and fasta:
            console.print("[red]Error:[/red] Provide only one of --sequence or --fasta")
            raise typer.Exit(1)

        # Validate temperature
        if temperature <= 0:
            console.print("[red]Error:[/red] Temperature must be positive")
            raise typer.Exit(1)

        # Validate num_steps
        if num_steps <= 0:
            console.print("[red]Error:[/red] Number of steps must be positive")
            raise typer.Exit(1)

        # Build params
        params = {
            "temperature": temperature,
            "num_steps": num_steps,
        }

        if sequence:
            params["sequence"] = sequence
        elif fasta:
            # CLI mode: send file content directly
            fasta_content = fasta.read_text()
            params["fasta_content"] = fasta_content
            params["fasta_filename"] = fasta.name
            console.print(f"Read FASTA from {fasta}")

        # Parse protected positions if provided
        if protected_positions:
            try:
                positions = [int(p.strip()) for p in protected_positions.split(",")]
                if not all(p > 0 for p in positions):
                    console.print("[red]Error:[/red] Protected positions must be positive integers (1-indexed)")
                    raise typer.Exit(1)
                params["protected_positions"] = positions
            except ValueError:
                console.print("[red]Error:[/red] Invalid protected positions format. Use comma-separated integers.")
                raise typer.Exit(1)

        if job_name:
            params["job_name"] = job_name

        # Execute
        run_tool_with_progress("protein-mc", params, output)
