from enum import Enum
from dataclasses import dataclass, field
from collections.abc import MutableMapping, Iterator
from typing import Generic, TypeVar, Optional, Dict, Callable
from .room_dims import RoomDimensions
from .lamp import Lamp
from .calc_zone import CalcZone, CalcPlane, CalcVol
from .reflectance import Surface
import warnings


# todo - integrate this
class OnCollision(str, Enum):
    ERROR = "error"
    OVERWRITE = "overwrite"
    INCREMENT = "increment"


T = TypeVar("T")


@dataclass
class Registry(Generic[T], MutableMapping[str, T]):
    """A thin wrapper around dict[str, T] with consistent ID/collision behavior."""

    base_id = "item"
    on_collision: str = "increment"  # "error" | "overwrite" | "increment"
    dims: Optional[Callable[[], "RoomDimensions"]] = None
    _items: Dict[str, T] = field(default_factory=dict)

    # ---- MutableMapping protocol ----
    def __getitem__(self, key: str) -> T:
        return self._items[key]

    def __setitem__(self, key: str, value: T) -> None:
        self._items[key] = value

    def __delitem__(self, key: str) -> None:
        del self._items[key]

    def __iter__(self) -> Iterator[str]:
        return iter(self._items)

    def __len__(self) -> int:
        return len(self._items)

    def require(self, key: str):
        try:
            return self._items[key]
        except KeyError as e:
            raise KeyError(
                f"Unknown id: {key!r}. Available: {list(self._items.keys())}"
            ) from e

    def remove(self, key: str):
        try:
            return self._items.pop(key)
        except KeyError as e:
            raise KeyError(f"Cannot remove {key!r}; not found.") from e

    def clear(self) -> None:
        self._items.clear()

    def _unique_id(self, base: str) -> str:
        if base not in self._items:
            return base
        prefix = base + "-"
        max_suffix = 1  # 1 corresponds to plain `base` being present
        for key in self.keys():
            if key == base:
                max_suffix = max(max_suffix, 1)
            elif key.startswith(prefix):
                rest = key[len(prefix) :]
                if rest.isdigit():
                    n = int(rest)
                    if n > max_suffix:
                        max_suffix = n
        # Next free number
        return f"{base}-{max_suffix + 1}"

    def _check_position(self, dims, obj):
        """
        check if an object's dimensions exceed the room's boundaries.
        """
        msg = None
        origin, roomdims = self.dims().origin, self.dims().dimensions
        for coord, pt1, pt2 in zip(dims, origin, roomdims):
            if coord > pt2 or coord < pt1:
                msg = f"{obj.name} exceeds room boundaries!"
                warnings.warn(msg, stacklevel=2)
        return msg

    def check_position(self, obj):
        """implement in subclasses"""
        raise NotImplementedError

    def get_position_warnings(self):
        dct = {}
        for obj_id, obj in self.items():
            dct[obj_id] = self.check_position(obj)
        return dct

    def _validate(self, obj):
        return obj

    def validate(self):
        for obj_id, obj in self.items():
            self._validate(obj)

    def add(self, obj, on_collision=None):
        policy = on_collision or self.on_collision
        key = getattr(obj, "id") or self.base_id
        if key in self._items:
            if policy == "error":
                raise KeyError(f"ID {key!r} already exists.")
            if policy == "increment":
                key = self._unique_id(key)
            # OVERWRITE keeps key as-is
        obj._assign_id(key)
        self._items[key] = self._validate(obj)
        # self.on_added(key, obj)  # hook
        return key

    @property
    def calc_state(self) -> tuple:
        tpl = ()
        for val in self.values():
            if hasattr(val, "calc_state"):
                tpl += val.calc_state
        return tpl

    @property
    def update_state(self) -> tuple:
        tpl = ()
        for val in self.values():
            if hasattr(val, "update_state"):
                tpl += val.update_state
        return tpl


@dataclass
class LampRegistry(Registry["Lamp"]):
    base_id = "Lamp"

    def _validate(self, lamp):
        if not isinstance(lamp, Lamp):
            raise TypeError(f"Must be type Lamp, not {type(lamp)}")
        if lamp.surface.units != self.dims().units:
            # automatically convert for now
            lamp.set_units(self.dims().units)
        # check dimensions
        self.check_position(lamp)
        return lamp

    def check_position(self, lamp) -> str:
        return self._check_position(lamp.position, lamp)

    @property
    def wavelengths(self) -> list:
        return {k: v.wavelength for k, v in self.items()}
        # lst = [v.wavelength for k,v in self.items() if v.wavelength is not None]
        # return sorted(set(lst))


@dataclass
class ZoneRegistry(Registry["CalcZone"]):
    base_id = "CalcZone"

    def _validate(self, zone):
        if not isinstance(zone, (CalcZone, CalcPlane, CalcVol)):
            raise TypeError(f"Must be CalcZone, CalcPlane, or CalcVol not {type(zone)}")
        self.check_position(zone)
        return zone

    def check_position(self, zone) -> str:
        if isinstance(zone, (CalcPlane, CalcVol)):
            x, y, z = zone.coords.T
            dimensions = x.max(), y.max(), z.max()
            return self._check_position(dimensions, zone)
        return None


@dataclass
class SurfaceRegistry(Registry["Surface"]):
    base_id = "Surface"

    def _validate(self, surface):
        if not isinstance(surface, Surface):
            raise TypeError(f"Must be Surface, not {type(surface)}")
        return surface

    def check_position(self, surface) -> str:
        x, y, z = surface.plane.coords.T
        dimensions = x.max(), y.max(), z.max()
        return self._check_position(dimensions, surface)
