from __future__ import annotations

import os
from dataclasses import dataclass

os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

import torch


DEVICE_ENV = "R1_MATH_DEVICE"
DTYPE_ENV = "R1_MATH_DTYPE"


@dataclass(frozen=True)
class RuntimeDevice:
    device: torch.device
    dtype: torch.dtype

    @property
    def type(self) -> str:
        return self.device.type

    @property
    def label(self) -> str:
        if self.type == "cuda":
            name = torch.cuda.get_device_name(self.device)
            return f"CUDA - {name}"
        if self.type == "mps":
            return "MPS - Apple GPU"
        return "CPU"


def is_mps_available() -> bool:
    backend = getattr(torch.backends, "mps", None)
    return bool(backend and backend.is_available())


def resolve_device(preferred: str | None = None) -> torch.device:
    requested = (preferred or os.getenv(DEVICE_ENV) or "auto").strip().lower()

    if requested in {"auto", ""}:
        if torch.cuda.is_available():
            return torch.device("cuda")
        if is_mps_available():
            return torch.device("mps")
        return torch.device("cpu")

    if requested.startswith("cuda"):
        if not torch.cuda.is_available():
            raise RuntimeError(f"{DEVICE_ENV}={requested!r}, but CUDA is not available.")
        return torch.device(requested)

    if requested == "mps":
        if not is_mps_available():
            raise RuntimeError(f"{DEVICE_ENV}=mps, but PyTorch MPS is not available.")
        return torch.device("mps")

    if requested == "cpu":
        return torch.device("cpu")

    raise ValueError(f"Unsupported {DEVICE_ENV}: {requested!r}. Use auto, cuda, mps, or cpu.")


def _dtype_from_name(name: str) -> torch.dtype:
    normalized = name.strip().lower()
    if normalized in {"float16", "fp16", "half"}:
        return torch.float16
    if normalized in {"bfloat16", "bf16"}:
        return torch.bfloat16
    if normalized in {"float32", "fp32", "full"}:
        return torch.float32
    raise ValueError(f"Unsupported {DTYPE_ENV}: {name!r}. Use auto, float16, bfloat16, or float32.")


def default_dtype_for_device(device: torch.device) -> torch.dtype:
    requested = (os.getenv(DTYPE_ENV) or "auto").strip().lower()
    if requested not in {"auto", ""}:
        return _dtype_from_name(requested)

    if device.type == "cuda":
        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    if device.type == "mps":
        return torch.float16
    return torch.float32


def get_runtime_device(preferred: str | None = None) -> RuntimeDevice:
    device = resolve_device(preferred)
    return RuntimeDevice(device=device, dtype=default_dtype_for_device(device))


def dtype_name(dtype: torch.dtype) -> str:
    return str(dtype).removeprefix("torch.")


def runtime_metadata(preferred: str | None = None) -> dict[str, str]:
    runtime = get_runtime_device(preferred)
    return {
        "device": runtime.type,
        "device_label": runtime.label,
        "torch_dtype": dtype_name(runtime.dtype),
    }


def mps_safe_training_defaults(device: torch.device) -> dict[str, bool | str]:
    if device.type == "mps":
        return {
            "bf16": False,
            "fp16": False,
            "tf32": False,
            "optim": "adamw_torch",
            "dataloader_pin_memory": False,
            "use_mps_device": True,
        }

    if device.type == "cuda":
        return {
            "bf16": torch.cuda.is_bf16_supported(),
            "fp16": not torch.cuda.is_bf16_supported(),
            "tf32": True,
            "optim": "paged_adamw_8bit",
            "dataloader_pin_memory": True,
            "use_mps_device": False,
        }

    return {
        "bf16": False,
        "fp16": False,
        "tf32": False,
        "optim": "adamw_torch",
        "dataloader_pin_memory": False,
        "use_mps_device": False,
    }