from __future__ import annotations

__all__ = ["resolve_device", "gpu_available", "mps_available"]

from typing import Optional, Any
import cv2
from .enums import ComputeDevice

# Синонимы (регистр не важен)
_ALIASES: dict[str, ComputeDevice] = {
    "cpu":   ComputeDevice.CPU,
    "gpu":   ComputeDevice.GPU,
    "cuda":  ComputeDevice.GPU,
    "cudo":  ComputeDevice.GPU,
    "nvidia":ComputeDevice.GPU,
    "mps":   ComputeDevice.MPS,
    "metal": ComputeDevice.MPS,
}

def _to_key(obj: Any) -> str:
    """
        Превращает что угодно в ключ для _ALIASES:
        - Enum с .value -> используем value
        - строка -> как есть
        - Enum без .value -> берём последний токен из str(obj), напр. 'ComputeDevice.GPU' -> 'gpu'
        - всё остальное -> str(obj).lower()
    """
    # строка
    if isinstance(obj, str):
        return obj.strip().lower()

    # наш Enum
    if isinstance(obj, ComputeDevice):
        return obj.value

    # похож на Enum (другой класс/другая копия)
    val = getattr(obj, "value", None)
    if isinstance(val, str):
        return val.strip().lower()

    # fallback: пытаемся вытащить 'GPU' из 'ComputeDevice.GPU'
    s = str(obj).strip().lower()
    if "." in s:
        s = s.split(".")[-1]
    return s

def resolve_device(device: Optional[str | ComputeDevice], default: ComputeDevice) -> ComputeDevice:
    """
        Нормализует устройство. Поддерживает синонимы и любой регистр.
        Примеры: "gpu", "GPU", "cuda", "CUDA", "cpu", "mps", "Metal", а также любые Enum'ы.
    """
    if device is None:
        return default

    key = _to_key(device)
    if key in _ALIASES:
        return _ALIASES[key]

    allowed = ", ".join(sorted(_ALIASES.keys()))
    raise ValueError(f"Unknown device '{device}'. Allowed: {allowed}")

def gpu_available() -> bool:
    """ Доступна ли CUDA (cv2.cuda). """
    try:
        return hasattr(cv2, "cuda") and cv2.cuda.getCudaEnabledDeviceCount() > 0
    except Exception:
        return False

def mps_available() -> bool:
    """ Заглушка под Apple Metal (MPS). Пока всегда False → фолбэк на CPU. """
    return False
