from __future__ import annotations

import ctypes
import os
import sys
from pathlib import Path
from typing import Sequence

import numpy as np

from .io import ensure_time_axis, percentile_normalize

_SESSION_CACHE: dict[tuple[str, tuple[str, ...]], object] = {}
_ORT_DLL_DIRS_READY = False
_ORT_DLL_DIR_HANDLES: list[object] = []


def _onnxruntime_windows_candidate_dirs() -> list[Path]:
    candidates: list[Path] = []
    meipass = getattr(sys, "_MEIPASS", None)
    if meipass:
        candidates.append(Path(meipass) / "onnxruntime" / "capi")
        candidates.append(Path(meipass) / "onnxruntime.libs")
        candidates.append(Path(meipass))

    exe_parent = Path(sys.executable).resolve().parent
    candidates.append(exe_parent / "_internal" / "onnxruntime" / "capi")
    candidates.append(exe_parent / "_internal" / "onnxruntime.libs")
    candidates.append(exe_parent / "onnxruntime" / "capi")
    candidates.append(exe_parent / "onnxruntime.libs")
    return candidates


def _prepare_onnxruntime_windows_dll_dirs() -> None:
    global _ORT_DLL_DIRS_READY
    if _ORT_DLL_DIRS_READY or sys.platform != "win32":
        return

    seen: set[str] = set()
    valid_dirs: list[str] = []
    for path in _onnxruntime_windows_candidate_dirs():
        try:
            resolved = str(path.resolve())
        except Exception:
            resolved = str(path)
        norm = os.path.normcase(resolved)
        if norm in seen or not path.exists():
            continue
        seen.add(norm)
        valid_dirs.append(resolved)

    if valid_dirs:
        existing_path = os.environ.get("PATH", "")
        os.environ["PATH"] = os.pathsep.join(valid_dirs + ([existing_path] if existing_path else []))
        add_dll_directory = getattr(os, "add_dll_directory", None)
        if add_dll_directory is not None:
            for dll_dir in valid_dirs:
                try:
                    _ORT_DLL_DIR_HANDLES.append(add_dll_directory(dll_dir))
                except OSError:
                    pass

    _ORT_DLL_DIRS_READY = True


def _onnxruntime_windows_debug_context() -> str:
    if sys.platform != "win32":
        return ""
    parts: list[str] = []
    probe_candidates: list[Path] = []
    for path in _onnxruntime_windows_candidate_dirs():
        exists = path.exists()
        parts.append(f"{path} exists={exists}")
        if exists:
            try:
                dlls = sorted(p.name for p in path.glob("*.dll"))
            except Exception:
                dlls = []
            try:
                pyds = sorted(p.name for p in path.glob("*.pyd"))
            except Exception:
                pyds = []
            if dlls:
                parts.append("  dlls=" + ", ".join(dlls[:12]))
            if pyds:
                parts.append("  pyds=" + ", ".join(pyds[:12]))
            for name in (
                "onnxruntime.dll",
                "onnxruntime_providers_cpu.dll",
                "onnxruntime_providers_shared.dll",
                "onnxruntime_pybind11_state.pyd",
                "vcruntime140.dll",
                "vcruntime140_1.dll",
                "msvcp140.dll",
                "msvcp140_1.dll",
                "msvcp140_atomic_wait.dll",
                "vcomp140.dll",
                "concrt140.dll",
                "libiomp5md.dll",
                "dnnl.dll",
                "tbb12.dll",
                "tbbmalloc.dll",
            ):
                candidate = path / name
                if candidate.exists():
                    probe_candidates.append(candidate)
    seen: set[str] = set()
    probe_results: list[str] = []
    for candidate in probe_candidates:
        key = str(candidate).lower()
        if key in seen:
            continue
        seen.add(key)
        try:
            ctypes.WinDLL(str(candidate))
            probe_results.append(f"{candidate.name}=ok")
        except Exception as exc:
            probe_results.append(f"{candidate.name}=FAIL({exc})")
    if probe_results:
        parts.append("probes=" + ", ".join(probe_results))
    return " | ".join(parts)


def _default_providers(ort) -> list[str]:
    available = set(ort.get_available_providers())
    preferred = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    return [p for p in preferred if p in available] or ["CPUExecutionProvider"]


def _parse_positive_int_env(name: str) -> int:
    raw = os.environ.get(name, "").strip()
    if not raw:
        return 0
    try:
        val = int(raw)
    except Exception:
        return 0
    return val if val > 0 else 0


def _session_options_from_env(ort):
    # Explicitly setting thread counts disables ORT's default thread affinity.
    intra = _parse_positive_int_env("TRACY_ORT_INTRA_THREADS")
    if intra <= 0:
        intra = _parse_positive_int_env("TRACY_ORT_THREADS")
    inter = _parse_positive_int_env("TRACY_ORT_INTER_THREADS")

    if intra <= 0 and inter <= 0:
        return None

    opts = ort.SessionOptions()
    if intra > 0:
        opts.intra_op_num_threads = int(intra)
    if inter > 0:
        opts.inter_op_num_threads = int(inter)
    return opts


def get_onnx_session(onnx_path: str | Path, providers: Sequence[str] | None = None):
    _prepare_onnxruntime_windows_dll_dirs()
    try:
        import onnxruntime as ort  # type: ignore
    except Exception as exc:
        detail = _onnxruntime_windows_debug_context()
        if detail:
            raise RuntimeError(f"onnxruntime import failed for ONNX inference: {exc} [{detail}]") from exc
        raise RuntimeError(f"onnxruntime import failed for ONNX inference: {exc}") from exc

    path = str(Path(onnx_path).expanduser().resolve())
    providers_tuple = tuple(providers) if providers else tuple(_default_providers(ort))
    key = (path, providers_tuple)
    sess = _SESSION_CACHE.get(key)
    if sess is None:
        sess_opts = _session_options_from_env(ort)
        sess = ort.InferenceSession(path, sess_options=sess_opts, providers=list(providers_tuple))
        _SESSION_CACHE[key] = sess
    return sess


def _run_outputs_on_tx_norm(
    tx_norm: np.ndarray,
    onnx_path: str | Path,
    *,
    providers: Sequence[str] | None = None,
    align_factor: int = 8,
) -> list[np.ndarray]:
    if tx_norm.ndim != 2:
        raise ValueError(f"Expected normalized (T, X) image, got shape {tx_norm.shape!r}")
    tx_norm = np.asarray(tx_norm, dtype=np.float32)
    orig_t, orig_x = tx_norm.shape

    pad_t = 0
    pad_x = 0
    if align_factor and align_factor > 1:
        pad_t = int((align_factor - (orig_t % align_factor)) % align_factor)
        pad_x = int((align_factor - (orig_x % align_factor)) % align_factor)
    if pad_t or pad_x:
        tx_norm = np.pad(tx_norm, ((0, pad_t), (0, pad_x)), mode="edge")

    inp = tx_norm[None, None, :, :].astype(np.float32)

    session = get_onnx_session(onnx_path, providers=providers)
    input_name = session.get_inputs()[0].name
    raw_outs = session.run(None, {input_name: inp})

    outs: list[np.ndarray] = []
    for out in raw_outs:
        arr = np.asarray(out)
        if arr.ndim != 4:
            raise ValueError(f"Expected ONNX output with shape [N,C,T,X], got {arr.shape!r}")
        chw = np.asarray(arr[0], dtype=np.float32)
        if pad_t or pad_x:
            chw = chw[:, :orig_t, :orig_x]
        outs.append(chw)
    return outs


def _run_logits_on_tx_norm(
    tx_norm: np.ndarray,
    onnx_path: str | Path,
    *,
    providers: Sequence[str] | None = None,
    align_factor: int = 8,
) -> np.ndarray:
    outs = _run_outputs_on_tx_norm(tx_norm, onnx_path, providers=providers, align_factor=align_factor)
    if not outs:
        raise RuntimeError("ONNX model returned no outputs")
    return outs[0]


def infer_logits_onnx(
    kymo_array: np.ndarray,
    onnx_path: str | Path,
    *,
    time_axis: str = "y",
    providers: Sequence[str] | None = None,
    align_factor: int = 8,
) -> np.ndarray:
    arr = np.asarray(kymo_array)
    if arr.ndim != 2:
        raise ValueError(f"Expected 2D kymograph input, got shape {arr.shape!r}")

    tx = ensure_time_axis(arr, time_axis=time_axis)
    tx_norm = percentile_normalize(tx)
    return _run_logits_on_tx_norm(tx_norm, onnx_path, providers=providers, align_factor=align_factor)


def infer_logits_and_embeddings_onnx(
    kymo_array: np.ndarray,
    onnx_path: str | Path,
    *,
    time_axis: str = "y",
    providers: Sequence[str] | None = None,
    align_factor: int = 8,
) -> tuple[np.ndarray, np.ndarray | None]:
    """Return segmentation logits and optional embedding map from ONNX.

    Returns:
        logits: [C, T, X]
        embeddings: [D, T, X] or None when model has a single output
    """

    arr = np.asarray(kymo_array)
    if arr.ndim != 2:
        raise ValueError(f"Expected 2D kymograph input, got shape {arr.shape!r}")

    tx = ensure_time_axis(arr, time_axis=time_axis)
    tx_norm = percentile_normalize(tx)
    outs = _run_outputs_on_tx_norm(tx_norm, onnx_path, providers=providers, align_factor=align_factor)
    if not outs:
        raise RuntimeError("ONNX model returned no outputs")
    logits = outs[0]
    embeddings = outs[1] if len(outs) > 1 else None
    return logits, embeddings


def softmax_np(logits: np.ndarray, axis: int = 0) -> np.ndarray:
    shifted = logits - np.max(logits, axis=axis, keepdims=True)
    ex = np.exp(shifted)
    return ex / np.sum(ex, axis=axis, keepdims=True)


def infer_prob_move_onnx(
    kymo_array: np.ndarray,
    onnx_path: str | Path,
    *,
    time_axis: str = "y",
    moving_class: int = 1,
    providers: Sequence[str] | None = None,
    dual_polarity: bool = False,
    dark_weight: float = 1.0,
) -> np.ndarray:
    arr = np.asarray(kymo_array)
    if arr.ndim != 2:
        raise ValueError(f"Expected 2D kymograph input, got shape {arr.shape!r}")

    tx = ensure_time_axis(arr, time_axis=time_axis)
    tx_norm = percentile_normalize(tx)

    logits = _run_logits_on_tx_norm(tx_norm, onnx_path, providers=providers)
    probs = softmax_np(logits, axis=0)
    if not (0 <= moving_class < probs.shape[0]):
        raise ValueError(f"moving_class={moving_class} out of range for logits with {probs.shape[0]} classes")
    prob_move = probs[moving_class].astype(np.float32)

    if dual_polarity:
        tx_inv = 1.0 - tx_norm
        logits_inv = _run_logits_on_tx_norm(tx_inv, onnx_path, providers=providers)
        probs_inv = softmax_np(logits_inv, axis=0)
        prob_move_inv = probs_inv[moving_class].astype(np.float32)
        w = float(max(0.0, dark_weight))
        if w == 0.0:
            pass
        elif w >= 1.0:
            prob_move = np.maximum(prob_move, prob_move_inv)
        else:
            prob_move = np.maximum(prob_move, w * prob_move_inv + (1.0 - w) * prob_move)

    # Return in same orientation as input.
    if (time_axis or "y").lower() == "x":
        prob_move = prob_move.T
    return prob_move.astype(np.float32)


def infer_prob_move_and_embeddings_onnx(
    kymo_array: np.ndarray,
    onnx_path: str | Path,
    *,
    time_axis: str = "y",
    moving_class: int = 1,
    providers: Sequence[str] | None = None,
) -> tuple[np.ndarray, np.ndarray | None]:
    """Return moving-class probability and optional embedding map.

    Returns:
        prob_move: [T, X] (same orientation as input)
        embeddings: [D, T, X] in input orientation (or None)
    """

    logits, embeddings = infer_logits_and_embeddings_onnx(
        kymo_array,
        onnx_path,
        time_axis=time_axis,
        providers=providers,
    )
    probs = softmax_np(logits, axis=0)
    if not (0 <= moving_class < probs.shape[0]):
        raise ValueError(f"moving_class={moving_class} out of range for logits with {probs.shape[0]} classes")
    prob_move = probs[moving_class].astype(np.float32)

    if (time_axis or "y").lower() == "x":
        prob_move = prob_move.T
        if embeddings is not None:
            embeddings = np.transpose(np.asarray(embeddings, dtype=np.float32), (0, 2, 1))
    elif embeddings is not None:
        embeddings = np.asarray(embeddings, dtype=np.float32)

    return prob_move.astype(np.float32), embeddings
