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

Supports protein, nucleic acid, and small molecule structure prediction
with automatic MSA generation via ColabFold.
"""

import json
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": "openfold3",
    "display_name": "OpenFold3",
    "category": "folding",
    "description": "Predict biomolecular structures (proteins, RNA, DNA, ligands)",
    "modal_function_name": "openfold3_worker",
    "modal_app_name": "openfold3-api",
    "status": "available",
    "outputs": {
        "structure_filepath": "Predicted structure in PDB/CIF format",
        "confidence_aggregated_filepath": "Aggregated confidence metrics (pLDDT, ptm, iptm)",
        "pae_filepath": "PAE matrix for inter-chain confidence",
        "ipsae_per_residue_filepath": "Per-residue interface SAE scores (for complexes)",
    },
}

console = Console()

# Valid character sets for validation (match schemas.py)
VALID_PROTEIN_CHARS = set("ACDEFGHIKLMNPQRSTVWY")
VALID_DNA_CHARS = set("ATCG")
VALID_RNA_CHARS = set("AUCG")


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 (A-Z)")


def _validate_protein_sequence(seq: str, source: str) -> str:
    """Validate and normalize protein sequence. Returns cleaned sequence or raises Exit."""
    cleaned = seq.strip().upper().replace(" ", "").replace("\n", "")
    invalid = set(cleaned) - VALID_PROTEIN_CHARS
    if invalid:
        console.print(
            f"[red]Error:[/red] Invalid amino acids in {source}: {', '.join(sorted(invalid))}\n"
            f"Valid: ACDEFGHIKLMNPQRSTVWY"
        )
        raise typer.Exit(1)
    if not cleaned:
        console.print(f"[red]Error:[/red] Empty sequence in {source}")
        raise typer.Exit(1)
    return cleaned


def _validate_dna_sequence(seq: str) -> str:
    """Validate and normalize DNA sequence."""
    cleaned = seq.strip().upper().replace(" ", "").replace("\n", "")
    invalid = set(cleaned) - VALID_DNA_CHARS
    if invalid:
        console.print(f"[red]Error:[/red] Invalid nucleotides in DNA: {', '.join(sorted(invalid))}\nValid: ATCG")
        raise typer.Exit(1)
    if not cleaned:
        console.print("[red]Error:[/red] Empty DNA sequence")
        raise typer.Exit(1)
    return cleaned


def _validate_rna_sequence(seq: str) -> str:
    """Validate and normalize RNA sequence."""
    cleaned = seq.strip().upper().replace(" ", "").replace("\n", "")
    invalid = set(cleaned) - VALID_RNA_CHARS
    if invalid:
        console.print(f"[red]Error:[/red] Invalid nucleotides in RNA: {', '.join(sorted(invalid))}\nValid: AUCG")
        raise typer.Exit(1)
    if not cleaned:
        console.print("[red]Error:[/red] Empty RNA sequence")
        raise typer.Exit(1)
    return cleaned


def _validate_smiles(smiles: str) -> str:
    """Basic SMILES validation (non-empty)."""
    cleaned = smiles.strip()
    if not cleaned:
        console.print("[red]Error:[/red] Empty SMILES string")
        raise typer.Exit(1)
    return cleaned


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("openfold3")
    def run_openfold3(
        sequence: 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: Optional[List[Path]] = typer.Option(
            None,
            "--fasta",
            "-f",
            help="FASTA file for protein sequence (can specify multiple)",
            exists=True,
        ),
        ligand: 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.",
        ),
        dna: 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: 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.",
        ),
        query_json: Optional[Path] = typer.Option(
            None,
            "--query",
            "-q",
            help="Path to JSON file with complex query (bypasses entity flags)",
            exists=True,
        ),
        quality: str = typer.Option(
            "quick",
            "--quality",
            help="Quality preset: 'quick' (faster) or 'best' (more samples)",
        ),
        use_msa: bool = typer.Option(
            False,
            "--use-msa",
            help="Enable ColabFold MSA server for automatic MSA generation",
        ),
        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",
            "-j",
            help="Custom job name for output files (default: random 4-letter code)",
        ),
        background: bool = typer.Option(
            False,
            "--background",
            "-b",
            help="Submit job and return immediately without waiting for completion",
        ),
    ):
        """
        Predict biomolecular structures using OpenFold3.

        Supports proteins, ligands, DNA, and RNA. Use repeatable flags to build
        multi-chain complexes, or --query for advanced JSON input.

        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 openfold3 -s "MKFLILLFNILCLFPVLAADNH" -o ./output/

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

            # Two proteins (heterodimer)
            amina run openfold3 -s "MKFLIL" -s "PNNTHEQ" -o ./output/

            # Protein + ligand complex
            amina run openfold3 -s "MKFLILLFNILCLFPVLAADNH" -l "CCO" -o ./output/

            # Protein + DNA complex
            amina run openfold3 -s "MKFLIL" --dna "ATCGATCG" -o ./output/

            # From FASTA files
            amina run openfold3 -f ./protein1.fasta -f ./protein2.fasta -o ./output/

            # With MSA generation (better accuracy, slower)
            amina run openfold3 -s "MKFLIL..." --use-msa -o ./output/

            # Best quality preset (more samples)
            amina run openfold3 -s "MKFLIL..." --quality best -o ./output/

            # Advanced query via JSON file
            amina run openfold3 --query ./complex.json -o ./output/
        """
        # 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)

        # Validate quality preset
        if quality not in ("quick", "best"):
            console.print(f"[red]Error:[/red] Invalid quality preset '{quality}'. Use 'quick' or 'best'.")
            raise typer.Exit(1)

        # Check if any entities provided via flags
        has_entities = any(
            [
                sequence,
                fasta,
                ligand,
                dna,
                rna,
            ]
        )

        # Must have either entities OR query JSON
        if not has_entities and not query_json:
            console.print(
                "[red]Error:[/red] Provide entities via flags (-s, -f, -l, --dna, --rna) or --query JSON file"
            )
            raise typer.Exit(1)

        if has_entities and query_json:
            console.print("[red]Error:[/red] Cannot combine entity flags with --query")
            raise typer.Exit(1)

        # Build params
        params = {}

        if query_json:
            # Load complex query from JSON file
            try:
                query_data = json.loads(query_json.read_text())
            except json.JSONDecodeError as e:
                console.print(f"[red]Error:[/red] Invalid JSON in query file: {e}")
                raise typer.Exit(1)

            # Validate structure
            if "queries" not in query_data or not query_data["queries"]:
                console.print("[red]Error:[/red] Query JSON must have a 'queries' array with at least one query")
                raise typer.Exit(1)

            params["queries"] = query_data["queries"]
            console.print(f"Loaded query from {query_json} with {len(params['queries'])} queries")

        else:
            # Build chains list from flags
            used_chains = set()
            chains = []

            # Process proteins from --sequence (with validation, supports X:SEQ format)
            for i, seq in enumerate(sequence or []):
                parsed_chain_id, raw_seq = parse_sequence_with_chain_id(seq, used_chains)
                validated_seq = _validate_protein_sequence(raw_seq, f"--sequence #{i + 1}")
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                chains.append(
                    {
                        "molecule_type": "protein",
                        "chain_ids": chain_id,
                        "sequence": validated_seq,
                    }
                )

            # Process proteins from --fasta (single-sequence FASTA only, with validation)
            for fasta_path in fasta or []:
                content = fasta_path.read_text()
                # Count headers to ensure single-sequence
                headers = [line for line in content.splitlines() if line.strip().startswith(">")]
                if len(headers) > 1:
                    console.print(
                        f"[red]Error:[/red] {fasta_path} contains multiple sequences. Use one --fasta per sequence."
                    )
                    raise typer.Exit(1)
                seq = "".join(
                    line.strip() for line in content.splitlines() if line.strip() and not line.startswith(">")
                )
                validated_seq = _validate_protein_sequence(seq, f"--fasta {fasta_path}")
                chain_id = _auto_chain_id(used_chains)
                chains.append(
                    {
                        "molecule_type": "protein",
                        "chain_ids": chain_id,
                        "sequence": validated_seq,
                    }
                )

            # Process ligands from --ligand (SMILES, with validation, supports X:SMILES format)
            for smiles in ligand or []:
                parsed_chain_id, raw_smiles = parse_sequence_with_chain_id(smiles, used_chains)
                validated_smiles = _validate_smiles(raw_smiles)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                chains.append(
                    {
                        "molecule_type": "ligand",
                        "chain_ids": chain_id,
                        "smiles": validated_smiles,
                    }
                )

            # Process DNA from --dna (with validation, supports X:SEQ format)
            for dna_seq in dna or []:
                parsed_chain_id, raw_seq = parse_sequence_with_chain_id(dna_seq, used_chains)
                validated_seq = _validate_dna_sequence(raw_seq)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                chains.append(
                    {
                        "molecule_type": "dna",
                        "chain_ids": chain_id,
                        "sequence": validated_seq,
                    }
                )

            # Process RNA from --rna (with validation, supports X:SEQ format)
            for rna_seq in rna or []:
                parsed_chain_id, raw_seq = parse_sequence_with_chain_id(rna_seq, used_chains)
                validated_seq = _validate_rna_sequence(raw_seq)
                chain_id = parsed_chain_id if parsed_chain_id else _auto_chain_id(used_chains)
                chains.append(
                    {
                        "molecule_type": "rna",
                        "chain_ids": chain_id,
                        "sequence": validated_seq,
                    }
                )

            params["queries"] = [{"chains": chains}]

            # Display what we're predicting
            entity_counts = []
            protein_count = len(sequence or []) + len(fasta or [])
            if protein_count:
                entity_counts.append(f"{protein_count} protein(s)")
            if ligand:
                entity_counts.append(f"{len(ligand)} ligand(s)")
            if dna:
                entity_counts.append(f"{len(dna)} DNA")
            if rna:
                entity_counts.append(f"{len(rna)} RNA")

            console.print(f"Predicting: {', '.join(entity_counts)} → chains {', '.join(sorted(used_chains))}")

        # Build config
        params["config"] = {
            "quality_preset": quality,
            "use_msa_server": use_msa,
        }

        if job_name:
            params["job_name"] = job_name

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