"""PDB Distance Calculator tool for the Amina CLI."""

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

METADATA = {
    "name": "pdb-distance-calculator",  # Match TOOL_REGISTRY naming convention
    "display_name": "PDB Distance Calculator",
    "category": "utilities",
    "description": "Calculate inter-residue distances in protein structures",
    "modal_function_name": "pdb_distance_calculator_worker",
    "modal_app_name": "pdb-distance-calculator-api",
    "status": "available",
    "outputs": {
        "json_filepath": "JSON file with distance calculation results",
    },
}

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("pdb-distance-calculator")
    def run_distance_calculator(
        pdb: Path = typer.Option(
            ...,
            "--pdb",
            "-p",
            help="Path to input PDB file",
            exists=True,
        ),
        # Single distance options
        chain1: Optional[str] = typer.Option(
            None,
            "--chain1",
            help="Chain ID for first residue (e.g., 'A')",
        ),
        resnum1: Optional[int] = typer.Option(
            None,
            "--resnum1",
            help="Residue number for first residue",
        ),
        chain2: Optional[str] = typer.Option(
            None,
            "--chain2",
            help="Chain ID for second residue (e.g., 'A')",
        ),
        resnum2: Optional[int] = typer.Option(
            None,
            "--resnum2",
            help="Residue number for second residue",
        ),
        atom1: str = typer.Option(
            "CA",
            "--atom1",
            help="Atom name for first residue (default: CA for alpha carbon)",
        ),
        atom2: str = typer.Option(
            "CA",
            "--atom2",
            help="Atom name for second residue (default: CA for alpha carbon)",
        ),
        # Contact map options
        contact_map: bool = typer.Option(
            False,
            "--contact-map",
            "-c",
            help="Generate contact map for a chain instead of single distance",
        ),
        chain: Optional[str] = typer.Option(
            None,
            "--chain",
            help="Chain ID for contact map calculation",
        ),
        cutoff: float = typer.Option(
            8.0,
            "--cutoff",
            help="Distance cutoff in Angstroms for contacts (default: 8.0)",
        ),
        # Model selection
        model_id: int = typer.Option(
            0,
            "--model",
            help="PDB model number to use (default: 0 for first model)",
        ),
        # Standard options
        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)",
        ),
    ):
        """
        Calculate inter-residue distances in protein structures.

        Supports two calculation modes:

        1. SINGLE DISTANCE: Measure distance between two specific residues
           Requires: --chain1, --resnum1, --chain2, --resnum2

        2. CONTACT MAP: Generate all pairwise distances within a cutoff
           Requires: --contact-map and --chain

        Examples:
            # Single distance between residue 10 and 50 on chain A
            amina run pdb-distance-calculator -p protein.pdb --chain1 A --resnum1 10 --chain2 A --resnum2 50 -o ./results/

            # Distance using CB atoms instead of CA
            amina run pdb-distance-calculator -p protein.pdb --chain1 A --resnum1 10 --chain2 A --resnum2 50 --atom1 CB --atom2 CB -o ./results/

            # Contact map for chain A with 8.0 Angstrom cutoff
            amina run pdb-distance-calculator -p protein.pdb --contact-map --chain A -o ./results/

            # Contact map with custom cutoff
            amina run pdb-distance-calculator -p protein.pdb --contact-map --chain A --cutoff 10.0 -o ./results/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

        # Read PDB file content
        pdb_content = pdb.read_text()
        console.print(f"Read structure from {pdb}")

        # Build params
        params = {
            "pdb_content": pdb_content,
            "pdb_filename": pdb.name,
            "model_id": model_id,
        }

        # Determine calculation type based on options
        if contact_map:
            # Contact map mode
            if not chain:
                console.print("[red]Error:[/red] --chain is required for contact map calculation")
                raise typer.Exit(1)

            params["calculation_type"] = "contact_map"
            params["contact_map_config"] = {
                "chain": chain.upper(),
                "cutoff": cutoff,
                "atom_type": atom1,  # Use atom1 as the atom type for contact map
            }
            console.print(f"Calculating contact map for chain {chain} with cutoff {cutoff} A")

        else:
            # Single distance mode
            if not all([chain1, resnum1, chain2, resnum2]):
                console.print(
                    "[red]Error:[/red] For single distance, provide all of: --chain1, --resnum1, --chain2, --resnum2"
                )
                console.print("[yellow]Hint:[/yellow] Use --contact-map for contact map calculation")
                raise typer.Exit(1)

            params["calculation_type"] = "single_distance"
            params["single_pair"] = {
                "chain1": chain1.upper(),
                "resnum1": resnum1,
                "chain2": chain2.upper(),
                "resnum2": resnum2,
                "atom1": atom1.upper(),
                "atom2": atom2.upper(),
            }
            console.print(f"Measuring distance: {chain1}:{resnum1}:{atom1} <-> {chain2}:{resnum2}:{atom2}")

        if job_name:
            params["job_name"] = job_name

        run_tool_with_progress("pdb-distance-calculator", params, output)
