"""Utility functions for validation, file I/O, and common operations."""

import os
from pathlib import Path
from typing import List, Optional, Tuple, Union

# Standard 20 amino acids
STANDARD_AMINO_ACIDS = [
    "A",
    "C",
    "D",
    "E",
    "F",
    "G",
    "H",
    "I",
    "K",
    "L",
    "M",
    "N",
    "P",
    "Q",
    "R",
    "S",
    "T",
    "V",
    "W",
    "Y",
]


def read_lines(filepath: Path) -> List[str]:
    """Read lines from a file and strip whitespace.

    Parameters
    ----------
    filepath : Path
        Path to the file to read.

    Returns
    -------
    List[str]
        List of stripped lines.
    """
    with open(filepath, "r", encoding="utf-8") as f:
        return [line.strip() for line in f if line.strip()]


def is_valid_amino_acid(letter: str) -> bool:
    """Check if a letter is a valid amino acid.

    Parameters
    ----------
    letter : str
        Single character to check.

    Returns
    -------
    bool
        True if valid amino acid, False otherwise.
    """
    return letter.upper() in STANDARD_AMINO_ACIDS


def parse_mutation(mutation_str: str) -> Tuple[str, str, int, Optional[str]]:
    """Parse a mutation string in format 'Chain.Wildtype.Position.Mutant'.

    This function only parses the format and does not validate the values.
    Use validate_mutation() for full validation.

    Parameters
    ----------
    mutation_str : str
        Mutation string (e.g., "B.H.68.F").

    Returns
    -------
    Tuple[str, str, int, str]
        Tuple of (chain, wildtype, position, mutant).

    Raises
    ------
    ValueError
        If mutation format is invalid (wrong number of parts or position not numeric).
    """
    parts = mutation_str.split(".")

    if len(parts) != 4 and len(parts) != 3:
        raise ValueError(
            f"Invalid mutation format: {mutation_str} with length {len(parts)}."
            f"Expected 'Chain.Wildtype.Position.Mutant' (e.g., 'B.H.68.F') or 'Chain.Wildtype.Position' (e.g., 'B.H.68')"
        )

    if len(parts) == 3:
        chain, wildtype, position = parts
        return chain, wildtype, int(position), None
    else:
        chain, wildtype, position, mutant = parts
        return chain, wildtype, int(position), mutant


def validate_mutation(
    mutation_str: str, protein_sequences: List[str]
) -> Tuple[bool, Optional[str]]:
    """Validate a mutation against protein sequences.

    Parameters
    ----------
    mutation_str : str
        Mutation string (e.g., "B.H.68.F").
    protein_sequences : List[str]
        List of protein sequences (one per chain).

    Returns
    -------
    Tuple[bool, Optional[str]]
        Tuple of (is_valid, error_message). If valid, error_message is None.

    """
    # Parse the mutation format
    try:
        chain, wildtype, position, mutant = parse_mutation(mutation_str)
    except ValueError as e:
        return False, str(e)

    # Validate chain (single uppercase letter)
    if not chain or len(chain) != 1 or not chain.isupper():
        return False, f"Invalid chain identifier: {chain}"

    # Validate wildtype and mutant amino acids
    if not (is_valid_amino_acid(wildtype) and is_valid_amino_acid(mutant)):
        invalid = wildtype if not is_valid_amino_acid(wildtype) else mutant
        kind = "wildtype" if not is_valid_amino_acid(wildtype) else "mutant"
        return False, f"Invalid {kind} residue: {invalid}"

    # Validate position
    if position <= 1:
        return False, f"Position must be > 1, got {position}"

    # Check chain index
    chain_idx = get_chain_index(chain)
    if chain_idx < 0 or chain_idx >= len(protein_sequences):
        return (
            False,
            f"Chain {chain} is out of range (available: A-{chr(ord('A') + len(protein_sequences) - 1)})",
        )

    # Check position
    sequence = protein_sequences[chain_idx]
    if position > len(sequence):
        return False, f"Position {position} exceeds chain {chain} length ({len(sequence)})"

    # Check wildtype matches
    actual_wildtype = sequence[position - 1]  # 0-indexed
    if actual_wildtype != wildtype:
        return False, (
            f"Wildtype mismatch at position {position} in chain {chain}: "
            f"expected {actual_wildtype}, got {wildtype}"
        )

    # Check mutant is different
    if wildtype == mutant:
        return False, f"Mutant residue {mutant} is the same as wildtype {wildtype}"

    return True, None


def validate_mutations(
    mutation_strings: List[str], protein_sequences: List[str]
) -> Tuple[bool, List[str]]:
    """Validate a list of mutations.

    Parameters
    ----------
    mutation_strings : List[str]
        List of mutation strings.
    protein_sequences : List[str]
        List of protein sequences.

    Returns
    -------
    Tuple[bool, List[str]]
        Tuple of (all_valid, error_messages).

    """
    errors = []
    for mut_str in mutation_strings:
        is_valid, error = validate_mutation(mut_str, protein_sequences)
        if not is_valid:
            errors.append(f"{mut_str}: {error}")

    return len(errors) == 0, errors


def ensure_directory(path: Union[str, Path]) -> Path:
    """Ensure a directory exists, creating it if necessary.

    Parameters
    ----------
    path : Union[str, Path]
        Directory path.

    Returns
    -------
    Path
        Path object for the directory.

    """

    dir_path = Path(path)
    dir_path.mkdir(parents=True, exist_ok=True)
    return dir_path


def get_chain_index(chain: str) -> int:
    """Convert chain identifier to index (A=0, B=1, etc.).

    Parameters
    ----------
    chain : str
        Chain identifier (single uppercase letter).

    Returns
    -------
    int
        Zero-based index.

    """
    return ord(chain) - ord("A")


def parse_mutations_file(filepath: Path) -> List[str]:
    """Parse mutations from a file.

    Each line can contain:
    - Single mutation: "B.H.68.F"
    - Multiple mutations (comma-separated): "B.H.68.F,A.K.42.R"
    - Saturation mutation (3 parts): "B.H.68"

    Parameters
    ----------
    filepath : Path
        Path to mutations file.

    Returns
    -------
    List[str]
        List of mutation strings.
    """
    lines = read_lines(str(filepath))
    mutations = []

    for line in lines:
        # Remove comments
        if "#" in line:
            line = line.split("#")[0].strip()
        if not line:
            continue

        mutations.append(line)

    return mutations
