"""ESM2 embedding extraction tool for the Amina CLI."""

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

METADATA = {
    "name": "esm2-embedding",
    "display_name": "ESM2 Embedding",
    "category": "properties",
    "description": "Extract sequence-level protein embeddings from FASTA files using ESM2 models",
    "modal_function_name": "esm2_embedding_worker",
    "modal_app_name": "esm2-embedding-api",
    "status": "available",
    "outputs": {
        "csv_filepath": "CSV file with sequence embeddings (sequence_id, sequence, dim_0...dim_N)",
        "tsne_plot_filepath": "Static t-SNE visualization of embeddings (PNG)",
        "tsne_plot_html_filepath": "Interactive t-SNE visualization with zoom/pan/hover (HTML)",
        "clusters_csv_filepath": "Cluster assignments (CSV, unless --no-cluster)",
        "cluster_summary_filepath": "Cluster statistics and quality metrics (JSON, unless --no-cluster)",
    },
}

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("esm2-embedding")
    def run_esm2_embedding(
        fasta: Path = typer.Option(
            ...,
            "--fasta",
            "-f",
            help="Path to FASTA file containing protein sequences",
            exists=True,
        ),
        output: Optional[Path] = typer.Option(
            None,
            "--output",
            "-o",
            help="Output directory for results (required unless --background)",
        ),
        esm_variant: str = typer.Option(
            "650M",
            "--model",
            "-m",
            help="ESM2 model variant: 650M (best), 150M, 35M, or 8M (fastest)",
        ),
        batch_size: int = typer.Option(
            8,
            "--batch-size",
            "-b",
            help="Batch size for processing (1-32)",
        ),
        no_tsne: bool = typer.Option(
            False,
            "--no-tsne",
            help="Skip t-SNE visualization generation",
        ),
        no_cluster: bool = typer.Option(
            False,
            "--no-cluster",
            help="Disable HDBSCAN clustering (clustering is ON by default)",
        ),
        cluster_min_size: Optional[int] = typer.Option(
            None,
            "--cluster-min-size",
            help="Minimum cluster size (auto-tuned via DBCV if not specified)",
        ),
        cluster_min_samples: Optional[int] = typer.Option(
            None,
            "--cluster-min-samples",
            help="Min samples for density estimation (auto-tuned if not specified)",
        ),
        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",
            help="Submit job and return immediately without waiting for completion",
        ),
    ):
        """
        Extract protein embeddings from a FASTA file using ESM2 models.

        Generates dense vector representations (embeddings) of protein sequences
        that capture evolutionary, structural, and functional information.

        By default, also performs HDBSCAN clustering to group similar sequences.
        Use --no-cluster to disable clustering.

        Model variants (embedding dimensions):
          - 650M: 1280-dim - Highest accuracy (default)
          - 150M: 640-dim  - Balanced
          - 35M:  480-dim  - Faster
          - 8M:   320-dim  - Fastest, lower accuracy

        Examples:
            amina run esm2-embedding --fasta sequences.fasta -o ./output/
            amina run esm2-embedding -f seqs.fa -m 8M -o ./output/  # Fast mode
            amina run esm2-embedding -f seqs.fa --no-tsne -o ./output/
            amina run esm2-embedding -f seqs.fa --no-cluster -o ./output/  # Skip clustering
            amina run esm2-embedding -f seqs.fa --cluster-min-size 5 -o ./output/  # Custom params
        """
        # 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 ESM variant
        valid_variants = ["650M", "150M", "35M", "8M"]
        if esm_variant.upper() not in valid_variants:
            console.print(
                f"[red]Error:[/red] Invalid model variant '{esm_variant}'. Choose from: {', '.join(valid_variants)}"
            )
            raise typer.Exit(1)

        # Validate batch size
        if batch_size < 1 or batch_size > 32:
            console.print("[red]Error:[/red] Batch size must be between 1 and 32")
            raise typer.Exit(1)

        # Read FASTA content
        fasta_content = fasta.read_text()
        if not fasta_content.strip().startswith(">"):
            console.print("[red]Error:[/red] File does not appear to be in FASTA format")
            raise typer.Exit(1)

        # Count sequences for display
        num_sequences = fasta_content.count(">")
        console.print(f"Read {num_sequences} sequence(s) from {fasta}")

        # Build params
        params = {
            "fasta_content": fasta_content,
            "fasta_filename": fasta.name,
            "esm_variant": esm_variant.upper(),
            "batch_size": batch_size,
            "generate_tsne": not no_tsne,
            "cluster_by_embeddings": not no_cluster,
            "cluster_min_size": cluster_min_size,
            "cluster_min_samples": cluster_min_samples,
        }

        if job_name:
            params["job_name"] = job_name

        # Execute
        run_tool_with_progress(METADATA["name"], params, output, background=background)
