# geney/pipelines/_mutation.py — Unified single-mutation pipeline.
from __future__ import annotations

import logging
import time

from ..results import (
    SplicingResult,
    TISResult,
    FoldingResult,
    TranscriptionResult,
    MutationResult,
)
from ..splicing import SpliceSimulator, TranscriptLibrary

from ._common import (
    SPLICING_CONTEXT_BP,
    SPLICE_MAX_DISTANCE,
    _prepare_mutation,
    _build_splice_site_index,
    _splicing_base_report,
    _oncosplice_dataframe,
    _prepare_tis,
    _tis_completed_dict,
    _run_folding_single,
    _transcription_result_from_dict,
)

logger = logging.getLogger(__name__)


def mutation_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
    splicing_engine: str = "spliceai",
    context_bp: int = 1000,
    run_splicing: bool = True,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    tis_proximity_window: int = 500,
    tis_probability_threshold: float = 0.1,
    folding_window_size: int = 39,
    gene_load_timeout: float | None = None,
    max_isoforms: int | None = None,
) -> MutationResult:
    """Run splicing, TIS, mRNA-folding, and transcription analyses for a single mutation.

    Loads the gene and selects a transcript once, then dispatches to each
    subsystem.  Returns a :class:`MutationResult` namedtuple.

    Args:
        mut_id: Mutation ID (``GENE:CHR:POS:REF:ALT``; pipe-separated for compound).
        transcript_id: Optional specific transcript ID.
        organism: Genome build (default ``"hg38"``).
        splicing_engine: Splicing engine (default ``"spliceai"``).
        context_bp: Base pairs on each side of mutation to extract (default 1000).
        run_splicing: Whether to run splicing/oncosplice analysis.
        run_tis: Whether to run TIS analysis.
        run_folding: Whether to run mRNA folding analysis.
        run_transcription: Whether to run Borzoi transcription-effect analysis.
        tis_proximity_window: Max distance from 5' end for TIS analysis (nt).
        tis_probability_threshold: Min probability to consider a TIS.
        folding_window_size: Local folding window for delta-MFE (nt).
        gene_load_timeout: Timeout for gene loading (seconds); *None* = no limit.
    """
    _t_total = time.perf_counter()
    timing: dict[str, float] = {}
    errors: dict[str, str] = {}

    # -- Prepare (parse → gene → transcript → mutate → context) ---------------
    _t0 = time.perf_counter()
    p = _prepare_mutation(mut_id, organism, transcript_id, context_bp, gene_load_timeout)
    timing["prepare"] = round(time.perf_counter() - _t0, 4)

    if p["status"] != "ok":
        timing["total"] = round(time.perf_counter() - _t_total, 4)
        return MutationResult(timing=timing)

    m = p["event"]
    central_pos = p["central_pos"]
    gene = p["gene"]
    selected = p["selected"]
    reference_transcript = p["ref_transcript"]
    mutated_transcript = p["mut_transcript"]
    mutation_context = p["mutation_context"]

    # -- 1. Splicing / Oncosplice ---------------------------------------------
    splicing_result = SplicingResult()
    if run_splicing:
        _t0 = time.perf_counter()
        try:
            # Use scoped pre-mRNA read for splicing prediction — only fetch ~15kb
            ref_for_splicing = selected.clone()
            ref_for_splicing.generate_pre_mrna(
                region_start=central_pos - SPLICING_CONTEXT_BP,
                region_end=central_pos + SPLICING_CONTEXT_BP,
            )
            tl = TranscriptLibrary(ref_for_splicing, m, individual_mutations=False)
            tl.predict_splicing(central_pos, engine=splicing_engine, inplace=True)
            splicing_preds = tl.get_event_columns("event")

            # Clear scoped pre-mRNA so isoform clones use direct-exon assembly
            tl.event._pre_mrna = None

            ss = SpliceSimulator(
                splicing_preds, tl.event, feature="event", max_distance=SPLICE_MAX_DISTANCE,
                mutations=list(m.mutation_args()),
            )

            base_report = _splicing_base_report(mut_id, m, reference_transcript, splicing_engine, central_pos)
            ss_metadata = ss.report(central_pos)
            splice_index = _build_splice_site_index(gene)
            splicing_result = _oncosplice_dataframe(reference_transcript, ss, ss_metadata, base_report, max_isoforms=max_isoforms, splice_index=splice_index)
        except Exception as e:
            splicing_result = SplicingResult()
            errors["splicing"] = str(e)
        timing["spliceai"] = round(time.perf_counter() - _t0, 4)

    # -- 1b. Oncosplice scoring (included in spliceai timing above) -----------
    timing["oncosplice"] = 0.0  # oncosplice runs inside the splicing block

    # -- 2. TIS ---------------------------------------------------------------
    tis_result = None
    if run_tis:
        _t0 = time.perf_counter()
        try:
            from ..tis import compute_tis_shift_metrics

            tis_inputs = _prepare_tis(reference_transcript, m, mut_id, tis_proximity_window)
            if isinstance(tis_inputs, TISResult):
                tis_result = tis_inputs
            else:
                ref_mrna_seq, var_mrna_seq, mrna_positions = tis_inputs
                result = compute_tis_shift_metrics(
                    ref_mrna=ref_mrna_seq,
                    var_mrna=var_mrna_seq,
                    canonical_start=None,
                    probability_threshold=tis_probability_threshold,
                )
                tis_result = _tis_completed_dict(
                    result, mut_id, m.gene, reference_transcript.transcript_id,
                    ref_mrna_seq, var_mrna_seq, mrna_positions,
                )
        except Exception as e:
            tis_result = TISResult(status="error", score=0.0, reason=str(e),
                                   mut_id=mut_id, gene=m.gene)
            errors["tis"] = str(e)
        timing["tis"] = round(time.perf_counter() - _t0, 4)

    # -- 3. mRNA Folding ------------------------------------------------------
    folding_result = None
    if run_folding:
        _t0 = time.perf_counter()
        try:
            folding_result = _run_folding_single(reference_transcript, m, mut_id, folding_window_size)
        except Exception as e:
            folding_result = FoldingResult(status="error", reason=str(e),
                                           mut_id=mut_id, gene=m.gene)
            errors["folding"] = str(e)
        timing["folding"] = round(time.perf_counter() - _t0, 4)

    # -- 4. Transcription (Borzoi) --------------------------------------------
    transcription_result = None
    if run_transcription:
        _t0 = time.perf_counter()
        try:
            from ..transcription.prediction import predict_transcription_effect
            raw = predict_transcription_effect(
                chrom=str(m.chrom), pos=central_pos,
                mutations=m.mutation_args(), organism=organism,
            )
            transcription_result = _transcription_result_from_dict(raw, mut_id, m.gene)
        except Exception as e:
            transcription_result = TranscriptionResult(
                status="error", reason=str(e), mut_id=mut_id, gene=m.gene)
            errors["borzoi"] = str(e)
        timing["borzoi"] = round(time.perf_counter() - _t0, 4)

    timing["total"] = round(time.perf_counter() - _t_total, 4)

    return MutationResult(
        splicing=splicing_result,
        tis=tis_result,
        folding=folding_result,
        transcription=transcription_result,
        transcript=mutated_transcript,
        mutation_context=mutation_context,
        timing=timing,
        errors=errors,
    )
