import json
import torch
from typing import Any

# DEVICES

if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")

print(f"[MultimodalReasoning] PyTorch is using primary device: {device}.")

# CONFIG PATHS
class ConfigNode(dict):
    """
    Dict that supports both:
        cfg["Path"]
        cfg.Path

    Nested dictionaries are automatically wrapped too.
    """

    def __init__(self, data: dict):
        super().__init__()

        for key, value in data.items():
            self[key] = self._wrap(value)

    @staticmethod
    def _wrap(value: Any) -> Any:
        if isinstance(value, dict):
            return ConfigNode(value)

        if isinstance(value, list):
            return [ConfigNode._wrap(x) for x in value]

        return value

    def __getattr__(self, key: str) -> Any:
        try:
            return self[key]
        except KeyError:
            raise AttributeError(f"Config key not found: {key}") from None

    def __setattr__(self, key: str, value: Any):
        self[key] = self._wrap(value)

    def __delattr__(self, key: str):
        try:
            del self[key]
        except KeyError:
            raise AttributeError(f"Config key not found: {key}") from None

with open("configs.json", "r", encoding="utf-8") as f:
    configs = ConfigNode(json.load(f))

def get(path: str | tuple[str, ...] | list[str], default: Any = None, *, strict: bool = True) -> Any:
    """
    Examples:
        get("LLM/MODEL_NAME")
        get(("DATASET", "REASONING_SFT"))
        get(["DATASET", "REASONING_SFT", "Path"])
    """

    if isinstance(path, str):
        parts = [p for p in path.split("/") if p]
    else:
        parts = list(path)

    current = configs

    for key in parts:
        try:
            current = current[key]
        except (KeyError, TypeError):
            if strict:
                raise KeyError(f"Config path not found: {'/'.join(parts)} at key '{key}'")
            return default

    return current

def get_device():
    return device