from types import NoneType
from typing import Any

Primitive = bool | float | int | str | NoneType


class Group(dict):
    """
    A dictionary that enforces:
    - All keys are strings
    - All values are JSON-serializable (primitives, lists, dicts, or nested JsonableDict)
    """

    PRIMITIVES = (str, int, float, bool, NoneType)

    def __init__(self, *args, **kwargs):
        # Create a temporary dict from input
        temp = dict(*args, **kwargs)
        # Validate all entries before accepting
        for key, value in temp.items():

            self._validate_key(key)
            self._validate_value(value)
        # Use dict's __init__ after validation
        super().__init__(temp)

    def _validate_key(self, key: Any) -> None:
        if not isinstance(key, str):
            raise TypeError(
                f"JsonableDict keys must be strings, got {type(key).__name__}: {key!r}"
            )

    def _validate_value(self, value: Any) -> None:
        if value is None or isinstance(value, (str, int, float, bool)):
            return

        if isinstance(value, (list, tuple)):
            for item in value:
                self._validate_value(item)
            return

        if isinstance(value, dict):
            # Allow regular dicts (will be recursively validated on assignment)
            for k, v in value.items():
                self._validate_key(k)
                self._validate_value(v)
            return

        if isinstance(value, Group):
            return  # Already validated internally

        raise TypeError(
            f"Value is not JSON serializable: {value!r} (type: {type(value).__name__})"
        )

    # Override mutation methods to enforce rules

    def __setitem__(self, key: str, value: Primitive) -> None:
        self._validate_key(key)
        self._validate_value(value)
        super().__setitem__(key, value)

    def update(self, other=(), /, **kwds) -> None:
        if hasattr(other, "keys"):
            for k in other.keys():
                if type(other[k]) is Group and type(self.get(k)) is Group:
                    self[k].update(other[k])
                else:
                    self[k] = other[k]
        else:
            for k, v in other:
                if type(self.get(k)) is Group and type(v) is Group:
                    self[k].update(v)
                else:
                    self[k] = v
        for k, v in kwds.items():
            if type(self.get(k)) is Group and type(v) is Group:
                self[k].update(v)
            else:
                self[k] = v

    def setdefault(self, key: str, default: Primitive = None) -> Primitive:
        if key not in self:
            self[key] = default
        return self[key]
