# geney/pipelines/_tis.py — TIS analysis pipelines.
from __future__ import annotations

import logging

from seqmat import Gene

from ..variants import MutationalEvent
from ..results import TISResult
from ..utils import select_transcript

from ._common import _prepare_tis, _tis_completed_dict

logger = logging.getLogger(__name__)


def tis_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
    proximity_window: int = 500,
    probability_threshold: float = 0.1,
    use_prior: bool = False,
    skip_if_not_near_start: bool = True,
) -> TISResult | None:
    """Run TIS (Translation Initiation Site) analysis for a mutation.

    Applies the mutation directly to the mature mRNA (without splicing changes)
    and compares TIS probabilities between reference and variant.
    """
    from ..tis import compute_tis_shift_metrics

    m = MutationalEvent(mut_id)
    if not m.compatible():
        raise ValueError(f"Mutations in event are incompatible: {mut_id}")

    gene = Gene.from_file(m.gene, organism=organism)
    central_pos = m.central_position

    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return None

    reference_transcript = (
        selected
        .generate_mature_mrna()
        .generate_protein()
    )

    tis_inputs = _prepare_tis(reference_transcript, m, mut_id,
                              proximity_window, skip_if_not_near_start)
    if isinstance(tis_inputs, TISResult):
        return tis_inputs

    ref_mrna_seq, var_mrna_seq, mrna_positions = tis_inputs

    try:
        result = compute_tis_shift_metrics(
            ref_mrna=ref_mrna_seq,
            var_mrna=var_mrna_seq,
            canonical_start=None,
            probability_threshold=probability_threshold,
            use_prior=use_prior,
        )
        return _tis_completed_dict(
            result, mut_id, m.gene, reference_transcript.transcript_id,
            ref_mrna_seq, var_mrna_seq, mrna_positions,
        )
    except Exception as e:
        return TISResult(status="error", score=0.0, reason=str(e),
                         mut_id=mut_id, gene=m.gene)


def tis_quick_check(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
) -> dict:
    """Quick check if a mutation is likely to affect TIS.

    Fast pre-screening that doesn't run the full TITER model.
    """
    from ..tis.helpers import _mutation_in_mature_mrna, _mutation_near_start_codon

    m = MutationalEvent(mut_id)
    gene = Gene.from_file(m.gene, organism=organism)
    central_pos = m.central_position

    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return {
            'should_analyze': False,
            'reason': 'no_transcript',
            'mutation_in_mrna': False,
            'near_start_codon': False,
            'affects_start_codon': False,
        }

    reference_transcript = (
        selected
        .generate_mature_mrna()
    )

    mutation_in_mrna = _mutation_in_mature_mrna(reference_transcript, m.positions)
    near_start = _mutation_near_start_codon(reference_transcript, m.positions, window=500)

    affects_start = False
    try:
        mrna = reference_transcript.mature_mrna
        if hasattr(mrna, 'seq') and hasattr(mrna, 'index'):
            seq = mrna.seq.upper()
            atg_pos = seq.find('ATG')
            if atg_pos != -1 and atg_pos < len(mrna.index):
                atg_genomic = set(mrna.index[atg_pos:atg_pos+3])
                affects_start = any(pos in atg_genomic for pos in m.positions)
    except Exception as e:
        logger.debug("ATG check failed for %s: %s", mut_id, e)

    return {
        'should_analyze': mutation_in_mrna and (near_start or affects_start),
        'mutation_in_mrna': mutation_in_mrna,
        'near_start_codon': near_start,
        'affects_start_codon': affects_start,
        'gene': m.gene,
        'transcript_id': getattr(reference_transcript, 'transcript_id', None),
    }
