"""DiffDock molecular docking tool for the Amina CLI."""

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

METADATA = {
    "name": "diffdock",
    "display_name": "DiffDock",
    "category": "interactions",
    "description": "Diffusion-based molecular docking for protein-ligand binding prediction",
    "modal_function_name": "diffdock_worker",
    "modal_app_name": "diffdock-api",
    "status": "available",
    "outputs": {
        "top_pose_filepath": "Top-ranked docked pose in SDF format",
        "confidence_json_filepath": "JSON file with confidence scores for all poses",
    },
}

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("diffdock")
    def run_diffdock(
        # Protein input options (mutually exclusive)
        protein_sequence: Optional[str] = typer.Option(
            None,
            "--protein-sequence",
            "-ps",
            help="Protein sequence (will use ESMFold for structure prediction)",
        ),
        protein_pdb: Optional[Path] = typer.Option(
            None,
            "--protein-pdb",
            "-pp",
            help="Path to protein PDB file",
            exists=True,
        ),
        # Ligand input options (mutually exclusive)
        ligand_smiles: Optional[str] = typer.Option(
            None,
            "--ligand-smiles",
            "-ls",
            help="SMILES string for the ligand molecule",
        ),
        ligand_sdf: Optional[Path] = typer.Option(
            None,
            "--ligand-sdf",
            "-lf",
            help="Path to ligand SDF file",
            exists=True,
        ),
        # Docking parameters
        samples: int = typer.Option(
            10,
            "--samples",
            "-n",
            help="Number of docking poses to generate (1-100)",
            min=1,
            max=100,
        ),
        inference_steps: int = typer.Option(
            20,
            "--inference-steps",
            "-i",
            help="Number of denoising steps (5-100)",
            min=5,
            max=100,
        ),
        # 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)",
        ),
    ):
        """
        Perform molecular docking using DiffDock's diffusion-based model.

        Predicts binding poses between proteins and small molecule ligands
        with confidence scores. Accepts protein as PDB file or sequence,
        ligand as SDF file or SMILES string.

        Examples:
            # Dock with SMILES ligand and protein PDB
            amina run diffdock --protein-pdb ./protein.pdb --ligand-smiles "CCO" -o ./output/

            # Dock with protein sequence (uses ESMFold) and SMILES
            amina run diffdock -ps "MKFLILLFNILCLFPVLAADNH" -ls "CC(=O)OC1=CC=CC=C1C(=O)O" -o ./output/

            # Dock with ligand SDF file
            amina run diffdock --protein-pdb ./protein.pdb --ligand-sdf ./ligand.sdf -n 20 -o ./output/

            # Custom job name and more poses
            amina run diffdock -pp ./protein.pdb -ls "CCO" -n 50 -j mydock -o ./output/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

        # Validate protein input (exactly one required)
        if not protein_sequence and not protein_pdb:
            console.print("[red]Error:[/red] Provide either --protein-sequence or --protein-pdb")
            raise typer.Exit(1)

        if protein_sequence and protein_pdb:
            console.print("[red]Error:[/red] Provide only one of --protein-sequence or --protein-pdb")
            raise typer.Exit(1)

        # Validate ligand input (exactly one required)
        if not ligand_smiles and not ligand_sdf:
            console.print("[red]Error:[/red] Provide either --ligand-smiles or --ligand-sdf")
            raise typer.Exit(1)

        if ligand_smiles and ligand_sdf:
            console.print("[red]Error:[/red] Provide only one of --ligand-smiles or --ligand-sdf")
            raise typer.Exit(1)

        # Build params
        params = {
            "samples_per_complex": samples,
            "inference_steps": inference_steps,
        }

        # Add protein input
        if protein_sequence:
            params["protein_sequence"] = protein_sequence
            console.print(f"Protein: sequence ({len(protein_sequence)} residues)")
        elif protein_pdb:
            # Read PDB content
            pdb_content = protein_pdb.read_text()
            params["protein_pdb_content"] = pdb_content
            params["input_filename"] = protein_pdb.stem
            console.print(f"Protein: {protein_pdb.name}")

        # Add ligand input
        if ligand_smiles:
            params["ligand_smiles"] = ligand_smiles
            console.print(f"Ligand: SMILES ({ligand_smiles[:30]}{'...' if len(ligand_smiles) > 30 else ''})")
        elif ligand_sdf:
            # Read SDF content
            sdf_content = ligand_sdf.read_text()
            params["ligand_sdf_content"] = sdf_content
            console.print(f"Ligand: {ligand_sdf.name}")

        if job_name:
            params["job_name"] = job_name

        console.print(f"Docking parameters: {samples} poses, {inference_steps} inference steps")

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