from __future__ import annotations
import logging
import os
import torch

logger = logging.getLogger("riemann.device")
try:
    import torch_directml

    _DML_AVAILABLE = True
except ImportError:
    torch_directml = None
    _DML_AVAILABLE = False


def directml_available() -> bool:
    return _DML_AVAILABLE and torch_directml.device_count() > 0


def main_device() -> torch.device:
    if directml_available():
        return torch_directml.device(0)
    if torch.cuda.is_available():
        return torch.device("cuda", 0)
    logger.warning("No DirectML or CUDA device available; training will run on CPU.")
    return torch.device("cpu")


def secondary_device() -> torch.device:
    if directml_available() and torch_directml.device_count() > 1:
        return torch_directml.device(1)
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        return torch.device("cuda", 1)
    return main_device()


def cpu_device() -> torch.device:
    return torch.device("cpu")


def configure_cpu_threads(num_threads: int | None = None) -> int:
    threads = num_threads or os.cpu_count() or 1
    torch.set_num_threads(threads)
    try:
        torch.set_num_interop_threads(max(1, threads // 2))
    except RuntimeError:
        pass
    return threads


def describe_devices() -> dict:
    info = {
        "directml_available": directml_available(),
        "cuda_available": torch.cuda.is_available(),
        "devices": [],
    }
    if directml_available():
        for i in range(torch_directml.device_count()):
            info["devices"].append(
                {
                    "backend": "directml",
                    "index": i,
                    "name": torch_directml.device_name(i),
                }
            )
    if torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            info["devices"].append(
                {"backend": "cuda", "index": i, "name": torch.cuda.get_device_name(i)}
            )
    info["cpu_logical_cores"] = os.cpu_count()
    info["cpu_threads_configured"] = torch.get_num_threads()
    return info


if __name__ == "__main__":
    import json

    print(json.dumps(describe_devices(), indent=2))