import atexit
import copy
import logging
import os
import threading
from pathlib import Path
from typing import Any, Dict, Optional, Sequence, Tuple

import onnxruntime as ort

from highlighter.core.config import HighlighterRuntimeConfig
from highlighter.core.shared_cache import SharedResourceCache

LOGGER = logging.getLogger(__name__)

_ALLOC_LOCK = threading.RLock()
_SESSION_CACHE: SharedResourceCache[ort.InferenceSession] = SharedResourceCache(maxsize=0)
_ALLOC_STATE: Dict[str, Any] = {
    "initialized": False,
    "enabled": False,
    "signature": None,
}


def _freeze(value: Any) -> Any:
    if isinstance(value, dict):
        return tuple((k, _freeze(v)) for k, v in sorted(value.items()))
    if isinstance(value, (list, tuple)):
        return tuple(_freeze(v) for v in value)
    if isinstance(value, set):
        return tuple(sorted(_freeze(v) for v in value))
    try:
        hash(value)
        return value
    except TypeError:
        return repr(value)


def _split_providers_and_options(
    providers: Sequence[Any], provider_options: Optional[Sequence[Any]]
) -> Tuple[Sequence[str], Optional[Sequence[Any]]]:
    if providers and isinstance(providers[0], (list, tuple)) and len(providers[0]) == 2:
        provider_names = [str(p[0]) for p in providers]
        provider_opts = [p[1] for p in providers]
        return provider_names, provider_opts
    return [str(p) for p in providers], provider_options


def _normalize_provider_options(
    providers: Sequence[str], provider_options: Optional[Sequence[Any]]
) -> list[dict[str, Any]]:
    """Normalize provider options to match providers length (ORT requires this)."""
    num_providers = len(providers)
    if num_providers == 0:
        return []

    if provider_options is None:
        opts: list[Any] = [{} for _ in range(num_providers)]
    elif isinstance(provider_options, dict):
        opts = [provider_options]
    else:
        opts = list(provider_options)

    if len(opts) < num_providers:
        opts.extend({} for _ in range(num_providers - len(opts)))
    return [dict(opt or {}) for opt in opts[:num_providers]]


def _normalize_provider_options_for_key(
    provider_names: Sequence[str], provider_options: Optional[Sequence[Any]]
) -> Tuple[Tuple[str, ...], Tuple[Any, ...]]:
    provider_names_tuple = tuple(provider_names) if provider_names else tuple()
    normalized = _normalize_provider_options(provider_names, provider_options)
    normalized_options = tuple(_freeze(opt) for opt in normalized)
    return provider_names_tuple, normalized_options


def get_ort_options() -> Optional[ort.RunOptions]:
    cfg = copy.deepcopy(HighlighterRuntimeConfig.load().onnxruntime)
    if not cfg.arena_shrinkage:
        return None

    run_options = ort.RunOptions()
    try:
        run_options.add_run_config_entry(
            "memory.enable_memory_arena_shrinkage",
            str(cfg.arena_shrinkage),
        )
    except Exception:
        LOGGER.warning("Unable to set arena shrinkage RunOptions entry.")
        return None
    return run_options


def _arena_cfg_signature(cfg) -> Tuple[Any, ...]:
    return (
        cfg.arena_max_mem_bytes,
        cfg.arena_extend_strategy,
        cfg.arena_initial_chunk_bytes,
    )


def _build_arena_cfg(cfg):
    if not hasattr(ort, "OrtArenaCfg"):
        return None
    if (
        cfg.arena_max_mem_bytes is None
        and cfg.arena_initial_chunk_bytes is None
        and cfg.arena_extend_strategy is None
    ):
        return None

    arena_extend_strategy = int(cfg.arena_extend_strategy or 0)
    max_mem = int(cfg.arena_max_mem_bytes or 0)
    initial_chunk = int(cfg.arena_initial_chunk_bytes or 0)
    max_dead_bytes_per_chunk = 0
    try:
        return ort.OrtArenaCfg(max_mem, arena_extend_strategy, initial_chunk, max_dead_bytes_per_chunk)
    except TypeError:
        LOGGER.warning("OrtArenaCfg signature mismatch; skipping arena config.")
        return None


def _ensure_shared_allocator(cfg) -> bool:
    if not cfg.shared_allocator:
        return False

    with _ALLOC_LOCK:
        signature = _arena_cfg_signature(cfg)
        if _ALLOC_STATE["initialized"]:
            if _ALLOC_STATE["enabled"] and signature != _ALLOC_STATE["signature"]:
                LOGGER.warning(
                    "Shared allocator already initialized with a different config; "
                    "ignoring new allocator settings."
                )
            return _ALLOC_STATE["enabled"]

        if not hasattr(ort, "create_and_register_allocator"):
            LOGGER.warning("onnxruntime.create_and_register_allocator not available; skipping.")
            _ALLOC_STATE["initialized"] = True
            _ALLOC_STATE["enabled"] = False
            _ALLOC_STATE["signature"] = signature
            return False

        arena_cfg = _build_arena_cfg(cfg)
        LOGGER.info(
            "Creating shared allocator (max_mem=%s, extend_strategy=%s, initial_chunk=%s)",
            cfg.arena_max_mem_bytes,
            cfg.arena_extend_strategy,
            cfg.arena_initial_chunk_bytes,
        )
        try:
            mem_info = ort.OrtMemoryInfo(
                "Cpu",
                ort.OrtAllocatorType.ORT_ARENA_ALLOCATOR,
                0,
                ort.OrtMemType.DEFAULT,
            )
            ort.create_and_register_allocator(mem_info, arena_cfg)
            _ALLOC_STATE["initialized"] = True
            _ALLOC_STATE["enabled"] = True
            _ALLOC_STATE["signature"] = signature
            return True
        except Exception:
            LOGGER.exception("Failed to create shared allocator; falling back to per-session.")
            _ALLOC_STATE["initialized"] = True
            _ALLOC_STATE["enabled"] = False
            _ALLOC_STATE["signature"] = signature
            return False


def _apply_cuda_provider_options(
    providers: Sequence[str], provider_options: list[dict[str, Any]], cfg
) -> list[dict[str, Any]]:
    """Apply CUDA config settings. Expects pre-normalized provider_options."""
    if cfg.cuda_arena_extend_strategy is None and cfg.cuda_gpu_mem_limit_bytes is None:
        return provider_options

    for idx, provider in enumerate(providers):
        if provider != "CUDAExecutionProvider":
            continue
        opts = dict(provider_options[idx])
        if cfg.cuda_arena_extend_strategy is not None:
            opts["arena_extend_strategy"] = cfg.cuda_arena_extend_strategy
        if cfg.cuda_gpu_mem_limit_bytes is not None:
            opts["gpu_mem_limit"] = int(cfg.cuda_gpu_mem_limit_bytes)
        provider_options[idx] = opts

    return provider_options


_MIN_TRT_COMPUTE_CAPABILITY = 7.5


def _gpu_supports_tensorrt() -> bool:
    """Check if the current GPU supports TensorRT (compute capability >= 7.5)."""
    try:
        import torch

        if not torch.cuda.is_available():
            return False
        capability = torch.cuda.get_device_capability()
        cc = float(f"{capability[0]}.{capability[1]}")
        if cc < _MIN_TRT_COMPUTE_CAPABILITY:
            LOGGER.info(
                "GPU compute capability %.1f < %.1f; TensorRT EP will be skipped.",
                cc,
                _MIN_TRT_COMPUTE_CAPABILITY,
            )
            return False
        return True
    except Exception:
        LOGGER.debug("Could not determine GPU compute capability; skipping TensorRT EP.")
        return False


def _filter_tensorrt_if_unsupported(
    providers: list[str], provider_options: list[dict[str, Any]]
) -> tuple[list[str], list[dict[str, Any]]]:
    """Remove TensorRT EP from provider list if the GPU doesn't support it."""
    if "TensorrtExecutionProvider" not in providers:
        return providers, provider_options
    if _gpu_supports_tensorrt():
        return providers, provider_options
    filtered = [(p, o) for p, o in zip(providers, provider_options) if p != "TensorrtExecutionProvider"]
    return [p for p, _ in filtered], [o for _, o in filtered]


def _apply_tensorrt_provider_options(
    providers: Sequence[str], provider_options: list[dict[str, Any]], cfg
) -> list[dict[str, Any]]:
    """Apply TensorRT config settings. Expects pre-normalized provider_options."""
    for idx, provider in enumerate(providers):
        if provider != "TensorrtExecutionProvider":
            continue
        opts = dict(provider_options[idx])
        opts["trt_fp16_enable"] = "True" if cfg.trt_fp16_enable else "False"
        opts["trt_engine_cache_enable"] = "True" if cfg.trt_engine_cache_enable else "False"
        cache_path = cfg.trt_engine_cache_path or str(Path.home() / ".cache" / "highlighter" / "trt_engines")
        os.makedirs(cache_path, exist_ok=True)
        opts["trt_engine_cache_path"] = cache_path
        provider_options[idx] = opts

    return provider_options


def _build_session_options(cfg, custom_ops_path: Optional[str], use_shared_allocator: bool):
    sess_options = ort.SessionOptions()
    sess_options.intra_op_num_threads = int(cfg.intra_op_num_threads)
    sess_options.inter_op_num_threads = int(cfg.inter_op_num_threads)
    if custom_ops_path:
        sess_options.register_custom_ops_library(str(custom_ops_path))

    if hasattr(sess_options, "enable_cpu_mem_arena"):
        sess_options.enable_cpu_mem_arena = bool(cfg.enable_cpu_mem_arena)
    else:
        LOGGER.warning("SessionOptions.enable_cpu_mem_arena not available; skipping.")

    if cfg.use_global_threadpool:
        if hasattr(sess_options, "use_per_session_threads"):
            sess_options.use_per_session_threads = False
        else:
            LOGGER.warning("SessionOptions.use_per_session_threads not available; skipping.")

    if use_shared_allocator:
        try:
            sess_options.add_session_config_entry("session.use_env_allocators", "1")
        except Exception:
            LOGGER.warning("Unable to enable env allocators for shared allocator.")

    return sess_options


def _cache_key(
    model_path: str,
    providers: Sequence[Any],
    provider_options: Optional[Sequence[Any]],
    custom_ops_path: Optional[str],
    cfg,
) -> Tuple[Any, ...]:
    provider_names, provider_opts = _split_providers_and_options(providers, provider_options)
    _, normalized_options = _normalize_provider_options_for_key(provider_names, provider_opts)
    return (
        str(model_path),
        tuple(provider_names),
        normalized_options,
        str(custom_ops_path) if custom_ops_path else None,
        cfg.shared_allocator,
        cfg.enable_cpu_mem_arena,
        cfg.use_global_threadpool,
        cfg.intra_op_num_threads,
        cfg.inter_op_num_threads,
        cfg.arena_extend_strategy,
        cfg.arena_max_mem_bytes,
        cfg.arena_initial_chunk_bytes,
        cfg.cuda_arena_extend_strategy,
        cfg.cuda_gpu_mem_limit_bytes,
        cfg.trt_fp16_enable,
        cfg.trt_engine_cache_enable,
        cfg.trt_engine_cache_path,
    )


def get_session(
    model_path: str,
    *,
    providers: Sequence[Any],
    provider_options: Optional[Sequence[Any]] = None,
    custom_ops_path: Optional[str] = None,
) -> ort.InferenceSession:
    cfg = copy.deepcopy(HighlighterRuntimeConfig.load().onnxruntime)

    key = _cache_key(model_path, providers, provider_options, custom_ops_path, cfg)
    provider_names_list: list[str] = []  # Populated only by the creator thread for logging.
    shared_allocator_enabled = False

    def _create_session() -> ort.InferenceSession:
        nonlocal shared_allocator_enabled, provider_names_list
        shared_allocator_enabled = _ensure_shared_allocator(cfg)
        sess_options = _build_session_options(cfg, custom_ops_path, shared_allocator_enabled)

        provider_names_list, provider_options_raw = _split_providers_and_options(providers, provider_options)
        provider_options_list = _normalize_provider_options(provider_names_list, provider_options_raw)
        provider_names_list, provider_options_list = _filter_tensorrt_if_unsupported(
            provider_names_list, provider_options_list
        )
        provider_options_list = _apply_cuda_provider_options(provider_names_list, provider_options_list, cfg)
        provider_options_list = _apply_tensorrt_provider_options(
            provider_names_list, provider_options_list, cfg
        )

        LOGGER.info(
            "ort session init model=%s providers=%s allocator=%s custom_ops_path=%s",
            model_path,
            provider_names_list,
            "shared" if shared_allocator_enabled else "per-session",
            custom_ops_path,
        )
        LOGGER.debug(
            "ort session config provider_options=%s shared_allocator_requested=%s enable_cpu_mem_arena=%s "
            "use_global_threadpool=%s arena_extend_strategy=%s arena_max_mem_bytes=%s "
            "arena_initial_chunk_bytes=%s arena_shrinkage=%s cuda_arena_extend_strategy=%s "
            "cuda_gpu_mem_limit_bytes=%s trt_fp16_enable=%s trt_engine_cache_enable=%s "
            "trt_engine_cache_path=%s session_cache_maxsize=%s",
            provider_options_list,
            cfg.shared_allocator,
            cfg.enable_cpu_mem_arena,
            cfg.use_global_threadpool,
            cfg.arena_extend_strategy,
            cfg.arena_max_mem_bytes,
            cfg.arena_initial_chunk_bytes,
            cfg.arena_shrinkage,
            cfg.cuda_arena_extend_strategy,
            cfg.cuda_gpu_mem_limit_bytes,
            cfg.trt_fp16_enable,
            cfg.trt_engine_cache_enable,
            cfg.trt_engine_cache_path,
            cfg.session_cache_maxsize,
        )

        return ort.InferenceSession(
            str(model_path),
            sess_options=sess_options,
            providers=provider_names_list,
            provider_options=provider_options_list,
        )

    try:
        session = _SESSION_CACHE.get_or_create(
            key,
            _create_session,
            maxsize=int(cfg.session_cache_maxsize),
        )
    except Exception:
        LOGGER.exception(
            "Failed to create InferenceSession for model_path=%s providers=%s",
            model_path,
            providers,
        )
        raise

    return session


def clear_cache() -> None:
    _SESSION_CACHE.clear()


atexit.register(clear_cache)
