"""ESMFold structure prediction tool for the Amina CLI."""

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

METADATA = {
    "name": "esmfold",
    "display_name": "ESMFold",
    "category": "folding",
    "description": "Predict protein structure from a sequence using ESMFold",
    "modal_function_name": "esmfold_worker",
    "modal_app_name": "esmfold-api",
    "status": "available",
    "outputs": {
        "pdb_filepath": "Predicted structure in PDB format",
        "csv_filepath": "Per-residue pLDDT confidence scores",
        "plot_filepath": "pLDDT score visualization plot",
    },
}

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("esmfold")
    def run_esmfold(
        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",
            exists=True,
        ),
        output: Optional[Path] = typer.Option(
            None,
            "--output",
            "-o",
            help="Output directory for results (required unless --background)",
        ),
        job_name: Optional[str] = typer.Option(
            None,
            "--job-name",
            "-j",
            help="Custom job name for output files (default: random 4-letter code)",
        ),
        background: bool = typer.Option(
            False,
            "--background",
            "-b",
            help="Submit job and return immediately without waiting for completion",
        ),
    ):
        """
        Predict protein structure from sequence using ESMFold.

        Provide either a sequence directly or a FASTA file.

        Examples:
            amina run esmfold --sequence "MKFLILLFNILCLFPVLAADNH" -o ./output/
            amina run esmfold --fasta ./input.fasta -o ./output/
            amina run esmfold --sequence "MKFLIL" -j myjob -o ./output/
            amina run esmfold --sequence "MKFLIL" -o ./output/ --background
        """
        # Validate required options
        if output is None and not background:
            console.print("[red]Error:[/red] --output / -o is required (unless using --background)")
            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)

        # Build params
        params = {}
        if sequence:
            params["sequence"] = sequence
        elif fasta:
            # Read sequence from FASTA
            content = fasta.read_text()
            # Simple FASTA parsing - skip header lines
            lines = [line.strip() for line in content.split("\n") if not line.startswith(">")]
            params["sequence"] = "".join(lines)
            console.print(f"Read sequence from {fasta} ({len(params['sequence'])} residues)")

        if job_name:
            params["job_name"] = job_name

        # Execute
        run_tool_with_progress("esmfold", params, output, background=background)
