"""Protein Relaxer tool for the Amina CLI."""

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

METADATA = {
    "name": "protein-relaxer",  # Use hyphens to match CLI command name convention
    "display_name": "Protein Relaxer",
    "category": "utilities",
    "description": "Clean PDB structures and minimize steric clashes via energy minimization",
    "modal_function_name": "protein_relaxer_worker",
    "modal_app_name": "protein_relaxer_api",
    "status": "available",
    "outputs": {
        "pdb_filepath": "Relaxed PDB structure after energy minimization",
        "cleaning_report_filepath": "CSV report with cleaning details",
    },
}

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("protein-relaxer")
    def run_protein_relaxer(
        pdb: Path = typer.Option(
            ...,
            "--pdb",
            "-p",
            help="Path to PDB file to relax",
            exists=True,
        ),
        temperature: float = typer.Option(
            300.0,
            "--temperature",
            "-t",
            help="Simulation temperature in Kelvin (default: 300.0)",
        ),
        nonbonded_cutoff: float = typer.Option(
            1.0,
            "--cutoff",
            help="Nonbonded interaction cutoff in nm (default: 1.0)",
        ),
        padding: float = typer.Option(
            1.0,
            "--padding",
            help="Solvent shell thickness in nm (default: 1.0)",
        ),
        ionic_strength: float = typer.Option(
            0.15,
            "--ionic-strength",
            help="NaCl concentration in M (default: 0.15)",
        ),
        max_iterations: int = typer.Option(
            0,
            "--max-iterations",
            help="Max minimization iterations, 0 for unlimited (default: 0)",
        ),
        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: uses input filename)",
        ),
    ):
        """
        Clean PDB structures and minimize steric clashes through energy minimization.

        Uses OpenMM to perform energy minimization on protein structures.
        Removes heterogens (ligands), adds missing atoms and hydrogens,
        and performs constrained energy minimization.

        Examples:
            amina run protein-relaxer --pdb ./1A2J.pdb -o ./output/
            amina run protein-relaxer --pdb ./structure.pdb --temperature 310 -o ./output/
            amina run protein-relaxer --pdb ./1A2J.pdb --max-iterations 1000 -j myjob -o ./output/
        """
        # 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()

        # Build relaxer_config dict
        relaxer_config = {
            "temperature": temperature,
            "nonbonded_cutoff": nonbonded_cutoff,
            "padding": padding,
            "ionic_strength": ionic_strength,
            "max_iterations": max_iterations,
            # Use defaults for other parameters
            "friction": 1.0,
            "timestep": 0.002,
            "platform": "CUDA",
            "precision": "mixed",
            "device_index": "0",
            "forcefield_files": ["amber14-all.xml", "amber14/tip3pfb.xml"],
            "only_protein": True,
        }

        params = {
            "pdb_content": pdb_content,
            "pdb_filename": pdb.name,
            "relaxer_config": relaxer_config,
        }

        if job_name:
            params["job_name"] = job_name

        run_tool_with_progress("protein_relaxer", params, output)
