"""GPU Manager - Singleton for CUDA detection and memory management."""

import logging
from threading import Lock
from typing import Optional

logger = logging.getLogger(__name__)


class GPUManager:
    """Singleton GPU manager for CUDA detection and memory pooling."""

    _instance: Optional["GPUManager"] = None
    _lock = Lock()

    def __init__(self) -> None:
        """Initialize GPU manager (use get_instance() instead)."""
        if GPUManager._instance is not None:
            raise RuntimeError("Use GPUManager.get_instance() instead of direct instantiation")

        self._cuda_available = False
        self._device_count = 0
        self._memory_pool: Optional[object] = None
        self._gpu_memory_fraction = 0.8

        self._detect_cuda()

    @classmethod
    def get_instance(cls) -> "GPUManager":
        """Get singleton instance of GPUManager."""
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = cls()
        return cls._instance

    def _detect_cuda(self) -> None:
        """Detect CUDA availability and initialize memory pool."""
        try:
            import cupy as cp

            self._cuda_available = True
            self._device_count = cp.cuda.runtime.getDeviceCount()

            # Initialize memory pool with 80% of GPU memory
            pool = cp.cuda.MemoryPool()
            total_memory = cp.cuda.Device().mem_info[1]
            pool.set_limit(size=int(total_memory * self._gpu_memory_fraction))
            cp.cuda.set_allocator(pool.malloc)
            self._memory_pool = pool

            logger.info(
                f"CUDA detected: {self._device_count} device(s), "
                f"memory pool: {total_memory * self._gpu_memory_fraction / 1e9:.2f} GB"
            )
        except ImportError:
            logger.info("CuPy not available - GPU acceleration disabled")
        except Exception as e:
            logger.warning(f"CUDA detection failed: {e} - falling back to CPU")

    @property
    def cuda_available(self) -> bool:
        """Check if CUDA is available."""
        return self._cuda_available

    @property
    def device_count(self) -> int:
        """Get number of CUDA devices."""
        return self._device_count

    def set_memory_fraction(self, fraction: float) -> None:
        """Set GPU memory pool fraction (0.0-1.0)."""
        if not 0.0 < fraction <= 1.0:
            raise ValueError("Memory fraction must be between 0.0 and 1.0")

        self._gpu_memory_fraction = fraction

        if self._cuda_available and self._memory_pool is not None:
            import cupy as cp

            total_memory = cp.cuda.Device().mem_info[1]
            self._memory_pool.set_limit(size=int(total_memory * fraction))
            logger.info(f"GPU memory pool updated to {fraction * 100}%")

    def get_memory_info(self) -> dict[str, int]:
        """Get current GPU memory usage."""
        if not self._cuda_available:
            return {"total": 0, "free": 0, "used": 0}

        try:
            import cupy as cp

            free, total = cp.cuda.Device().mem_info
            return {"total": total, "free": free, "used": total - free}
        except Exception as e:
            logger.error(f"Failed to get memory info: {e}")
            return {"total": 0, "free": 0, "used": 0}

    def clear_memory_pool(self) -> None:
        """Clear GPU memory pool."""
        if self._cuda_available and self._memory_pool is not None:
            self._memory_pool.free_all_blocks()
            logger.info("GPU memory pool cleared")
