"""PDB Chain Selection tool for the Amina CLI."""

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

METADATA = {
    "name": "pdb-chain-select",
    "display_name": "PDB Chain Select",
    "category": "utilities",
    "description": "Extract or remove specific protein chains from PDB files",
    "modal_function_name": "pdb_chain_select_worker",
    "modal_app_name": "pdb-chain-select-api",
    "status": "available",
    "outputs": {
        "supabase_filepath": "Processed PDB file with selected chains",
        "statistics_report_filepath": "CSV report with chain selection statistics",
    },
}

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-chain-select")
    def run_chain_select(
        pdb: Path = typer.Option(
            ...,
            "--pdb",
            "-p",
            help="Path to PDB file",
            exists=True,
        ),
        chains: str = typer.Option(
            ...,
            "--chains",
            "-c",
            help="Comma-separated chain IDs to KEEP (e.g., 'A,B'). Use '_' for blank chain.",
        ),
        invert: bool = typer.Option(
            False,
            "--invert/--keep",
            help="Invert selection: REMOVE listed chains instead of keeping them",
        ),
        renumber_atoms: bool = typer.Option(
            False,
            "--renumber-atoms/--keep-atoms",
            help="Renumber atom serials sequentially (1..N)",
        ),
        renumber_residues: bool = typer.Option(
            False,
            "--renumber-residues/--keep-residues",
            help="Renumber residue numbers per chain (1..N)",
        ),
        rebuild_conect: bool = typer.Option(
            False,
            "--rebuild-conect/--keep-conect",
            help="Rebuild CONECT records for kept atoms",
        ),
        preserve_cofactors: bool = typer.Option(
            False,
            "--preserve-cofactors/--remove-cofactors",
            help="Keep HETATM records even if associated chains are removed",
        ),
        altloc_policy: str = typer.Option(
            "keep-all",
            "--altloc",
            "-a",
            help="Alternative location policy: 'keep-all', 'highest-occupancy', or 'first'",
        ),
        validate: bool = typer.Option(
            True,
            "--validate/--no-validate",
            help="Perform structure validation checks",
        ),
        warn_insertions: bool = typer.Option(
            True,
            "--warn-insertions/--no-warn-insertions",
            help="Warn about insertion codes that may be lost during renumbering",
        ),
        quiet: bool = typer.Option(
            False,
            "--quiet/--verbose",
            "-q",
            help="Suppress warnings and statistics output",
        ),
        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)",
        ),
    ):
        """
        Extract or remove specific protein chains from PDB files.

        Supports chain extraction/removal, atom/residue renumbering, alternative
        location handling, and cofactor preservation for protein engineering workflows.

        Examples:
            # Keep only chains A and B
            amina run pdb-chain-select --pdb ./1A2J.pdb --chains A,B -o ./output/

            # Remove chain C (invert selection)
            amina run pdb-chain-select --pdb ./1A2J.pdb --chains C --invert -o ./output/

            # Keep chain A with residue renumbering
            amina run pdb-chain-select --pdb ./1A2J.pdb --chains A --renumber-residues -o ./output/

            # Keep chain A, use highest occupancy for alt locs
            amina run pdb-chain-select --pdb ./1A2J.pdb --chains A --altloc highest-occupancy -o ./output/

            # Custom job name
            amina run pdb-chain-select --pdb ./1A2J.pdb --chains A -j myjob -o ./output/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

        if not chains or not chains.strip():
            console.print("[red]Error:[/red] --chains / -c is required")
            raise typer.Exit(1)

        # Validate altloc_policy
        valid_policies = ["keep-all", "highest-occupancy", "first"]
        if altloc_policy not in valid_policies:
            console.print(f"[red]Error:[/red] --altloc must be one of: {', '.join(valid_policies)}")
            raise typer.Exit(1)

        # Read PDB file content
        pdb_content = pdb.read_text()

        # Build params dict matching worker's expected fields
        params = {
            "pdb_content": pdb_content,
            "pdb_filename": pdb.name,  # e.g., "1A2J.pdb"
            "chain_config": {
                "chains": chains,
                "invert": invert,
                "renumber_atom_serials": renumber_atoms,
                "renumber_residues": renumber_residues,
                "rebuild_conect": rebuild_conect,
                "preserve_cofactors": preserve_cofactors,
                "altloc_policy": altloc_policy,
                "validate_structure": validate,
                "warn_insertions": warn_insertions,
                "quiet": quiet,
            },
        }

        # Pass job_name if provided
        if job_name:
            params["job_name"] = job_name

        # Execute (handles spinner, errors, downloads)
        run_tool_with_progress("pdb-chain-select", params, output)
