from __future__ import annotations

import platform
import socket
import sys
import uuid
from collections.abc import Callable
from contextlib import suppress
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, TypeVar

import psutil  # type: ignore[import-untyped]

from modelforge_api import __version__
from modelforge_api.domain.enums import Availability
from modelforge_api.domain.hardware import (
    AcceleratorInventory,
    AcceleratorTelemetry,
    HostInventory,
    NvidiaCollection,
    ObservedValue,
    StorageObservation,
)

T = TypeVar("T")


class NodeIdentityProvider:
    """Stable node identity: explicit override, OS machine ID, then persisted UUID."""

    def __init__(
        self,
        identity_file: Path,
        explicit_identity: str | None = None,
        force_persisted: bool = False,
    ) -> None:
        self.identity_file = identity_file
        self.explicit_identity = explicit_identity
        self.force_persisted = force_persisted

    def resolve(self) -> tuple[str, str]:
        if self.explicit_identity:
            return self.explicit_identity, "configured"
        if not self.force_persisted:
            system_id = self._system_machine_id()
            if system_id:
                return str(uuid.uuid5(uuid.NAMESPACE_OID, system_id)), "os_machine_id"
        try:
            if self.identity_file.exists():
                return self.identity_file.read_text(encoding="utf-8").strip(), "persisted_uuid"
            value = str(uuid.uuid4())
            self.identity_file.parent.mkdir(parents=True, exist_ok=True)
            self.identity_file.write_text(value, encoding="utf-8")
            return value, "persisted_uuid"
        except OSError as exc:
            raise RuntimeError("unable to establish stable node identity") from exc

    @staticmethod
    def _system_machine_id() -> str | None:
        # `sys.platform`, not `os.name`: a type checker narrows on the former and not on the latter,
        # so under `os.name` this block was analysed on Linux too — where winreg has no attributes.
        # The images run on Linux, so the platform the release actually ships on was the one the
        # type check could not see. They mean the same thing at runtime.
        if sys.platform == "win32":
            try:
                import winreg

                with winreg.OpenKey(
                    winreg.HKEY_LOCAL_MACHINE,
                    r"SOFTWARE\Microsoft\Cryptography",
                ) as key:
                    return str(winreg.QueryValueEx(key, "MachineGuid")[0])
            except (OSError, ImportError):
                return None
        for path in (Path("/etc/machine-id"), Path("/var/lib/dbus/machine-id")):
            try:
                value = path.read_text(encoding="utf-8").strip()
                if value:
                    return value
            except OSError:
                continue
        return None


class SystemHostCollector:
    def __init__(
        self,
        identity_provider: NodeIdentityProvider,
        storage_paths: dict[str, Path],
    ) -> None:
        self.identity_provider = identity_provider
        self.storage_paths = storage_paths

    @staticmethod
    def _known_or_unknown(value: T | None, reason: str) -> ObservedValue[T]:
        return (
            ObservedValue.known(value)
            if value is not None
            else ObservedValue.absent(Availability.UNKNOWN, reason)
        )

    def collect(self) -> HostInventory:
        identity_key, identity_source = self.identity_provider.resolve()
        memory = psutil.virtual_memory()
        cpu_model = platform.processor().strip() or None
        physical = psutil.cpu_count(logical=False)
        logical = psutil.cpu_count(logical=True)
        storage: list[StorageObservation] = []
        for purpose, path in self.storage_paths.items():
            try:
                usage = psutil.disk_usage(str(path))
                storage.append(
                    StorageObservation(
                        purpose=purpose,
                        path=str(path.resolve()),
                        total_bytes=ObservedValue.known(usage.total),
                        used_bytes=ObservedValue.known(usage.used),
                        free_bytes=ObservedValue.known(usage.free),
                    )
                )
            except OSError as exc:
                reason = f"{type(exc).__name__}: path unavailable"
                unavailable = ObservedValue[int].absent(Availability.UNAVAILABLE, reason)
                storage.append(
                    StorageObservation(
                        purpose=purpose,
                        path=str(path),
                        total_bytes=unavailable,
                        used_bytes=unavailable,
                        free_bytes=unavailable,
                    )
                )
        release = platform.release()
        environment = "wsl2" if "microsoft" in release.lower() else "native"
        return HostInventory(
            identity_key=identity_key,
            identity_source=identity_source,
            hostname=socket.gethostname(),
            display_name=socket.gethostname(),
            os_name=platform.system() or "unknown",
            os_version=self._known_or_unknown(platform.version() or None, "OS version unavailable"),
            architecture=platform.machine() or "unknown",
            kernel_version=self._known_or_unknown(release or None, "kernel unavailable"),
            cpu_model=self._known_or_unknown(cpu_model, "CPU model unavailable"),
            logical_cpu_count=self._known_or_unknown(logical, "logical CPU count unavailable"),
            physical_core_count=self._known_or_unknown(physical, "physical core count unavailable"),
            total_ram_bytes=ObservedValue.known(memory.total),
            available_ram_bytes=ObservedValue.known(memory.available),
            agent_version=__version__,
            storage=storage,
            metadata={"environment": environment},
        )


class NvidiaNvmlCollector:
    def __init__(self, api: Any) -> None:
        self.api = api
        self.nvml_error = getattr(api, "NVMLError", Exception)
        self.not_supported = getattr(api, "NVMLError_NotSupported", ())

    def _optional(
        self, call: Callable[[], T], transform: Callable[[T], Any] | None = None
    ) -> ObservedValue[Any]:
        try:
            value = call()
            return ObservedValue.known(transform(value) if transform else value)
        except self.not_supported:
            return ObservedValue.absent(Availability.UNSUPPORTED, "NVML metric not supported")
        except self.nvml_error as exc:
            return ObservedValue.absent(Availability.TEMPORARILY_FAILED, type(exc).__name__)

    @staticmethod
    def _text(value: Any) -> str:
        return value.decode("utf-8", errors="replace") if isinstance(value, bytes) else str(value)

    @staticmethod
    def _cuda_version(value: int) -> str:
        return f"{value // 1000}.{(value % 1000) // 10}"

    def _architecture(self, handle: Any) -> ObservedValue[str]:
        def transform(value: Any) -> str:
            for name in dir(self.api):
                if name.startswith("NVML_DEVICE_ARCH_") and getattr(self.api, name) == value:
                    return name.removeprefix("NVML_DEVICE_ARCH_").lower()
            return f"nvml_arch_{value}"

        function = getattr(self.api, "nvmlDeviceGetArchitecture", None)
        if function is None:
            return ObservedValue.absent(Availability.UNSUPPORTED, "architecture API unavailable")
        return self._optional(lambda: function(handle), transform)

    def collect(self) -> NvidiaCollection:
        initialized = False
        try:
            self.api.nvmlInit()
            initialized = True
            count = self.api.nvmlDeviceGetCount()
            driver = self._optional(self.api.nvmlSystemGetDriverVersion, self._text)
            cuda_function = getattr(self.api, "nvmlSystemGetCudaDriverVersion_v2", None) or getattr(
                self.api, "nvmlSystemGetCudaDriverVersion", None
            )
            cuda_driver = (
                self._optional(cuda_function, self._cuda_version)
                if cuda_function
                else ObservedValue.absent(Availability.UNSUPPORTED, "CUDA driver API unavailable")
            )
            inventory: list[AcceleratorInventory] = []
            telemetry: list[AcceleratorTelemetry] = []
            failures: list[str] = []
            observed_at = datetime.now(UTC)
            for index in range(count):
                try:
                    handle = self.api.nvmlDeviceGetHandleByIndex(index)
                    device_uuid = self._text(self.api.nvmlDeviceGetUUID(handle))
                    name = self._text(self.api.nvmlDeviceGetName(handle))
                    memory = self.api.nvmlDeviceGetMemoryInfo(handle)
                    pci = self._optional(
                        lambda: self.api.nvmlDeviceGetPciInfo(handle),
                        lambda value: self._text(value.busId),
                    )
                    compute = self._optional(
                        lambda: self.api.nvmlDeviceGetCudaComputeCapability(handle)
                    )
                    major = (
                        ObservedValue.known(int(compute.value[0]))
                        if compute.availability is Availability.KNOWN and compute.value is not None
                        else ObservedValue.absent(compute.availability, compute.reason)
                    )
                    minor = (
                        ObservedValue.known(int(compute.value[1]))
                        if compute.availability is Availability.KNOWN and compute.value is not None
                        else ObservedValue.absent(compute.availability, compute.reason)
                    )
                    mig_function = getattr(self.api, "nvmlDeviceGetMigMode", None)
                    mig = (
                        self._optional(lambda: mig_function(handle), lambda value: bool(value[0]))
                        if mig_function
                        else ObservedValue.absent(Availability.UNSUPPORTED, "MIG API unavailable")
                    )
                    inventory.append(
                        AcceleratorInventory(
                            device_index=index,
                            device_uuid=device_uuid,
                            pci_bus_id=pci,
                            name=name,
                            architecture=self._architecture(handle),
                            compute_capability_major=major,
                            compute_capability_minor=minor,
                            total_vram_bytes=ObservedValue.known(int(memory.total)),
                            driver_version=driver,
                            cuda_driver_version=cuda_driver,
                            mig_mode_current=mig,
                            inventory_at=observed_at,
                        )
                    )
                    utilization = self._optional(
                        lambda: self.api.nvmlDeviceGetUtilizationRates(handle)
                    )
                    gpu_util = (
                        ObservedValue.known(int(utilization.value.gpu))
                        if utilization.availability is Availability.KNOWN
                        and utilization.value is not None
                        else ObservedValue.absent(utilization.availability, utilization.reason)
                    )
                    mem_util = (
                        ObservedValue.known(int(utilization.value.memory))
                        if utilization.availability is Availability.KNOWN
                        and utilization.value is not None
                        else ObservedValue.absent(utilization.availability, utilization.reason)
                    )
                    telemetry.append(
                        AcceleratorTelemetry(
                            device_uuid=device_uuid,
                            observed_at=observed_at,
                            used_vram_bytes=ObservedValue.known(int(memory.used)),
                            free_vram_bytes=ObservedValue.known(int(memory.free)),
                            gpu_utilization_percent=gpu_util,
                            memory_utilization_percent=mem_util,
                            temperature_c=self._optional(
                                lambda: self.api.nvmlDeviceGetTemperature(
                                    handle, self.api.NVML_TEMPERATURE_GPU
                                ),
                                int,
                            ),
                            power_draw_w=self._optional(
                                lambda: self.api.nvmlDeviceGetPowerUsage(handle),
                                lambda value: round(value / 1000, 3),
                            ),
                            power_limit_w=self._optional(
                                lambda: self.api.nvmlDeviceGetEnforcedPowerLimit(handle),
                                lambda value: round(value / 1000, 3),
                            ),
                            graphics_clock_mhz=self._optional(
                                lambda: self.api.nvmlDeviceGetClockInfo(
                                    handle, self.api.NVML_CLOCK_GRAPHICS
                                ),
                                int,
                            ),
                            memory_clock_mhz=self._optional(
                                lambda: self.api.nvmlDeviceGetClockInfo(
                                    handle, self.api.NVML_CLOCK_MEM
                                ),
                                int,
                            ),
                            fan_speed_percent=self._optional(
                                lambda: self.api.nvmlDeviceGetFanSpeed(handle), int
                            ),
                            performance_state=self._optional(
                                lambda: self.api.nvmlDeviceGetPerformanceState(handle),
                                lambda value: f"P{value}",
                            ),
                        )
                    )
                except self.nvml_error as exc:
                    failures.append(f"device {index}: {type(exc).__name__}")
            availability = Availability.KNOWN if not failures else Availability.TEMPORARILY_FAILED
            return NvidiaCollection(
                availability=availability,
                reason="; ".join(failures)
                or ("no NVIDIA devices detected" if count == 0 else None),
                inventory=inventory,
                telemetry=telemetry,
                observed_at=observed_at,
            )
        except self.nvml_error as exc:
            return NvidiaCollection(
                availability=Availability.UNAVAILABLE, reason=type(exc).__name__
            )
        finally:
            if initialized:
                with suppress(self.nvml_error):
                    self.api.nvmlShutdown()


def build_nvml_collector() -> NvidiaNvmlCollector:
    import pynvml  # type: ignore[import-untyped]

    return NvidiaNvmlCollector(pynvml)