import atexit
import copy
import logging
import threading
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


def _build_session_options(cfg, custom_ops_path: Optional[str], use_shared_allocator: bool):
    sess_options = ort.SessionOptions()
    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.arena_extend_strategy,
        cfg.arena_max_mem_bytes,
        cfg.arena_initial_chunk_bytes,
        cfg.cuda_arena_extend_strategy,
        cfg.cuda_gpu_mem_limit_bytes,
    )


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_options_list = _apply_cuda_provider_options(provider_names_list, provider_options_list, cfg)

        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

    if provider_names_list:
        LOGGER.debug(
            "Created ORT session (model=%s, providers=%s, allocator=%s)",
            model_path,
            provider_names_list,
            "shared" if shared_allocator_enabled else "per-session",
        )
    return session


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


atexit.register(clear_cache)
