"""RFDiffusion protein structure generation."""

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

METADATA = {
    "name": "rfdiffusion",
    "display_name": "RFDiffusion",
    "category": "design",
    "description": "Protein structure generation with 6 modes: unconditional, binder-design, binder-redesign, motif-scaffolding, partial-diffusion, custom-contigs",
    "modal_function_name": None,  # Core handles routing based on mode
    "modal_app_name": "rfdiffusion-api",
    "status": "available",
    "outputs": {
        "pdb_filepath": "Generated protein structure(s)",
        "metadata_filepath": "RFDiffusion metadata (.trb)",
    },
}

console = Console()

# Valid modes (CLI uses kebab-case)
VALID_MODES = [
    "unconditional",
    "binder-design",
    "binder-redesign",
    "motif-scaffolding",
    "partial-diffusion",
    "custom-contigs",
]


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("rfdiffusion")
    def run_rfdiffusion(
        # Required
        mode: str = typer.Option(
            ...,
            "--mode",
            "-m",
            help=f"Generation mode: {', '.join(VALID_MODES)}",
        ),
        output: Path = typer.Option(
            ...,
            "--output",
            "-o",
            help="Output directory for results",
        ),
        # Common options
        num_designs: int = typer.Option(
            10,
            "--num-designs",
            "-n",
            help="Number of designs to generate (1-100)",
        ),
        job_name: Optional[str] = typer.Option(
            None,
            "--job-name",
            "-j",
            help="Custom job name",
        ),
        # Input file (for modes that need it)
        input_pdb: Optional[Path] = typer.Option(
            None,
            "--input",
            "-i",
            help="Input PDB file (for binder/motif/partial modes)",
            exists=True,
        ),
        # Unconditional options
        length_range: Optional[str] = typer.Option(
            None,
            "--length",
            "-l",
            help="Length range, e.g., '100-200' (unconditional mode)",
        ),
        symmetry: Optional[str] = typer.Option(
            None,
            "--symmetry",
            "-sym",
            help="Symmetry type, e.g., 'c4', 'd2' (unconditional/motif modes)",
        ),
        # Binder design options
        hotspots: Optional[str] = typer.Option(
            None,
            "--hotspots",
            "-h",
            help="Hotspot residues, comma-separated, e.g., 'A30,A33' (binder-design)",
        ),
        binder_length: Optional[str] = typer.Option(
            None,
            "--binder-length",
            "-bl",
            help="Binder length range, e.g., '80-120' (binder-design)",
        ),
        use_beta: bool = typer.Option(
            False,
            "--beta",
            help="Use beta model for diverse topologies (binder-design)",
        ),
        use_fold_conditioning: bool = typer.Option(
            False,
            "--fold-conditioning",
            help="Enable fold conditioning (binder-design)",
        ),
        # Binder redesign options
        binder_chain: Optional[str] = typer.Option(
            None,
            "--binder-chain",
            "-bc",
            help="Binder chain ID (binder-redesign)",
        ),
        target_chain: Optional[str] = typer.Option(
            None,
            "--target-chain",
            "-tc",
            help="Target chain ID (binder-redesign)",
        ),
        # Partial diffusion / redesign options
        noise_level: Optional[int] = typer.Option(
            None,
            "--noise",
            help="Noise level 1-50 (partial-diffusion, binder-redesign)",
        ),
        preserve: Optional[str] = typer.Option(
            None,
            "--preserve",
            "-p",
            help="Residues to preserve, comma-separated (partial-diffusion, binder-redesign)",
        ),
        # Motif scaffolding options
        contigs: Optional[str] = typer.Option(
            None,
            "--contigs",
            help="Contig specification, e.g., '10-40/A163-181/10-40' (motif-scaffolding, custom-contigs)",
        ),
        inpaint_seq: Optional[str] = typer.Option(
            None,
            "--inpaint-seq",
            help="Residues for sequence inpainting, e.g., 'A1,A30-40' (motif-scaffolding)",
        ),
        provide_seq: Optional[str] = typer.Option(
            None,
            "--provide-seq",
            help="Residue ranges with known sequence, e.g., '100-119' (motif-scaffolding)",
        ),
        active_site: bool = typer.Option(
            False,
            "--active-site",
            help="Use active site model for small motifs (motif-scaffolding)",
        ),
        # Custom contigs options
        custom_model: Optional[str] = typer.Option(
            None,
            "--model",
            help="Custom model checkpoint (custom-contigs)",
        ),
        advanced_config: Optional[str] = typer.Option(
            None,
            "--advanced-config",
            help="Raw Hydra overrides as JSON string (custom-contigs)",
        ),
    ):
        """
        Generate protein structures using RFDiffusion.

        6 modes available:

        - unconditional: Generate novel proteins from scratch

        - binder-design: Design binders for target hotspots

        - binder-redesign: Redesign existing binders

        - motif-scaffolding: Scaffold around fixed motifs

        - partial-diffusion: Diversify existing structures

        - custom-contigs: Advanced custom contig control

        Examples:

            amina run rfdiffusion -m unconditional --length 100-150 -n 5 -o ./out/

            amina run rfdiffusion -m binder-design -i target.pdb --hotspots A30,A33 --binder-length 80-120 -o ./out/

            amina run rfdiffusion -m partial-diffusion -i protein.pdb --noise 25 -o ./out/

            amina run rfdiffusion -m motif-scaffolding -i motif.pdb --contigs "10-40/A163-181/10-40" -o ./out/
        """
        # Validate mode
        if mode not in VALID_MODES:
            console.print(f"[red]Invalid mode: {mode}[/red]")
            console.print(f"Valid modes: {', '.join(VALID_MODES)}")
            raise typer.Exit(1)

        # Map CLI mode (kebab-case) to backend mode (snake_case)
        mode_map = {
            "unconditional": "unconditional_generation",
            "binder-design": "binder_design",
            "binder-redesign": "binder_redesign",
            "motif-scaffolding": "motif_scaffolding",
            "partial-diffusion": "partial_diffusion",
            "custom-contigs": "custom_contigs",
        }
        mode_value = mode_map[mode]

        # Build params based on mode
        params = {"mode": mode_value, "num_designs": num_designs}
        if job_name:
            params["job_name"] = job_name

        # Read PDB content if provided
        pdb_content = input_pdb.read_text() if input_pdb else None

        # Mode-specific validation and param building
        if mode == "unconditional":
            if not length_range:
                length_range = "150-150"
            params["length_range"] = length_range
            if symmetry:
                params["symmetry_type"] = symmetry

        elif mode == "binder-design":
            if not input_pdb:
                console.print("[red]--input/-i required for binder-design[/red]")
                raise typer.Exit(1)
            if not hotspots:
                console.print("[red]--hotspots/-h required for binder-design[/red]")
                raise typer.Exit(1)
            params["pdb_content"] = pdb_content
            params["hotspot_residues"] = [h.strip() for h in hotspots.split(",")]
            params["binder_length_range"] = binder_length or "100-100"
            params["use_beta_model"] = use_beta
            if use_fold_conditioning:
                params["fold_conditioning"] = {"use_fold_conditioning": True}

        elif mode == "binder-redesign":
            if not input_pdb:
                console.print("[red]--input/-i required for binder-redesign[/red]")
                raise typer.Exit(1)
            if not noise_level:
                console.print("[red]--noise required for binder-redesign[/red]")
                raise typer.Exit(1)
            params["pdb_content"] = pdb_content
            params["binder_chain"] = binder_chain or "A"
            params["target_chain"] = target_chain or "B"
            params["noise_level"] = noise_level
            if preserve:
                params["preserve_sequences"] = [p.strip() for p in preserve.split(",")]

        elif mode == "motif-scaffolding":
            if not input_pdb:
                console.print("[red]--input/-i required for motif-scaffolding[/red]")
                raise typer.Exit(1)
            if not contigs:
                console.print("[red]--contigs required for motif-scaffolding[/red]")
                raise typer.Exit(1)
            params["pdb_content"] = pdb_content
            # Build ContigSpecification
            contig_spec = {"contigs": [contigs]}
            if inpaint_seq:
                contig_spec["inpaint_seq"] = [s.strip() for s in inpaint_seq.split(",")]
            if provide_seq:
                contig_spec["provide_seq"] = [s.strip() for s in provide_seq.split(",")]
            params["contigs"] = contig_spec
            params["use_active_site_model"] = active_site
            if symmetry:
                params["symmetry_type"] = symmetry

        elif mode == "partial-diffusion":
            if not input_pdb:
                console.print("[red]--input/-i required for partial-diffusion[/red]")
                raise typer.Exit(1)
            if not noise_level:
                console.print("[red]--noise required for partial-diffusion[/red]")
                raise typer.Exit(1)
            params["pdb_content"] = pdb_content
            params["noise_level"] = noise_level
            if preserve:
                params["preserve_sequences"] = [p.strip() for p in preserve.split(",")]

        elif mode == "custom-contigs":
            if not contigs:
                console.print("[red]--contigs required for custom-contigs[/red]")
                raise typer.Exit(1)
            params["contig_string"] = contigs
            if pdb_content:
                params["pdb_content"] = pdb_content
            if custom_model:
                params["custom_model"] = custom_model
            if advanced_config:
                try:
                    params["advanced_config"] = json.loads(advanced_config)
                except json.JSONDecodeError as e:
                    console.print(f"[red]Invalid JSON in --advanced-config: {e}[/red]")
                    raise typer.Exit(1)

        # Run the tool
        run_tool_with_progress("rfdiffusion", params, output)
