#!/usr/bin/env python3
"""
helix_cdc/runtime/device_registry.py

Device Registry for HELIX64 v0.2.0 Universal Driver.

Decides which backends/devices are allowed and sets up environment
for child processes.

Usage:
    from helix_cdc.runtime.device_registry import DeviceRegistry

    registry = DeviceRegistry.from_probe()  # Probes system
    registry = DeviceRegistry.from_env()    # Reads HELIX_BACKEND/HELIX_DEVICES

    # Get allowed devices
    print(registry.allowed_targets())  # ["cpu", "cuda:0", ...]

    # Get environment to set before exec
    env = registry.get_env()  # {"HELIX_BACKEND": "cpu", ...}
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Dict, Any, List, Optional

from .device_probe import probe_all, probe_cuda, probe_cpu


@dataclass
class DeviceRegistry:
    """
    Registry of allowed compute devices.

    Acts as the single source of truth for "which devices can we use?"
    so that FlowTorch, FibPi, and experiments all agree.
    """

    cpu_available: bool = True
    cuda_available: bool = False
    cuda_devices: List[str] = field(default_factory=list)
    cuda_error: Optional[str] = None
    cpu_cores: int = 1
    cpu_affinity: List[int] = field(default_factory=list)
    selected_backend: str = "cpu"  # "cpu" or "cuda"
    probe_result: Dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_probe(cls, force_backend: Optional[str] = None) -> "DeviceRegistry":
        """
        Create registry by probing system.

        Args:
            force_backend: Force "cpu" or "cuda" (None = auto-detect)
        """
        probe = probe_all()

        cpu_info = probe.get("cpu", {})
        cuda_info = probe.get("cuda", {})

        cpu_available = cpu_info.get("available", True)
        cuda_available = cuda_info.get("available", False)
        cuda_devices = cuda_info.get("visible_devices", [])
        cuda_error = cuda_info.get("error")
        cpu_cores = cpu_info.get("cores", 1)
        cpu_affinity = cpu_info.get("affinity", [0])

        # Select backend
        if force_backend == "cpu":
            selected = "cpu"
        elif force_backend == "cuda":
            if cuda_available:
                selected = "cuda"
            else:
                # Requested CUDA but unavailable - fall back to CPU
                selected = "cpu"
        else:
            # Auto: prefer CUDA if available
            selected = "cuda" if cuda_available else "cpu"

        return cls(
            cpu_available=cpu_available,
            cuda_available=cuda_available,
            cuda_devices=cuda_devices,
            cuda_error=cuda_error,
            cpu_cores=cpu_cores,
            cpu_affinity=cpu_affinity,
            selected_backend=selected,
            probe_result=probe,
        )

    @classmethod
    def from_env(cls) -> "DeviceRegistry":
        """
        Create registry from environment variables.

        Reads:
            HELIX_BACKEND: "cpu" or "cuda" or "auto"
            HELIX_DEVICES: comma-separated devices (e.g. "cpu,cuda:0")
        """
        backend = os.environ.get("HELIX_BACKEND", "auto")
        devices_str = os.environ.get("HELIX_DEVICES", "")

        # Parse devices
        if devices_str:
            devices = [d.strip() for d in devices_str.split(",") if d.strip()]
        else:
            devices = []

        # Determine backend from devices or explicit setting
        if backend == "auto":
            cuda_devices = [d for d in devices if d.startswith("cuda")]
            if cuda_devices:
                selected = "cuda"
            else:
                selected = "cpu"
        else:
            selected = backend

        # Build registry
        cuda_devices_list = [d for d in devices if d.startswith("cuda")]
        cuda_available = len(cuda_devices_list) > 0

        return cls(
            cpu_available=True,  # CPU always available
            cuda_available=cuda_available,
            cuda_devices=cuda_devices_list,
            cuda_error=None,
            cpu_cores=os.cpu_count() or 1,
            cpu_affinity=list(range(os.cpu_count() or 1)),
            selected_backend=selected,
            probe_result={},
        )

    def allowed_targets(self) -> List[str]:
        """
        Get list of allowed device targets.

        Returns:
            ["cpu"] or ["cpu", "cuda:0", "cuda:1", ...]
        """
        targets = ["cpu"]  # CPU always allowed
        if self.cuda_available and self.cuda_devices:
            targets.extend(self.cuda_devices)
        return targets

    def get_env(self, for_subprocess: bool = True) -> Dict[str, str]:
        """
        Get environment variables to set.

        Args:
            for_subprocess: If True, include vars that disable CUDA when CPU-only

        Returns:
            {
                "HELIX_BACKEND": "cpu" or "cuda",
                "HELIX_DEVICES": "cpu,cuda:0,...",
                "USE_TF": "0",
                "USE_TORCH": "1",
                "CUDA_VISIBLE_DEVICES": "" (if CPU-only)
            }
        """
        env = {
            "HELIX_BACKEND": self.selected_backend,
            "HELIX_DEVICES": ",".join(self.allowed_targets()),
            "USE_TF": "0",  # Disable TensorFlow by default
            "USE_TORCH": "1",  # Enable PyTorch
        }

        # If CPU-only, disable CUDA completely to avoid spurious warnings
        if for_subprocess and self.selected_backend == "cpu":
            env["CUDA_VISIBLE_DEVICES"] = ""

        # Set HELIX_CORES for FibPi
        if self.cpu_affinity:
            env["HELIX_CORES"] = ",".join(str(c) for c in self.cpu_affinity)

        return env

    def to_dict(self) -> Dict[str, Any]:
        """Serialize registry to dict (for receipts)."""
        return {
            "cpu_available": self.cpu_available,
            "cuda_available": self.cuda_available,
            "cuda_devices": self.cuda_devices,
            "cuda_error": self.cuda_error,
            "cpu_cores": self.cpu_cores,
            "cpu_affinity": self.cpu_affinity,
            "selected_backend": self.selected_backend,
            "allowed_targets": self.allowed_targets(),
        }

    def describe(self) -> str:
        """Human-readable description."""
        lines = [
            f"Backend: {self.selected_backend}",
            f"CPU: {self.cpu_cores} cores, {len(self.cpu_affinity)} available",
        ]
        if self.cuda_available:
            lines.append(f"CUDA: {len(self.cuda_devices)} devices ({', '.join(self.cuda_devices)})")
        elif self.cuda_error:
            lines.append(f"CUDA: unavailable ({self.cuda_error})")
        else:
            lines.append("CUDA: not detected")

        return "\n".join(lines)


def get_registry(force_backend: Optional[str] = None) -> DeviceRegistry:
    """
    Get the device registry (singleton-like, but fresh probe each time).

    Args:
        force_backend: Force "cpu" or "cuda" (None = auto)

    Returns:
        DeviceRegistry instance
    """
    # Check if env vars are set
    if os.environ.get("HELIX_BACKEND"):
        return DeviceRegistry.from_env()
    else:
        return DeviceRegistry.from_probe(force_backend=force_backend)


if __name__ == "__main__":
    import json

    print("=" * 60)
    print("HELIX64 Device Registry")
    print("=" * 60)
    print()

    registry = DeviceRegistry.from_probe()

    print(registry.describe())
    print()
    print("Allowed targets:", registry.allowed_targets())
    print()
    print("Environment for subprocess:")
    for k, v in registry.get_env().items():
        print(f"  {k}={v}")
    print()
    print("Full dict:")
    print(json.dumps(registry.to_dict(), indent=2))
