# geney/utils.py — Shared utilities used across subpackages.
from __future__ import annotations

import logging

logger = logging.getLogger(__name__)


def _position_in_transcript(transcript, position: int) -> bool:
    """Check if a genomic position falls within a transcript's bounds."""
    try:
        start = transcript.transcript_start
        end = transcript.transcript_end
        return min(start, end) <= position <= max(start, end)
    except AttributeError:
        return False


def _is_protein_coding(transcript) -> bool:
    """Check if transcript is protein-coding."""
    try:
        # Check biotype if available
        biotype = getattr(transcript, 'biotype', None)
        if biotype:
            return biotype.lower() == 'protein_coding'
        # Fallback: check if it can generate a protein
        return hasattr(transcript, 'generate_protein')
    except Exception:
        return False


def select_transcript(gene, position: int, preferred_transcript_id: str | None = None):
    """
    Select the appropriate transcript for a mutation position.

    Logic:
    1. If preferred_transcript_id is given, use it if the position is within bounds
    2. Otherwise, try the primary/MANE transcript first
    3. If position is outside primary transcript, find a protein-coding transcript that contains it
    4. Return None if no suitable transcript found

    Args:
        gene: seqmat Gene object
        position: genomic position of the mutation
        preferred_transcript_id: optional specific transcript ID to use

    Returns:
        Transcript object or None if no suitable transcript found
    """
    # If a specific transcript is requested, try it first
    if preferred_transcript_id:
        try:
            t = gene.transcript(preferred_transcript_id)
            if _position_in_transcript(t, position):
                return t
        except Exception as e:
            logger.debug("Preferred transcript %s lookup failed: %s", preferred_transcript_id, e)

    # Try the primary/MANE transcript (default when no ID specified)
    try:
        primary = gene.transcript(None)
        if _position_in_transcript(primary, position):
            return primary
    except Exception as e:
        logger.debug("Primary transcript lookup failed: %s", e)

    # Position not in primary transcript - search for a protein-coding transcript that contains it
    try:
        transcripts = getattr(gene, 'transcripts', None) or getattr(gene, '_transcripts', {})
        if hasattr(transcripts, 'items'):
            transcript_list = list(transcripts.values())
        elif hasattr(transcripts, '__iter__'):
            transcript_list = list(transcripts)
        else:
            transcript_list = []

        for t in transcript_list:
            if _position_in_transcript(t, position) and _is_protein_coding(t):
                return t
    except Exception as e:
        logger.debug("Transcript search failed at position %d: %s", position, e)

    return None


def get_all_protein_coding_transcripts(gene, position: int | None = None) -> list:
    """Get all protein-coding transcripts for a gene.

    Args:
        gene: seqmat Gene object.
        position: If provided, only return transcripts containing this position.

    Returns:
        List of transcript objects, with primary/MANE first if present.
    """
    transcripts_attr = getattr(gene, 'transcripts', None) or getattr(gene, '_transcripts', {})
    if hasattr(transcripts_attr, 'values'):
        all_transcripts = list(transcripts_attr.values())
    elif hasattr(transcripts_attr, '__iter__'):
        all_transcripts = list(transcripts_attr)
    else:
        return []

    result = []
    primary_id = None
    try:
        primary = gene.transcript(None)
        primary_id = getattr(primary, 'transcript_id', None)
    except Exception:
        pass

    for t in all_transcripts:
        if not _is_protein_coding(t):
            continue
        if position is not None and not _position_in_transcript(t, position):
            continue
        result.append(t)

    # Sort: primary/MANE first, then by transcript_id for determinism
    if primary_id:
        result.sort(key=lambda t: (
            0 if getattr(t, 'transcript_id', '') == primary_id else 1,
            getattr(t, 'transcript_id', ''),
        ))
    else:
        result.sort(key=lambda t: getattr(t, 'transcript_id', ''))

    return result


def get_all_transcripts(gene, position: int | None = None) -> tuple[list, list]:
    """Get all transcripts for a gene, split into protein-coding and non-coding.

    Args:
        gene: seqmat Gene object.
        position: If provided, only return transcripts containing this position.

    Returns:
        Tuple of (protein_coding_list, non_coding_list).
        Primary/MANE transcript is first in the protein-coding list.
    """
    transcripts_attr = getattr(gene, 'transcripts', None) or getattr(gene, '_transcripts', {})
    if hasattr(transcripts_attr, 'values'):
        all_t = list(transcripts_attr.values())
    elif hasattr(transcripts_attr, '__iter__'):
        all_t = list(transcripts_attr)
    else:
        return [], []

    primary_id = None
    try:
        primary = gene.transcript(None)
        primary_id = getattr(primary, 'transcript_id', None)
    except Exception:
        pass

    coding = []
    non_coding = []
    for t in all_t:
        if position is not None and not _position_in_transcript(t, position):
            continue
        if _is_protein_coding(t):
            coding.append(t)
        else:
            non_coding.append(t)

    if primary_id:
        coding.sort(key=lambda t: (
            0 if getattr(t, 'transcript_id', '') == primary_id else 1,
            getattr(t, 'transcript_id', ''),
        ))
    else:
        coding.sort(key=lambda t: getattr(t, 'transcript_id', ''))

    non_coding.sort(key=lambda t: getattr(t, 'transcript_id', ''))
    return coding, non_coding


def _apply_mutation_safe(pre_mrna, pos: int, ref: str, alt: str, permissive: bool = True):
    """
    Apply a mutation to pre_mrna with optional permissive mode for ref mismatches.

    Args:
        pre_mrna: pre-mRNA sequence object with apply_mutations method
        pos: genomic position
        ref: reference allele
        alt: alternate allele
        permissive: if True, ignore reference mismatches (default True)
    """
    try:
        # Try standard apply_mutations first
        pre_mrna.apply_mutations((pos, ref, alt))
    except Exception as e:
        error_msg = str(e).lower()
        if permissive and ('mismatch' in error_msg or 'ref' in error_msg):
            # Ref mismatch - try applying without ref check if seqmat supports it
            try:
                # Some versions of seqmat support permissive flag
                pre_mrna.apply_mutations((pos, ref, alt), permissive=True)
            except TypeError:
                # Fallback: apply as insertion at position (ignoring ref)
                try:
                    pre_mrna.apply_mutations((pos, '', alt))
                except Exception as e2:
                    logger.debug("Mutation apply last-resort failed pos=%d ref=%s alt=%s: %s",
                                 pos, ref, alt, e2)
        else:
            raise
