# geney/pipelines/_common.py — Constants, caches, and shared helpers.
from __future__ import annotations

import logging
import threading
import time
from collections import OrderedDict
from datetime import datetime
from typing import Optional

import numpy as np
import pandas as pd

from seqmat import Gene

from ..results import (
    SplicingResult,
    TISResult,
    FoldingResult, FoldingWindow,
    TranscriptionResult,
    MutationResult,
)
from ..splicing import SpliceSimulator, TranscriptLibrary
from ..variants import MutationalEvent
from ..oncosplice import Oncosplice
from ..utils import select_transcript, _is_protein_coding, _position_in_transcript

logger = logging.getLogger(__name__)

# -- Constants ---------------------------------------------------------------
SPLICING_CONTEXT_BP = 7500          # Half-width of scoped pre-mRNA for splicing
SPLICE_MAX_DISTANCE = 100_000_000   # SpliceSimulator max_distance sentinel
_GENE_CACHE_MAX = 32


# -- LRU Gene Cache ----------------------------------------------------------

_gene_cache: OrderedDict[tuple[str, str], Gene] = OrderedDict()
_gene_cache_lock = threading.Lock()


def _load_gene_with_timeout(gene_name: str, organism: str, timeout_seconds: Optional[float]) -> Optional[Gene]:
    """Load gene via Gene.from_file, optionally with a timeout.

    Caches up to _GENE_CACHE_MAX genes with LRU eviction.
    """
    key = (gene_name, organism)
    with _gene_cache_lock:
        if key in _gene_cache:
            _gene_cache.move_to_end(key)      # mark as recently used
            return _gene_cache[key]

    result: list = []
    exc: list = []

    def _run():
        try:
            result.append(Gene.from_file(gene_name, organism=organism))
        except Exception as e:
            exc.append(e)

    if timeout_seconds is not None and timeout_seconds > 0:
        thread = threading.Thread(target=_run, daemon=True)
        thread.start()
        thread.join(timeout=timeout_seconds)
        if thread.is_alive():
            raise TimeoutError(
                f"Loading gene '{gene_name}' from seqmat did not complete within {timeout_seconds}s."
            )
        if exc:
            raise exc[0]
        gene = result[0] if result else None
    else:
        gene = Gene.from_file(gene_name, organism=organism)

    if gene is not None:
        with _gene_cache_lock:
            if key in _gene_cache:
                _gene_cache.move_to_end(key)
            else:
                if len(_gene_cache) >= _GENE_CACHE_MAX:
                    _gene_cache.popitem(last=False)   # evict least-recently used
                _gene_cache[key] = gene
    return gene


# -- Shared helpers -----------------------------------------------------------

def _build_splice_site_index(gene) -> dict[frozenset, str]:
    """Build a lookup from splice-site signature to transcript ID.

    Each key is a frozenset of (donor, acceptor) tuples; the value is the
    transcript_id of the annotated transcript that uses that exact splice
    pattern.  O(T) construction, O(1) lookup per isoform.
    """
    index: dict[frozenset, str] = {}
    transcripts = getattr(gene, 'transcripts', None) or getattr(gene, '_transcripts', {})
    if hasattr(transcripts, 'values'):
        transcript_list = list(transcripts.values())
    elif hasattr(transcripts, '__iter__'):
        transcript_list = list(transcripts)
    else:
        return index

    for t in transcript_list:
        tid = getattr(t, 'transcript_id', None)
        donors = tuple(sorted(getattr(t, 'donors', [])))
        acceptors = tuple(sorted(getattr(t, 'acceptors', [])))
        if tid and (donors or acceptors):
            index[frozenset(donors + acceptors)] = tid
    return index


def _match_known_transcript(variant_transcript, splice_index: dict[frozenset, str]) -> str | None:
    """Return the transcript ID if the variant's splice sites match a known transcript."""
    if not splice_index:
        return None
    donors = tuple(sorted(getattr(variant_transcript, 'donors', [])))
    acceptors = tuple(sorted(getattr(variant_transcript, 'acceptors', [])))
    return splice_index.get(frozenset(donors + acceptors))


def _get_conservation_vector(transcript) -> np.ndarray:
    """Get conservation vector for a transcript, with fallback to ones."""
    cons_vector = getattr(transcript, 'cons_vector', None)
    if cons_vector is not None:
        return cons_vector
    protein_length = len(getattr(transcript, 'protein', ''))
    if protein_length == 0:
        protein_length = 1
    return np.ones(protein_length)


def _extract_gained_alternatives(result) -> list[dict]:
    """Extract alternative TIS sites that gained significant probability."""
    df = result.analysis_df
    if df is None or df.empty:
        return []
    alts = df[
        (df["location"] != "canonical")
        & (df["prob_change"] > 0)
        & (df["var_probability"] > 0)
    ].sort_values("prob_change", ascending=False)
    return alts.to_dict(orient="records")


def _splicing_base_report(mut_id, m, ref_transcript, splicing_engine, central_pos):
    """Build the shared base report Series for splicing results."""
    return pd.Series({
        "mut_id": mut_id,
        "gene": m.gene,
        "transcript_id": ref_transcript.transcript_id,
        "primary_transcript": getattr(ref_transcript, "primary_transcript", False),
        "splicing_engine": splicing_engine,
        "central_position": central_pos,
        "mutation_count": len(m.positions),
        "time_of_execution": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    })


def _oncosplice_dataframe(ref_transcript, ss, ss_metadata, base_report, max_isoforms=None, splice_index=None):
    """Run Oncosplice on all viable transcripts, return SplicingResult."""
    cons_vector = _get_conservation_vector(ref_transcript)
    rows = []
    for variant_transcript, isoform_metadata in ss.get_viable_transcripts(metadata=True, max_isoforms=max_isoforms):
        onco = Oncosplice(
            ref_transcript.protein,
            variant_transcript.protein,
            cons_vector,
        )
        matched_tid = _match_known_transcript(variant_transcript, splice_index) if splice_index else None
        row = pd.concat([
            isoform_metadata,
            pd.Series({
                "variant_mrna": variant_transcript.mature_mrna.seq,
                "known_transcript_id": matched_tid,
            }),
            onco.get_analysis_series(),
        ])
        rows.append(row)

    isoforms_df = pd.DataFrame(rows)

    # Drop columns that are stored once on SplicingResult, not per-isoform
    for col in ("reference_protein", "reference_length"):
        if col in isoforms_df.columns:
            isoforms_df = isoforms_df.drop(columns=col)

    missplicing = float(ss_metadata.get("missplicing", 0.0) or 0.0)

    max_oncosplice = 0.0
    if not isoforms_df.empty and "oncosplice_score" in isoforms_df.columns:
        max_oncosplice = float(isoforms_df["oncosplice_score"].max())

    return SplicingResult(
        score=abs(missplicing),
        missplicing=missplicing,
        max_oncosplice_score=max_oncosplice,
        isoforms=isoforms_df,
        reference_protein=ref_transcript.protein,
        reference_mrna=ref_transcript.mature_mrna.seq,
        conservation_vector=cons_vector,
        mut_id=base_report.get("mut_id", ""),
        gene=base_report.get("gene", ""),
        transcript_id=base_report.get("transcript_id"),
        splicing_engine=base_report.get("splicing_engine", ""),
        central_position=int(base_report.get("central_position", 0)),
        donor_df=ss.donor_df,
        acceptor_df=ss.acceptor_df,
        splice_events=ss_metadata,
    )


def _mrna_mutation_positions(mrna, mutation_args):
    """Map genomic mutation positions to 0-based mRNA positions."""
    if not hasattr(mrna, "index"):
        return []
    idx_list = mrna.index.tolist() if hasattr(mrna.index, "tolist") else list(mrna.index)
    pos_to_mrna = {pos: i for i, pos in enumerate(idx_list)}
    return [pos_to_mrna[pos] for pos, _, _ in mutation_args if pos in pos_to_mrna]


def _tis_completed_dict(result, mut_id, gene, transcript_id,
                        ref_mrna_seq, var_mrna_seq, mrna_positions):
    """Build a TISResult from a TISShiftResult."""
    from ..tis import get_tis_summary

    atg_sites = result.analysis_df.to_dict("records") if not result.analysis_df.empty else []

    # Build canonical TISSite
    from ..results import TISSite
    canonical_row = next((s for s in atg_sites if s.get("location") == "canonical"), None)
    canonical_site = None
    if canonical_row:
        canonical_site = TISSite(
            position=canonical_row["position"],
            location="canonical",
            in_frame=True,
            ref_probability=canonical_row["ref_probability"],
            var_probability=canonical_row["var_probability"],
            prob_change=canonical_row["prob_change"],
            ref_usage=canonical_row.get("ref_usage", 0.0),
            var_usage=canonical_row.get("var_usage", 0.0),
        )

    return TISResult(
        status="completed",
        score=result.tis_shift_score,
        canonical=canonical_site,
        canonical_prob_ref=result.canonical_ref_prob,
        canonical_prob_var=result.canonical_var_prob,
        canonical_prob_change=result.canonical_prob_change,
        canonical_usage_ref=result.canonical_ref_usage,
        canonical_usage_var=result.canonical_var_usage,
        canonical_usage_change=result.canonical_usage_change,
        initiation_intact=result.initiation_intact,
        gained_alternative_tis=(gained := _extract_gained_alternatives(result)),
        num_alternative_tis=len(gained),
        top_alternative_position=gained[0]["position"] if gained else None,
        top_alternative_prob=gained[0]["var_probability"] if gained else 0.0,
        top_alternative_location=gained[0]["location"] if gained else None,
        top_alternative_in_frame=gained[0]["in_frame"] if gained else None,
        atg_sites=atg_sites,
        ref_mrna_length=len(ref_mrna_seq),
        var_mrna_length=len(var_mrna_seq),
        ref_mrna_window=ref_mrna_seq[:600],
        var_mrna_window=var_mrna_seq[:600],
        mutation_mrna_positions=mrna_positions,
        mut_id=mut_id,
        gene=gene,
        transcript_id=transcript_id,
        summary=get_tis_summary(result),
        tis_result=result,
    )


def _prepare_tis(ref_transcript, m, mut_id, proximity_window=500,
                 skip_if_not_near_start=True):
    """Prepare TIS inputs. Returns (ref_mrna, var_mrna, positions) or a TISResult for skip/error."""
    from ..tis.helpers import (
        _mutation_in_mature_mrna,
        _mutation_near_start_codon,
        _apply_mutation_to_mrna,
        _get_mature_mrna_seq,
    )

    tid = getattr(ref_transcript, "transcript_id", None)

    if not _mutation_in_mature_mrna(ref_transcript, m.positions):
        return TISResult(status="skipped", score=0.0, reason="mutation_not_in_mature_mrna",
                         mut_id=mut_id, gene=m.gene, transcript_id=tid)

    if skip_if_not_near_start:
        if not _mutation_near_start_codon(ref_transcript, m.positions, proximity_window):
            return TISResult(status="skipped", score=0.0, reason="mutation_far_from_start_codon",
                             mut_id=mut_id, gene=m.gene, transcript_id=tid,
                             proximity_window=proximity_window)

    ref_mrna_seq = _get_mature_mrna_seq(ref_transcript)
    if not ref_mrna_seq:
        return TISResult(status="error", score=0.0, reason="no_mature_mrna", mut_id=mut_id)

    mrna = ref_transcript.mature_mrna
    var_mrna_seq = (_apply_mutation_to_mrna(ref_mrna_seq, mrna.index, m.mutation_args())
                    if hasattr(mrna, "index") else ref_mrna_seq)

    positions = _mrna_mutation_positions(mrna, m.mutation_args())
    return ref_mrna_seq, var_mrna_seq, positions


def _transcription_result_from_dict(raw: dict, mut_id: str, gene: str) -> TranscriptionResult:
    """Convert a raw transcription prediction dict to a TranscriptionResult."""
    return TranscriptionResult(
        status=raw.get("status", "completed"),
        score=raw.get("regulatory_score", 0.0),
        tissue_fold_changes=raw.get("tissue_fold_changes", {}),
        tissue_predicted_levels=raw.get("tissue_predicted_levels", {}),
        max_lfc_tissue=raw.get("max_lfc_tissue", ""),
        max_lfc_value=raw.get("max_lfc_value", 0.0),
        max_abs_lfc=raw.get("max_abs_lfc", 0.0),
        num_affected_tissues=raw.get("num_affected_tissues", 0),
        mut_id=mut_id,
        gene=gene,
    )


def _run_folding_single(ref_transcript, m, mut_id, window_size):
    """Run folding analysis for a single mutation. Returns FoldingResult."""
    from ..folding.delta_mfe import compute_folding_profile as _cfp
    from ..folding.vienna import fold_available
    from ..folding.reference import complement

    tid = getattr(ref_transcript, "transcript_id", None)

    if not fold_available():
        return FoldingResult(status="error", reason="ViennaRNA not installed", mut_id=mut_id)

    mrna = ref_transcript.mature_mrna
    if mrna is None or not hasattr(mrna, "seq") or not hasattr(mrna, "index"):
        return FoldingResult(status="error", reason="no_mature_mrna", mut_id=mut_id)

    ref_seq = mrna.seq.upper()
    idx_list = mrna.index.tolist() if hasattr(mrna.index, "tolist") else list(mrna.index)
    pos_to_mrna = {gpos: i for i, gpos in enumerate(idx_list)}

    per_mut = []
    central_window = None
    for pos, ref, alt in m.mutation_args():
        mrna_pos = pos_to_mrna.get(pos)
        if mrna_pos is None:
            per_mut.append({"genomic_position": pos, "ref": ref, "alt": alt,
                            "status": "not_in_mrna", "delta_mfe": None})
            continue

        actual_ref = ref_seq[mrna_pos]
        coding_ref, coding_alt = ref.upper(), alt.upper()
        if actual_ref != coding_ref:
            coding_ref = complement(coding_ref)
            coding_alt = complement(coding_alt)
            if actual_ref != coding_ref:
                per_mut.append({"genomic_position": pos, "ref": ref, "alt": alt,
                                "status": "ref_mismatch", "delta_mfe": None,
                                "expected_ref": actual_ref})
                continue

        var_seq = ref_seq[:mrna_pos] + coding_alt + ref_seq[mrna_pos + 1:]
        try:
            profile = _cfp(ref_seq, var_seq, mrna_pos, window_size)

            # Build FoldingWindow from the best-centered window
            fw = None
            if profile.windows:
                best = min(profile.windows,
                           key=lambda w: abs(w.mutation_offset - len(w.ref_seq) // 2))
                fw = FoldingWindow(
                    ref_sequence=best.ref_seq,
                    var_sequence=best.var_seq,
                    ref_structure=best.ref_structure,
                    var_structure=best.var_structure,
                    ref_mfe=best.ref_mfe,
                    var_mfe=best.var_mfe,
                    delta_mfe=best.var_mfe - best.ref_mfe,
                    mutation_offset=best.mutation_offset,
                )
                if central_window is None:
                    central_window = fw

            # Build legacy-compatible dict entry
            entry = {"genomic_position": pos, "mrna_position": mrna_pos,
                     "ref": ref, "alt": alt, "coding_alt": coding_alt,
                     "status": "ok", "delta_mfe": profile.delta_mfe,
                     "score": profile.score}
            if fw:
                entry.update({
                    "ref_window": fw.ref_sequence,
                    "var_window": fw.var_sequence,
                    "ref_structure": fw.ref_structure,
                    "var_structure": fw.var_structure,
                    "mfe_ref": fw.ref_mfe,
                    "mfe_var": fw.var_mfe,
                    "mutation_offset": fw.mutation_offset,
                })
            per_mut.append(entry)
        except Exception as e:
            per_mut.append({"genomic_position": pos, "ref": ref, "alt": alt,
                            "status": "fold_error", "delta_mfe": None, "error": str(e)})

    first_delta = per_mut[0]["delta_mfe"] if per_mut else None
    first_score = next((e.get("score", 0.0) for e in per_mut if e.get("status") == "ok"), 0.0)

    return FoldingResult(
        status="completed",
        score=first_score or 0.0,
        delta_mfe=first_delta,
        mutations=per_mut,
        central_window=central_window,
        mut_id=mut_id,
        gene=m.gene,
        transcript_id=tid,
        window_size=window_size,
    )


# -- Shared mutation preparation ----------------------------------------------

def _prepare_mutation(
    mut_id: str,
    organism: str,
    transcript_id: str | None,
    context_bp: int,
    gene_load_timeout: float | None,
) -> dict:
    """Parse, load gene, select transcript, mutate, and build context.

    Returns a dict with either ``status='ok'`` and all prepared objects,
    or a short-circuit status (``'no_gene'`` / ``'no_transcript'``) suitable
    for building an empty MutationResult.

    This is the single source of truth for mutation setup — used by
    ``mutation_pipeline``, ``mutation_pipeline_batch``, and ``_concurrent_worker``.
    """
    from ..utils import _apply_mutation_safe

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

    gene = _load_gene_with_timeout(m.gene, organism, gene_load_timeout)
    if gene is None:
        return {"status": "no_gene", "event": m, "mut_id": mut_id}

    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return {"status": "no_transcript", "event": m, "mut_id": mut_id}

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

    mutated_transcript = selected.clone().generate_mature_mrna()
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(mutated_transcript.mature_mrna, pos, ref, alt)
        except Exception as e:
            logger.debug("Mutation apply failed (mature_mrna) %s pos=%d: %s", mut_id, pos, e)

    # Scoped pre-mRNA for mutation context (includes intronic sequence)
    ctx_transcript = selected.clone()
    ctx_transcript.generate_pre_mrna(
        region_start=central_pos - context_bp,
        region_end=central_pos + context_bp,
    )
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(ctx_transcript.pre_mrna, pos, ref, alt)
        except Exception as e:
            logger.debug("Mutation apply failed (pre_mrna ctx) %s pos=%d: %s", mut_id, pos, e)
    mutation_context = ctx_transcript.pre_mrna

    mutated_transcript.generate_protein()

    return {
        "status": "ok",
        "event": m,
        "mut_id": mut_id,
        "central_pos": central_pos,
        "gene": gene,
        "selected": selected,
        "ref_transcript": reference_transcript,
        "mut_transcript": mutated_transcript,
        "mutation_context": mutation_context,
    }
