# geney/tis/helpers.py — TIS-related helper functions extracted from pipelines.py.
from __future__ import annotations


def _get_mature_mrna_seq(transcript) -> str:
    """Extract mature mRNA sequence as string."""
    mrna = getattr(transcript, 'mature_mrna', None)
    if mrna is None:
        return ""
    if hasattr(mrna, 'seq'):
        return mrna.seq
    return str(mrna)


def _mutation_in_mature_mrna(transcript, mutation_positions: list) -> bool:
    """
    Check if any mutation position falls within the mature mRNA.

    Args:
        transcript: Transcript with mature_mrna attribute
        mutation_positions: List of genomic positions

    Returns:
        True if at least one mutation is in the mature mRNA
    """
    try:
        mrna = transcript.mature_mrna
        if mrna is None:
            return False

        # Get the genomic indices covered by mature mRNA
        if hasattr(mrna, 'index'):
            mrna_indices = set(mrna.index)
        else:
            # Fallback: assume transcript coordinates
            return True  # Can't determine, assume yes

        for pos in mutation_positions:
            if pos in mrna_indices:
                return True
        return False
    except Exception:
        return True  # Can't determine, assume yes to be safe


def _mutation_near_start_codon(transcript, mutation_positions: list, window: int = 500) -> bool:
    """
    Check if any mutation is near the start codon (within window of 5' end).

    Args:
        transcript: Transcript object
        mutation_positions: List of genomic positions
        window: Distance in nucleotides from 5' end to consider "near" (default 500)

    Returns:
        True if mutation is near the start codon region
    """
    try:
        mrna = transcript.mature_mrna
        if mrna is None or not hasattr(mrna, 'index'):
            return True  # Can't determine, assume yes

        indices = mrna.index
        if len(indices) == 0:
            return True

        # Get the 5' region (start of mature mRNA + window)
        # Account for strand direction
        rev = getattr(transcript, 'rev', False)
        if rev:
            # Reverse strand: 5' end is at higher coordinates
            five_prime_region = set(indices[:window]) if len(indices) > window else set(indices)
        else:
            # Forward strand: 5' end is at lower coordinates
            five_prime_region = set(indices[:window]) if len(indices) > window else set(indices)

        for pos in mutation_positions:
            if pos in five_prime_region:
                return True
        return False
    except Exception:
        return True  # Can't determine, assume yes


def _apply_mutation_to_mrna(ref_mrna_seq: str, ref_mrna_index, mutations: list) -> str:
    """
    Apply mutations directly to mature mRNA sequence.

    Args:
        ref_mrna_seq: Reference mature mRNA sequence
        ref_mrna_index: Genomic indices for each position in mRNA
        mutations: List of (pos, ref, alt) tuples

    Returns:
        Variant mRNA sequence with mutations applied
    """
    # Create position to mRNA index mapping
    if hasattr(ref_mrna_index, 'tolist'):
        idx_list = ref_mrna_index.tolist()
    else:
        idx_list = list(ref_mrna_index)

    pos_to_mrna_idx = {pos: i for i, pos in enumerate(idx_list)}

    # Convert to list for mutation
    mrna_list = list(ref_mrna_seq)

    # Sort mutations by position (reverse order to handle indels correctly)
    sorted_mutations = sorted(mutations, key=lambda x: x[0], reverse=True)

    for pos, ref, alt in sorted_mutations:
        if pos not in pos_to_mrna_idx:
            continue  # Mutation not in mature mRNA

        mrna_idx = pos_to_mrna_idx[pos]

        # Handle SNV
        if len(ref) == 1 and len(alt) == 1:
            if mrna_idx < len(mrna_list):
                mrna_list[mrna_idx] = alt
        # Handle deletions
        elif len(ref) > len(alt):
            del_len = len(ref) - len(alt)
            for i in range(del_len):
                if mrna_idx + len(alt) + i < len(mrna_list):
                    mrna_list[mrna_idx + len(alt) + i] = ''
            if len(alt) > 0 and mrna_idx < len(mrna_list):
                mrna_list[mrna_idx] = alt
        # Handle insertions
        elif len(alt) > len(ref):
            if mrna_idx < len(mrna_list):
                mrna_list[mrna_idx] = alt

    return ''.join(mrna_list)
