"""US-align universal structure alignment tool for the Amina CLI."""

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

METADATA = {
    "name": "usalign",
    "display_name": "US-align",
    "category": "analysis",
    "description": "Calculate TM-scores between protein structures using US-align",
    "modal_function_name": "usalign_worker",
    "modal_app_name": "usalign-api",
    "status": "available",
    "outputs": {
        "aligned_pdb_filepath": "Superimposed structure aligned to target",
    },
}

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("usalign")
    def run_usalign(
        mobile: Path = typer.Option(
            ...,
            "--mobile",
            "-m",
            help="Path to mobile/query structure (PDB file to be aligned)",
            exists=True,
        ),
        target: Path = typer.Option(
            ...,
            "--target",
            "-t",
            help="Path to target/reference structure (PDB file to align against)",
            exists=True,
        ),
        mode: str = typer.Option(
            "monomer",
            "--mode",
            help="Alignment mode: 'monomer' (single chains) or 'complex' (multi-chain with chain mapping)",
        ),
        mobile_chains: Optional[str] = typer.Option(
            None,
            "--mobile-chains",
            help="Chains to use from mobile structure (e.g., 'A' or 'A,B'). Default: all chains",
        ),
        target_chains: Optional[str] = typer.Option(
            None,
            "--target-chains",
            help="Chains to use from target structure (e.g., 'A' or 'A,B'). Default: all chains",
        ),
        ter_option: int = typer.Option(
            1,
            "--ter",
            help="TER handling: 0=all chains combined, 1=TER separates chains, 2=MODEL separates chains",
        ),
        mm_option: int = typer.Option(
            2,
            "--mm",
            help="Complex mode chain mapping: 1=sequential, 2=optimal mapping, 4=greedy (only for --mode complex)",
        ),
        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 tracking (default: random 4-letter code)",
        ),
    ):
        """
        Calculate TM-scores between two protein structures using US-align.

        US-align calculates the Template Modeling score (TM-score) which measures
        structural similarity independent of sequence. TM-score ranges from 0 to 1:
        - TM-score > 0.5: Generally same fold
        - TM-score > 0.17: Better than random

        Monomer mode (-mode monomer): Compares single chains or structures.
        Complex mode (-mode complex): Compares multi-chain complexes with optimal
        chain mapping to find the best superposition.

        Examples:
            # Basic monomer alignment
            amina run usalign -m designed.pdb -t native.pdb -o ./results/

            # Align specific chains
            amina run usalign -m model.pdb -t reference.pdb --mobile-chains A --target-chains A -o ./results/

            # Complex alignment with optimal chain mapping
            amina run usalign -m complex.pdb -t reference.pdb --mode complex -o ./results/

            # Complex alignment with greedy mapping
            amina run usalign -m complex.pdb -t reference.pdb --mode complex --mm 4 -o ./results/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

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

        # Validate ter_option
        if ter_option not in [0, 1, 2]:
            console.print(f"[red]Error:[/red] Invalid --ter option '{ter_option}'. Must be 0, 1, or 2")
            raise typer.Exit(1)

        # Validate mm_option
        if mm_option not in [1, 2, 4]:
            console.print(f"[red]Error:[/red] Invalid --mm option '{mm_option}'. Must be 1, 2, or 4")
            raise typer.Exit(1)

        # Read file contents
        mobile_content = mobile.read_text()
        target_content = target.read_text()
        console.print(f"Read mobile structure from {mobile}")
        console.print(f"Read target structure from {target}")

        # Build params
        params = {
            "mobile_pdb_content": mobile_content,
            "mobile_pdb_filename": mobile.name,
            "target_pdb_content": target_content,
            "target_pdb_filename": target.name,
            "alignment_mode": mode.lower(),
            "ter_option": ter_option,
            "mm_option": mm_option,
        }

        # Parse chain selections
        if mobile_chains:
            chains = [c.strip().upper() for c in mobile_chains.split(",") if c.strip()]
            params["mobile_chains"] = chains
            console.print(f"Using chains from mobile structure: {chains}")

        if target_chains:
            chains = [c.strip().upper() for c in target_chains.split(",") if c.strip()]
            params["target_chains"] = chains
            console.print(f"Using chains from target structure: {chains}")

        if job_name:
            params["job_name"] = job_name

        run_tool_with_progress("usalign", params, output)
