"""Protenix structure prediction tool for the Amina CLI."""

import json
import typer
from pathlib import Path
from typing import List, Optional, Set, Tuple, Union
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()


def _auto_chain_id(used_chains: Set[str]) -> str:
    """Get next available chain ID (A, B, C, ...)."""
    single_char_ids = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
    for char in single_char_ids:
        if char not in used_chains:
            used_chains.add(char)
            return char
    raise ValueError("Exceeded 62 chains (A-Z, a-z, 0-9)")


METADATA = {
    "name": "protenix",
    "display_name": "Protenix",
    "category": "folding",
    "description": "Predict protein structures using Protenix (AlphaFold 3 implementation)",
    "modal_function_name": "protenix_worker",
    "modal_app_name": "protenix-api",
    "status": "available",
    "outputs": {
        "pdb_filepath": "Predicted structure in PDB format",
        "cif_filepath": "Predicted structure in mmCIF format",
        "confidence_filepath": "Confidence scores JSON",
        "pae_filepath": "PAE matrix JSON",
    },
}

console = Console()


def build_protenix_json(
    sequences: List[str],
    fasta_files: List[Path],
    ligands: List[str],
    ligand_ccds: List[str],
    dna_sequences: List[str],
    rna_sequences: List[str],
    ions: List[str],
    count: int,
    job_name: Optional[str],
) -> List[dict]:
    """
    Build Protenix JSON from CLI flags.

    Supports X:SEQ format for explicit chain ID assignment:
    - --sequence A:MVLSPAD... -> chain A
    - --sequence MVLSPAD... -> auto-assigned chain

    Args:
        sequences: List of protein sequences (may include X:SEQ format)
        fasta_files: List of FASTA file paths
        ligands: List of ligand SMILES (may include X:SMILES format)
        ligand_ccds: List of CCD codes (may include X:CCD format)
        dna_sequences: List of DNA sequences (may include X:SEQ format)
        rna_sequences: List of RNA sequences (may include X:SEQ format)
        ions: List of ion codes (may include X:ION format)
        count: Copy count for all entities
        job_name: Job name for output

    Returns:
        Protenix JSON content
    """
    sequences_list = []
    used_chains: Set[str] = set()

    # Proteins from --sequence (supports X:SEQ format)
    for seq in sequences:
        chain_id, sequence = parse_sequence_with_chain_id(seq, used_chains)
        entity = {"sequence": sequence, "count": count}
        if chain_id:
            entity["chain_id"] = chain_id
        sequences_list.append({"proteinChain": entity})

    # Proteins from --fasta (auto-assign chain ID)
    for fasta_path in fasta_files:
        content = fasta_path.read_text()
        seq = "".join(line.strip() for line in content.splitlines() if line.strip() and not line.startswith(">"))
        entity = {"sequence": seq.upper(), "count": count}
        sequences_list.append({"proteinChain": entity})

    # DNA from --dna (supports X:SEQ format)
    for dna in dna_sequences:
        chain_id, sequence = parse_sequence_with_chain_id(dna, used_chains)
        entity = {"sequence": sequence, "count": count}
        if chain_id:
            entity["chain_id"] = chain_id
        sequences_list.append({"dnaSequence": entity})

    # RNA from --rna (supports X:SEQ format)
    for rna in rna_sequences:
        chain_id, sequence = parse_sequence_with_chain_id(rna, used_chains)
        entity = {"sequence": sequence, "count": count}
        if chain_id:
            entity["chain_id"] = chain_id
        sequences_list.append({"rnaSequence": entity})

    # Ligands from --ligand (SMILES, supports X:SMILES format)
    for smiles in ligands:
        chain_id, smiles_clean = parse_sequence_with_chain_id(smiles, used_chains)
        entity = {"ligand": smiles_clean, "count": count}
        if chain_id:
            entity["chain_id"] = chain_id
        sequences_list.append({"ligand": entity})

    # Ligands from --ligand-ccd (CCD codes, supports X:CCD format)
    for ccd in ligand_ccds:
        chain_id, ccd_raw = parse_sequence_with_chain_id(ccd, used_chains)
        ccd_clean = ccd_raw.strip().upper()
        ccd_value = f"CCD_{ccd_clean}" if not ccd_clean.startswith("CCD_") else ccd_clean
        entity = {"ligand": ccd_value, "count": count}
        if chain_id:
            entity["chain_id"] = chain_id
        sequences_list.append({"ligand": entity})

    # Ions from --ion (converted to CCD ligands - Protenix v1.0.0 removed "ion" type)
    for ion in ions:
        chain_id, ion_raw = parse_sequence_with_chain_id(ion, used_chains)
        ion_clean = ion_raw.strip().upper()
        entity = {"ligand": f"CCD_{ion_clean}", "count": count}
        if chain_id:
            entity["chain_id"] = chain_id
        sequences_list.append({"ligand": entity})

    return [{"name": job_name or "prediction", "sequences": sequences_list}]


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("protenix")
    def run_protenix(
        # Entity flags (repeatable)
        sequences: Optional[List[str]] = typer.Option(
            None,
            "--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: Optional[List[Path]] = typer.Option(
            None,
            "--fasta",
            "-f",
            help="FASTA file path (repeatable)",
        ),
        ligands: Optional[List[str]] = typer.Option(
            None,
            "--ligand",
            "-l",
            help="Ligand SMILES. Prefix with X: to set chain ID (e.g., 'L:CCO'), or omit for auto. Repeatable.",
        ),
        ligand_ccds: Optional[List[str]] = typer.Option(
            None,
            "--ligand-ccd",
            help="Ligand CCD code (e.g., ATP, MG). Prefix with X: for chain ID (e.g., 'L:ATP'), or omit for auto. Repeatable.",
        ),
        dna_sequences: Optional[List[str]] = typer.Option(
            None,
            "--dna",
            help="DNA sequence. Prefix with X: to set chain ID (e.g., 'D:ATGC'), or omit for auto. Repeatable.",
        ),
        rna_sequences: Optional[List[str]] = typer.Option(
            None,
            "--rna",
            help="RNA sequence. Prefix with X: to set chain ID (e.g., 'R:AUGC'), or omit for auto. Repeatable.",
        ),
        ions: Optional[List[str]] = typer.Option(
            None,
            "--ion",
            help="Ion code (e.g., ZN, MG). Prefix with X: for chain ID (e.g., 'Z:ZN'), or omit for auto. Converted to CCD ligands.",
        ),
        count: int = typer.Option(
            1,
            "--count",
            "-c",
            help="Copy count for all entities (default: 1, use for homodimers)",
            min=1,
        ),
        # JSON input (alternative to entity flags)
        json_file: Optional[Path] = typer.Option(
            None,
            "--json",
            "-j",
            help="Path to Protenix JSON input file (alternative to entity flags)",
            exists=True,
        ),
        # Output options
        output: Optional[Path] = typer.Option(
            None,
            "--output",
            "-o",
            help="Output directory for results (required unless --background)",
        ),
        job_name: Optional[str] = typer.Option(
            None,
            "--job-name",
            help="Custom job name for output files (default: random 4-letter code)",
        ),
        # Model options
        model: str = typer.Option(
            "protenix_base_20250630_v1.0.0",
            "--model",
            "-m",
            help="Model to use (protenix_base_20250630_v1.0.0 or protenix_base_default_v1.0.0)",
        ),
        # Template options (v1.0.0 feature)
        template_file: Optional[Path] = typer.Option(
            None,
            "--template",
            "-t",
            help="Template file (.a3m/.hhr) for template-based prediction",
            exists=True,
        ),
        sample: int = typer.Option(
            5,
            "--sample",
            "-n",
            help="Number of samples to generate (1-10)",
            min=1,
            max=10,
        ),
        seeds: str = typer.Option(
            "101",
            "--seeds",
            help="Random seed(s), comma-separated for multiple",
        ),
        output_format: str = typer.Option(
            "pdb",
            "--format",
            help="Output structure format: pdb or cif",
        ),
        background: bool = typer.Option(
            False,
            "--background",
            "-b",
            help="Submit job and return immediately without waiting for completion",
        ),
    ):
        """
        Predict protein structures using Protenix (AlphaFold 3 implementation).

        Supports multi-chain complexes, ligands, DNA/RNA, and covalent modifications.
        Use entity flags for simple cases or --json for complex inputs.

        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 protenix -s "MVLSPADKTNVKAAWGKVGAHAGEYGAEAL" -o ./results/

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

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

            # Homodimer (2 copies):
            amina run protenix -s "MKFLILLFNILCLFPVLAADNH" --count 2 -o ./results/

            # Protein-ligand (SMILES):
            amina run protenix -s "MKFLIL" -l "CCO" -o ./results/

            # Protein with ATP:
            amina run protenix -s "MKFLIL" --ligand-ccd ATP -o ./results/

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

            # Protein-DNA complex:
            amina run protenix -s "MKFLIL" --dna "ATGCATGCATGC" -o ./results/

            # Protein with metal ion (preferred - using CCD code):
            amina run protenix -s "MKFLIL" --ligand-ccd MG -o ./results/

            # Protein with zinc ion (shorthand):
            amina run protenix -s "MKFLIL" --ion ZN -o ./results/

            # Complex input (JSON):
            amina run protenix --json complex_input.json -o ./results/
        """
        # Normalize None to empty lists for entity flags
        sequences = sequences or []
        fasta_files = fasta_files or []
        ligands = ligands or []
        ligand_ccds = ligand_ccds or []
        dna_sequences = dna_sequences or []
        rna_sequences = rna_sequences or []
        ions = ions or []

        # Validate required options
        if output is None and not background:
            console.print("[red]Error:[/red] --output / -o is required (unless using --background)")
            raise typer.Exit(1)

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

        # Validate: need JSON or at least one entity
        if not json_file and total_entities == 0:
            console.print(
                "[red]Error:[/red] Provide at least one entity "
                "(--sequence, --fasta, --ligand, --dna, --rna, --ion) or use --json"
            )
            raise typer.Exit(1)

        # Validate output format
        if output_format.lower() not in ["pdb", "cif"]:
            console.print("[red]Error:[/red] --format must be 'pdb' or 'cif'")
            raise typer.Exit(1)

        # Build JSON content
        if json_file:
            # Existing path: load from file
            try:
                with open(json_file, "r") as f:
                    json_content = json.load(f)
                console.print(f"Loaded input from {json_file}")
            except json.JSONDecodeError as e:
                console.print(f"[red]Error:[/red] Invalid JSON file: {e}")
                raise typer.Exit(1)

            # Validate JSON structure
            if not isinstance(json_content, list):
                console.print("[red]Error:[/red] JSON must be a list of job definitions")
                raise typer.Exit(1)
        else:
            # New path: build from flags
            json_content = build_protenix_json(
                sequences, fasta_files, ligands, ligand_ccds, dna_sequences, rna_sequences, ions, count, job_name
            )
            entity_summary = []
            if sequences:
                entity_summary.append(f"{len(sequences)} protein(s)")
            if fasta_files:
                entity_summary.append(f"{len(fasta_files)} FASTA file(s)")
            if ligands:
                entity_summary.append(f"{len(ligands)} ligand(s)")
            if ligand_ccds:
                entity_summary.append(f"{len(ligand_ccds)} CCD ligand(s)")
            if dna_sequences:
                entity_summary.append(f"{len(dna_sequences)} DNA sequence(s)")
            if rna_sequences:
                entity_summary.append(f"{len(rna_sequences)} RNA sequence(s)")
            if ions:
                entity_summary.append(f"{len(ions)} ion(s)")
            console.print(f"Built input from flags: {', '.join(entity_summary)}")

        # Build params
        params = {
            "json_content": json_content,
            "model_name": model,
            "seeds": seeds,
            "sample": sample,
            "output_format": output_format.lower(),
        }

        if job_name:
            params["job_name"] = job_name

        # Handle template file (v1.0.0 feature)
        if template_file:
            # Template file will be uploaded to Supabase by the CLI framework
            # and the path will be set in template_supabase_path
            params["use_template"] = True
            params["template_local_path"] = str(template_file)
            console.print(f"Using template file: {template_file}")

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