"""Protein solubility prediction using AminoSol."""

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

METADATA = {
    "name": "aminosol",  # Must match gateway tool_registry key
    "display_name": "AminoSol",
    "category": "properties",
    "description": "Predict protein solubility for E. coli expression using fine-tuned ESM-2 models",
    "modal_function_name": "aminosol_worker",
    "modal_app_name": "aminosol-api",
    "status": "available",
    "outputs": {
        "csv_filepath": "CSV file with solubility predictions for each sequence",
    },
}

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("aminosol")
    def run_aminosol(
        sequences: Optional[str] = typer.Option(
            None,
            "--sequences",
            "-s",
            help="Comma-separated protein sequences (e.g., 'MVLSPADKTNVK,MKFLILLFNILC')",
        ),
        fasta: Optional[Path] = typer.Option(
            None,
            "--fasta",
            "-f",
            help="Path to FASTA file with one or more sequences",
            exists=True,
        ),
        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)",
        ),
        batch_size: int = typer.Option(
            8,
            "--batch-size",
            "-b",
            help="Batch size for prediction (1-64, higher = faster but more memory)",
            min=1,
            max=64,
        ),
    ):
        """
        Predict protein solubility for E. coli expression.

        Uses a fine-tuned ESM-2 model to predict whether proteins will be
        soluble when expressed in E. coli. Supports single sequences,
        comma-separated multiple sequences, or FASTA files.

        Provide either sequences directly or a FASTA file.

        Examples:
            amina run aminosol --sequences "MKFLILLFNILCLFPVLAADNH" -o ./output/
            amina run aminosol --sequences "MVLSPADKTNVK,MKFLILLFNILC" -o ./output/
            amina run aminosol --fasta ./proteins.fasta -o ./output/
            amina run aminosol --fasta ./proteins.fasta -j myjob -b 16 -o ./output/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

        # Validate input - must provide one of sequences or fasta
        if not sequences and not fasta:
            console.print("[red]Error:[/red] Provide either --sequences or --fasta")
            raise typer.Exit(1)

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

        # Build params
        params = {"batch_size": batch_size}

        if sequences:
            # Parse comma-separated sequences
            seq_list = [seq.strip().upper() for seq in sequences.split(",") if seq.strip()]
            if not seq_list:
                console.print("[red]Error:[/red] No valid sequences provided")
                raise typer.Exit(1)
            params["sequences"] = seq_list
            console.print(f"Processing {len(seq_list)} sequence(s)")

        elif fasta:
            # Parse FASTA file to extract sequences
            seq_list = _parse_fasta_file(fasta)
            if not seq_list:
                console.print(f"[red]Error:[/red] No sequences found in {fasta}")
                raise typer.Exit(1)
            params["sequences"] = seq_list
            console.print(f"Read {len(seq_list)} sequence(s) from {fasta}")

        if job_name:
            params["job_name"] = job_name

        # Execute
        run_tool_with_progress("aminosol", params, output)


def _parse_fasta_file(fasta_path: Path) -> list[str]:
    """Parse a FASTA file and return list of sequences."""
    try:
        content = fasta_path.read_text()
    except UnicodeDecodeError:
        raise typer.BadParameter(
            f"Could not read {fasta_path}: file has invalid encoding. Please ensure it's a valid UTF-8 text file."
        )
    except PermissionError:
        raise typer.BadParameter(f"Permission denied reading {fasta_path}")
    except OSError as e:
        raise typer.BadParameter(f"Could not read {fasta_path}: {e}")

    sequences = []
    current_sequence = []
    for line in content.split("\n"):
        line = line.strip()
        if not line:
            continue
        if line.startswith(">"):
            # Save previous sequence if exists
            if current_sequence:
                sequences.append("".join(current_sequence).upper())
                current_sequence = []
        else:
            current_sequence.append(line)

    # Don't forget the last sequence
    if current_sequence:
        sequences.append("".join(current_sequence).upper())

    return sequences
