from __future__ import annotations

from dataclasses import fields
import os

import torch

from config import V2Config


def resolve_device(requested: str | torch.device) -> torch.device:
    name = str(requested).lower()
    if name == "auto":
        if torch.cuda.is_available():
            return torch.device("cuda")
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")
    if name.startswith("cuda"):
        return torch.device(requested if torch.cuda.is_available() else "cpu")
    if name == "mps":
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")
    return torch.device(requested)


def sync_device(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    elif device.type == "mps":
        torch.mps.synchronize()


def configure_threads(cpu_threads: int | None = None) -> int:
    count = int(cpu_threads or os.cpu_count() or 1)
    count = max(1, count)
    for key in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_MAX_THREADS"):
        os.environ.setdefault(key, str(count))
    torch.set_num_threads(count)
    # Interop threads > 1 often makes small Torch ops fight each other.
    torch.set_num_interop_threads(1)
    return count


def make_config(**kwargs: object) -> V2Config:
    allowed = {field.name for field in fields(V2Config)}
    return V2Config(**{key: value for key, value in kwargs.items() if key in allowed})


def normalize_config(config: V2Config) -> V2Config:
    defaults = V2Config()
    values = {
        field.name: getattr(config, field.name, getattr(defaults, field.name))
        for field in fields(V2Config)
    }
    return V2Config(**values)