map

class Map[K, V](CGType):
    """A keyed collection.

    ``Map[K, V]`` is storable when ``V`` is (so a keyed fan-out
    ``Map[K, Path[V]]`` is non-storable); a bare ``Map`` is a schemaless keyed
    view. The key is always a keyable — hence storable — scalar. The type
    parameters type the entry surface (``keys()``/``values()``, values by
    key); broadcasts carry ``object`` elements, narrowed by the caller. Key
    *parameters* accept ``object``: a schema map's runtime keys are the
    decoded values (``datetime`` for a ``Date`` key), not the scalar class
    the subscription names.
    """

    __slots__ = ("_entries",)
    _entries: dict[K, V]
    _key: ClassVar[type[Keyable] | None] = None
    _element: ClassVar[type[CGType] | None] = None

    def __init__(self, entries: Iterable[tuple[K, V]] | dict[K, V] = ()) -> None:
        """Wrap ``entries`` (any ``(key, value)`` iterable) as a Map."""
        object.__setattr__(self, "_entries", dict(entries))

    def __class_getitem__(cls, params: tuple[type[CGType], type[CGType]]) -> type[Map]:
        """Build ``Map[key, element]``, a storable subclass keyed by ``key``."""
        if not (isinstance(params, tuple) and len(params) == 2):
            raise TypeError("Map takes two parameters: Map[key_scalar, element]")
        key, element = params
        if not is_keyable(key):
            raise TypeError(
                f"Map key must be a keyable scalar (Time, String, UserID, Bool), got {key!r}"
            )
        check_cgtype(element, "Map element")
        return type(
            f"Map[{key.__name__}, {element.__name__}]",
            (cls,),
            {
                "__slots__": (),
                "_key": key,
                "_element": element,
                "storable": element.storable,
            },
        )

    # --- storage (a schema'd Map[K, V] only) --------------------------------

    @classmethod
    def schema(cls) -> dict[str, JSONValue]:
        """This map's wire form: its key and element schemas, tagged ``"map"``."""
        if cls._key is None or cls._element is None:
            raise TypeError("a schemaless Map has no wire form")
        return {
            "type": "map",
            "key": cls._key.schema(),
            "element": cls._element.schema(),
        }

    @classmethod
    def load(cls, base: Path) -> Self:
        """Read the key-encoded entries under ``base`` as this map type.

        Typed by the annotation naming this class — ``Map[Date, Session]``
        loads as itself. The schema supplies the runtime ``_key``/``_element``;
        the cast below is the one boundary where the type-erased registry meets
        the static entry types the schema declares.
        """
        if cls._key is None or cls._element is None:
            raise TypeError(f"{cls.__name__} cannot load without a key and element")
        if not base.exists():
            return cls()
        loaded = (
            (cls._key.key_decode(stem), cls._element.load(base / stem))
            for stem in entry_stems(base)
        )
        return cls(cast("Iterable[tuple[K, V]]", loaded))

    @classmethod
    def store(cls, base: Path, value: object) -> None:
        """Write ``value`` (a ``dict``) under ``base`` as key-encoded entries."""
        if cls._key is None or cls._element is None:
            raise TypeError(f"{cls.__name__} cannot store without a key and element")
        if isinstance(value, Map):
            entries = value.to_dict()
        elif isinstance(value, dict):
            entries = value
        else:
            raise TypeError(f"{cls.__name__}.store needs a dict or Map, got {value!r}")
        base.mkdir(parents=True, exist_ok=True)
        for key, item in entries.items():
            cls._element.store(base / cls._key.key_encode(key), item)

    # --- keyed access (any Map instance) ------------------------------------

    def __iter__(self) -> Iterator[V]:
        """Iterate the entry values, in insertion order."""
        return iter(self._entries.values())

    def __len__(self) -> int:
        """The number of entries."""
        return len(self._entries)

    def __bool__(self) -> bool:
        """Falsy when empty, like a plain dict."""
        return bool(self._entries)

    def __contains__(self, key: object) -> bool:
        """Whether ``key`` has an entry."""
        return key in self._entries

    def __eq__(self, other: object) -> bool:
        """Equal to another ``Map`` or a plain ``dict`` with the same entries."""
        if isinstance(other, Map):
            return self._entries == other._entries
        if isinstance(other, dict):
            return self._entries == other
        return NotImplemented

    __hash__ = None  # type: ignore[assignment]

    def __repr__(self) -> str:
        """The wrapped entries' own ``repr``."""
        return f"Map({self._entries!r})"

    def keys(self) -> Array[K]:
        """The decoded keys, as an ``Array``."""
        return Array(self._entries.keys())

    def values(self) -> Array[V]:
        """The entry values, as an ``Array``."""
        return Array(self._entries.values())

    def items(self) -> Array[tuple[K, V]]:
        """The ``(key, value)`` pairs, as an ``Array``."""
        return Array(self._entries.items())

    def __getitem__(self, key: object) -> V:
        """The value at ``key``; an absent key chains as a reasoned ``Missing``."""
        try:
            return _keyed(self)[key]
        except KeyError as error:
            return _Missing(error)

    def __setitem__(self, key: K, value: V) -> None:
        """Add or replace the entry at ``key``, in place."""
        self._entries[key] = value

    def __delitem__(self, key: object) -> None:
        """Remove the entry at ``key``, in place."""
        del _keyed(self)[key]

    # --- broadcast over entry values ----------------------------------------

    def __getattr__(self, name: str) -> Array[object]:
        """Broadcast attribute access over the entry values.

        A name no value has raises (a typo), a name some lack chains as
        ``Missing``.
        """
        if name.startswith("_"):
            raise AttributeError(name)
        return Array(broadcast_attr(list(self._entries.values()), name))

    def __dir__(self) -> list[str]:
        """Standard members plus the broadcastable entry-value field names.

        Surfaces ``sessions.rating`` in REPL/notebook autocomplete; the
        broadcast itself stays dynamically typed.
        """
        return [*super().__dir__(), *self.values().dkeys]

    # --- runtime helpers ----------------------------------------------------

    def get_or_insert(self, key: object, value: V) -> V:
        """The value at ``key``, inserting ``value`` first if it is absent."""
        entries = _keyed(self)
        if key not in entries:
            entries[key] = value
        return entries[key]

    def with_key(self, key: K, value: V) -> Map[K, V]:
        """A copy with ``key`` added or replaced."""
        return Map({**self._entries, key: value})

    def map(self, function: Callable[[V], object]) -> Map[K, object]:
        """Apply ``function`` to each value, keeping keys; failures become ``Missing``.

        Raises if the call fails on every entry for a non-data reason (a broken
        lambda fails loudly); a sparse data-absence failure chains as ``Missing``.
        """
        keys = list(self._entries)
        results: list[object] = []
        for value in self._entries.values():
            try:
                results.append(function(value))
            except Exception as error:  # noqa: BLE001 — per-entry isolation
                results.append(_Missing(error))
        guard_broadcast(results, "map")
        return Map(zip(keys, results, strict=True))

    def to_dict(self) -> dict[K, V]:
        """A plain ``dict`` of the entries (keys preserved)."""
        return dict(self._entries)