"""Boltz-2 structure prediction tool for the Amina CLI.

Boltz-2 is a diffusion-based model for predicting 3D structures of biomolecular complexes.
It supports proteins, DNA, RNA, and small molecule ligands.
"""

import typer
from pathlib import Path
from typing import List, Optional, Set, Tuple
from rich.console import Console


def parse_sequence_with_chain_id(value: str, used_chains: Set[str]) -> Tuple[Optional[str], str]:
    """
    Parse 'X:SEQUENCE' or 'SEQUENCE' format.

    Supports explicit chain ID assignment via the X:SEQ format:
    - "A:MVLSPAD..." -> chain_id="A", sequence="MVLSPAD..."
    - "MVLSPAD..." -> chain_id=None (auto-assign), sequence="MVLSPAD..."

    Args:
        value: Input string, either "CHAIN_ID:SEQUENCE" or just "SEQUENCE"
        used_chains: Set of already-used chain IDs (for duplicate detection)

    Returns:
        Tuple of (chain_id or None, sequence)

    Raises:
        typer.Exit: If duplicate chain ID is detected
    """
    # Check for X:SEQ format (chain ID is 1-2 characters before colon)
    if ":" in value:
        parts = value.split(":", 1)
        potential_chain = parts[0].strip()
        # Only treat as chain ID if it's 1-2 characters (not a long prefix)
        if len(potential_chain) <= 2 and len(potential_chain) >= 1:
            chain_id = potential_chain.upper()
            if chain_id in used_chains:
                Console().print(f"[red]Error:[/red] Duplicate chain ID: {chain_id}")
                raise typer.Exit(1)
            used_chains.add(chain_id)
            sequence = parts[1].strip().upper()
            return chain_id, sequence
    # No chain ID specified - auto-assign
    return None, value.strip().upper()


METADATA = {
    "name": "boltz2",
    "display_name": "Boltz-2",
    "category": "folding",
    "description": "Predict 3D structures of biomolecular complexes using Boltz-2",
    "modal_function_name": "boltz2_worker",
    "modal_app_name": "boltz2-api",
    "status": "available",
    "outputs": {
        "pdb_filepath": "Predicted structure in PDB format",
        "json_filepath": "Confidence scores (pLDDT, pTM, ipTM)",
        "yaml_filepath": "Boltz-2 input configuration",
        "affinity_filepath": "Binding affinity predictions (if enabled)",
        "pae_filepath": "Predicted Aligned Error matrix",
    },
    "output_display": {
        "data_path": "predictions[0]",
        "sections": [
            {
                "title": "Confidence Metrics",
                "fields": [
                    {"key": "avg_plddt", "label": "pLDDT"},
                    {"key": "ptm", "label": "pTM"},
                    {"key": "iptm", "label": "ipTM"},
                    {"key": "confidence_score", "label": "Confidence Score"},
                ],
            },
            {
                "title": "Interface Metrics",
                "fields": [
                    {"key": "complex_iplddt", "label": "Complex ipLDDT"},
                    {"key": "protein_iptm", "label": "Protein ipTM"},
                    {"key": "ligand_iptm", "label": "Ligand ipTM"},
                ],
            },
            {
                "title": "Error Metrics",
                "fields": [
                    {"key": "min_pae", "label": "Min PAE", "format": "{:.2f} \u00c5"},
                    {"key": "ipsae", "label": "ipSAE"},
                    {"key": "complex_pde", "label": "Complex PDE"},
                    {"key": "complex_ipde", "label": "Complex iPDE"},
                ],
            },
            {
                "title": "Docking Quality Scores",
                "fields": [
                    {"key": "lis", "label": "LIS"},
                    {"key": "pdockq", "label": "pDockQ"},
                    {"key": "pdockq2", "label": "pDockQ2"},
                ],
            },
            {
                "title": "Per-Chain-Pair Metrics",
                "fields": [
                    {"key": "ipsae_per_chain_pair", "label": "ipSAE per chain pair"},
                    {"key": "lis_per_chain_pair", "label": "LIS per chain pair"},
                    {"key": "pdockq_per_chain_pair", "label": "pDockQ per chain pair"},
                    {"key": "pdockq2_per_chain_pair", "label": "pDockQ2 per chain pair"},
                ],
            },
            {
                "title": "Binding Affinity Prediction",
                "fields": [
                    {"key": "affinity_pred_value", "label": "Predicted pKd"},
                    {"key": "affinity_probability_binary", "label": "Binding Probability"},
                ],
            },
        ],
    },
}

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("boltz2")
    def run_boltz2(
        sequences: List[str] = typer.Option(
            [],
            "--sequence",
            "-s",
            help="Protein sequence. Prefix with X: to set chain ID (e.g., 'A:MVLSPAD'), or omit for auto (A,B,C...). Repeatable.",
        ),
        fasta_files: List[Path] = typer.Option(
            [],
            "--fasta",
            "-f",
            help="Path to FASTA file with single protein sequence (repeatable)",
        ),
        yaml_file: Optional[Path] = typer.Option(
            None,
            "--yaml",
            "-y",
            help="Path to Boltz-2 YAML input file (for complex inputs)",
            exists=True,
        ),
        ligands: List[str] = typer.Option(
            [],
            "--ligand",
            "-l",
            help="Ligand SMILES. Prefix with X: to set chain ID (e.g., 'L:CCO'), or omit for auto. Repeatable.",
        ),
        ligand_ccds: List[str] = typer.Option(
            [],
            "--ligand-ccd",
            help="Ligand CCD code (e.g., ATP, NAD). Prefix with X: for chain ID (e.g., 'L:ATP'), or omit for auto. Repeatable.",
        ),
        dna_sequences: List[str] = typer.Option(
            [],
            "--dna",
            help="DNA sequence. Prefix with X: to set chain ID (e.g., 'D:ATGC'), or omit for auto. Repeatable.",
        ),
        rna_sequences: List[str] = typer.Option(
            [],
            "--rna",
            help="RNA sequence. Prefix with X: to set chain ID (e.g., 'R:AUGC'), or omit for auto. Repeatable.",
        ),
        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)",
        ),
        recycling_steps: int = typer.Option(
            3,
            "--recycling-steps",
            "-r",
            help="Number of recycling iterations (1-20)",
            min=1,
            max=20,
        ),
        sampling_steps: int = typer.Option(
            200,
            "--sampling-steps",
            help="Number of diffusion sampling steps (1-1000)",
            min=1,
            max=1000,
        ),
        diffusion_samples: int = typer.Option(
            1,
            "--samples",
            "-n",
            help="Number of independent structure samples (1-50)",
            min=1,
            max=50,
        ),
        enable_affinity: bool = typer.Option(
            False,
            "--affinity",
            help="Enable binding affinity prediction (protein-ligand complexes only)",
        ),
        background: bool = typer.Option(
            False,
            "--background",
            "-b",
            help="Submit job and return immediately without waiting for completion",
        ),
    ):
        """
        Predict 3D structures of biomolecular complexes using Boltz-2.

        Entity flags (--sequence, --fasta, --ligand, etc.) are repeatable to build
        multi-chain complexes.

        Chain IDs can be specified using the X:SEQ format (e.g., "A:MVLSPAD...").
        If not specified, chain IDs are auto-assigned (A, B, C, ...).

        Examples:
            # Single protein:
            amina run boltz2 -s "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSH" -o ./results/

            # With explicit chain IDs (X:SEQ format):
            amina run boltz2 -s "A:TARGETSEQUENCE" -s "X:BINDERSEQUENCE" -o ./results/

            # Protein from FASTA file:
            amina run boltz2 -f ./protein.fasta -o ./results/

            # Protein-ligand complex with affinity prediction:
            amina run boltz2 -s "MKFLIL" -l "CCO" --affinity -o ./results/

            # Multi-chain protein complex:
            amina run boltz2 -s "MKFLIL" -s "PNNTHEQHLRK" -o ./results/

            # Two proteins + ligand + affinity:
            amina run boltz2 -s "MKFLIL" -s "PNNTHEQHLRK" -l "CCO" --affinity -o ./results/

            # Using YAML for complex inputs (DNA, RNA, custom chain IDs):
            amina run boltz2 --yaml ./complex.yaml -r 5 --samples 3 -o ./results/
        """
        # Validate required options
        if output is None:
            console.print("[red]Error:[/red] --output / -o is required")
            raise typer.Exit(1)

        # Helper function for auto chain ID assignment
        def _auto_chain_id(used: set) -> str:
            """Get next available chain ID (A, B, C, ...)."""
            for i in range(26):
                chain = chr(65 + i)
                if chain not in used:
                    used.add(chain)
                    return chain
            raise ValueError("Exceeded 26 chains")

        # Count total entities from flags
        total_entities = (
            len(sequences)
            + len(fasta_files)
            + len(ligands)
            + len(ligand_ccds)
            + len(dna_sequences)
            + len(rna_sequences)
        )

        # Validate input: need either YAML file or at least one entity
        if not yaml_file and total_entities == 0:
            console.print(
                "[red]Error:[/red] Provide at least one entity "
                "(--sequence, --fasta, --ligand, --ligand-ccd, --dna, --rna) or use --yaml"
            )
            raise typer.Exit(1)

        # Build YAML content
        yaml_content = None

        if yaml_file:
            # Load YAML directly (ignores other entity flags if YAML provided)
            import yaml

            try:
                with open(yaml_file) as f:
                    yaml_content = yaml.safe_load(f)
                console.print(f"Loaded complex definition from {yaml_file}")
            except Exception as e:
                console.print(f"[red]Error:[/red] Failed to parse YAML file: {e}")
                raise typer.Exit(1)
        else:
            # Build YAML from entity flags
            used_chains: set = set()
            sequences_list = []

            # Add proteins from --sequence flags (supports X:SEQ format)
            for seq in sequences:
                parsed_chain_id, sequence = parse_sequence_with_chain_id(seq, used_chains)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                sequences_list.append({"protein": {"id": chain_id, "sequence": sequence}})
                console.print(f"Added protein chain {chain_id} ({len(sequence)} residues)")

            # Add proteins from --fasta flags (auto-assign chain ID)
            for fasta_path in fasta_files:
                if not fasta_path.exists():
                    console.print(f"[red]Error:[/red] FASTA file not found: {fasta_path}")
                    raise typer.Exit(1)
                content = fasta_path.read_text()
                seq = "".join(
                    line.strip() for line in content.splitlines() if line.strip() and not line.startswith(">")
                )
                chain_id = _auto_chain_id(used_chains)
                sequences_list.append({"protein": {"id": chain_id, "sequence": seq.upper()}})
                console.print(f"Added protein chain {chain_id} from {fasta_path.name} ({len(seq)} residues)")

            # Track ligand chain IDs for affinity prediction
            ligand_chain_ids: list = []

            # Add ligands from --ligand flags (SMILES, supports X:SMILES format)
            for smiles in ligands:
                parsed_chain_id, smiles_clean = parse_sequence_with_chain_id(smiles, used_chains)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                ligand_chain_ids.append(chain_id)
                sequences_list.append({"ligand": {"id": chain_id, "smiles": smiles_clean}})
                smiles_preview = smiles_clean[:40] + "..." if len(smiles_clean) > 40 else smiles_clean
                console.print(f"Added ligand chain {chain_id} (SMILES: {smiles_preview})")

            # Add ligands from --ligand-ccd flags (supports X:CCD format)
            for ccd in ligand_ccds:
                parsed_chain_id, ccd_clean = parse_sequence_with_chain_id(ccd, used_chains)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                ligand_chain_ids.append(chain_id)
                sequences_list.append({"ligand": {"id": chain_id, "ccd": ccd_clean.upper()}})
                console.print(f"Added ligand chain {chain_id} (CCD: {ccd_clean.upper()})")

            # Add DNA from --dna flags (supports X:SEQ format)
            for dna in dna_sequences:
                parsed_chain_id, sequence = parse_sequence_with_chain_id(dna, used_chains)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                sequences_list.append({"dna": {"id": chain_id, "sequence": sequence}})
                console.print(f"Added DNA chain {chain_id} ({len(sequence)} bases)")

            # Add RNA from --rna flags (supports X:SEQ format)
            for rna in rna_sequences:
                parsed_chain_id, sequence = parse_sequence_with_chain_id(rna, used_chains)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                sequences_list.append({"rna": {"id": chain_id, "sequence": sequence}})
                console.print(f"Added RNA chain {chain_id} ({len(sequence)} bases)")

            yaml_content = {"version": 1, "sequences": sequences_list}

            # Add properties section for affinity prediction
            if enable_affinity:
                if not ligand_chain_ids:
                    console.print(
                        "[red]Error:[/red] --affinity requires at least one ligand (use --ligand or --ligand-ccd)"
                    )
                    raise typer.Exit(1)
                # Use first ligand as binder
                binder_chain = ligand_chain_ids[0]
                yaml_content["properties"] = [{"affinity": {"binder": binder_chain}}]
                console.print(f"Enabled affinity prediction (binder: chain {binder_chain})")

        # Build params for gateway
        params = {
            "yaml_content": yaml_content,
            "recycling_steps": recycling_steps,
            "sampling_steps": sampling_steps,
            "diffusion_samples": diffusion_samples,
            "step_scale": 1.638,
            "use_potentials": False,
            "enable_affinity_prediction": enable_affinity,
        }

        if job_name:
            params["job_name"] = job_name

        # Execute
        run_tool_with_progress("boltz2", params, output, background=background)
