"""PDB B-factor Overwrite tool for the Amina CLI."""

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

METADATA = {
    "name": "pdb-bfactor-overwrite",
    "display_name": "PDB B-factor Overwrite",
    "category": "utilities",
    "description": "Overwrite B-factor values in PDB files with custom scores",
    "modal_function_name": "pdb_bfactor_overwrite_worker",
    "modal_app_name": "pdb-bfactor-overwrite-api",
    "status": "available",
    "outputs": {
        "pdb_filepath": "PDB file with overwritten B-factor values",
    },
}

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-bfactor-overwrite")
    def run_pdb_bfactor_overwrite(
        pdb: Path = typer.Option(
            ...,
            "--pdb",
            "-p",
            help="Path to PDB file to modify",
            exists=True,
        ),
        bfactor_json: Path = typer.Option(
            ...,
            "--bfactor-json",
            "-b",
            help="JSON file with B-factor mappings (e.g., {'A:10': 85.5, 'A:11': 92.3})",
            exists=True,
        ),
        description: str = typer.Option(
            "modified",
            "--description",
            "-d",
            help="Description suffix for output filename (e.g., 'plddt', 'confidence')",
        ),
        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)",
        ),
    ):
        """
        Overwrite B-factor values in PDB files with custom scores.

        Useful for visualizing residue-level metrics (pLDDT, confidence, binding scores)
        in molecular visualization software like PyMOL or ChimeraX.

        The B-factor JSON file should map residue identifiers to new values:
        - Integer keys: {"10": 85.5} - applies to residue 10 on all chains
        - Chain-specific: {"A:10": 85.5, "B:10": 72.3} - chain-specific values

        Examples:
            amina run pdb-bfactor-overwrite --pdb ./structure.pdb --bfactor-json ./scores.json -o ./output/
            amina run pdb-bfactor-overwrite --pdb ./structure.pdb --bfactor-json ./plddt.json -d plddt -o ./output/
            amina run pdb-bfactor-overwrite --pdb ./protein.pdb --bfactor-json ./confidence.json -j myjob -o ./output/

        B-factor JSON format:
            {
                "A:10": 95.5,
                "A:11": 87.2,
                "A:12": 72.8,
                "B:10": 68.3
            }
        """
        # 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()

        # Parse B-factor JSON file
        try:
            bfactor_text = bfactor_json.read_text()
            bfactor_dict = json.loads(bfactor_text)

            if not isinstance(bfactor_dict, dict):
                console.print("[red]Error:[/red] B-factor JSON must be a dictionary")
                raise typer.Exit(1)

            if not bfactor_dict:
                console.print("[red]Error:[/red] B-factor dictionary cannot be empty")
                raise typer.Exit(1)

            # Validate all values are numeric
            for key, value in bfactor_dict.items():
                if not isinstance(value, (int, float)):
                    console.print(
                        f"[red]Error:[/red] B-factor value for '{key}' must be numeric, got {type(value).__name__}"
                    )
                    raise typer.Exit(1)

        except json.JSONDecodeError as e:
            console.print(f"[red]Error:[/red] Invalid JSON in B-factor file: {e}")
            raise typer.Exit(1)

        params = {
            "pdb_content": pdb_content,
            "pdb_filename": pdb.name,  # Include extension for output naming
            "bfactor_dict": bfactor_dict,
            "bfactor_description": description,
        }

        if job_name:
            params["job_name"] = job_name

        run_tool_with_progress("pdb-bfactor-overwrite", params, output)
