# CelerisLab/cuda/context.py
"""
CUDA device selection, context lifecycle, and device info queries.
"""

import subprocess
from typing import Optional

import pycuda.driver as cuda


class CudaContext:
    """Manages a single CUDA device context.

    Usage::

        ctx = CudaContext(device_id=0)
        # ... use PyCUDA ...
        ctx.close()

    Or as a context-manager::

        with CudaContext(0) as ctx:
            ...
    """

    def __init__(self, device_id: int = 0):
        cuda.init()
        if device_id < 0 or device_id >= cuda.Device.count():
            raise ValueError(
                f"device_id {device_id} invalid; "
                f"{cuda.Device.count()} device(s) available."
            )
        self.device_id = device_id
        self.device = cuda.Device(device_id)
        self._ctx = self.device.make_context()

    # -- properties ----------------------------------------------------------
    @property
    def name(self) -> str:
        return self.device.name()

    @property
    def compute_capability(self) -> str:
        cc = self.device.compute_capability()
        return f"{cc[0]}.{cc[1]}"

    @property
    def sm_arch(self) -> str:
        """Return 'sm_XX' string for nvcc -arch."""
        cc = self.device.compute_capability()
        return f"sm_{cc[0]}{cc[1]}"

    @property
    def total_memory(self) -> int:
        return self.device.total_memory()

    # -- lifecycle -----------------------------------------------------------
    def close(self):
        if self._ctx is not None:
            self._ctx.pop()
            self._ctx = None

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    def __del__(self):
        self.close()


def detect_compute_capability(device_id: int = 0) -> str:
    """Return 'sm_XX' for the given device without creating a persistent context."""
    cuda.init()
    dev = cuda.Device(device_id)
    cc = dev.compute_capability()
    return f"sm_{cc[0]}{cc[1]}"
