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), and ``map``/``pmap`` carry the function's return type; a broadcast
    names a field no type system can resolve, so it carries ``Any`` values.
    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", "_cg_base", "_cg_pin")
    _entries: dict[K, V]
    _key: ClassVar[type[Keyable] | None] = None
    _element: ClassVar[type[CGType] | None] = None

    @overload
    def __init__(self, entries: Iterable[tuple[K, V]] | dict[K, V] = ()) -> None: ...
    @overload
    def __init__[E](
        self: Map[Time, E],
        entries: Iterable[tuple[datetime, E]] | dict[datetime, E] = (),
    ) -> None: ...
    @overload
    def __init__[E](
        self: Map[String, E] | Map[UserID, E],
        entries: Iterable[tuple[str, E]] | dict[str, E] = (),
    ) -> None: ...
    @overload
    def __init__[E](
        self: Map[Bool, E],
        entries: Iterable[tuple[bool, E]] | dict[bool, E] = (),
    ) -> None: ...
    def __init__(
        self: Map[_KS, _E],
        entries: Iterable[tuple[_KD, _E]] | dict[_KD, _E] = (),
    ) -> 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):
            keyable = ", ".join(scalar.__name__ for scalar in keyable_scalars())
            raise TypeError(f"Map key must be a keyable scalar ({keyable}), got {key!r}")
        check_cgtype(element, "Map element")
        subscript = (cls, key, element)
        if subscript not in _SUBSCRIPTS:
            _SUBSCRIPTS[subscript] = type(
                f"Map[{key.__name__}, {element.__name__}]",
                (cls,),
                {
                    "__slots__": (),
                    "_key": key,
                    "_element": element,
                    "storable": element.storable,
                },
            )
        return _SUBSCRIPTS[subscript]

    # --- 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 _children(cls) -> list[tuple[str, type[CGType]]]:
        """The element type, labelled ``[*]``; a schemaless map declares none."""
        return [("[*]", cls._element)] if cls._element is not None else []

    @classmethod
    def _descend(cls, component: str) -> type[CGType]:
        """The element type, reached by any canonical key encoding.

        Raises:
            SchemaPathError: If this map is schemaless, or ``component`` is not a
                key of its key type in canonical encoding.
        """
        element = cls._element
        if element is None:
            raise SchemaPathError(f"{cls.__name__} is a schemaless map")
        cls.key_of(component)
        return element

    @classmethod
    def key_of(cls, component: str) -> object:
        """The key the on-disk entry name ``component`` spells, decoded by the key type.

        ```python
        assert cg.Map[cg.Time, cg.Number].key_of("0006380070917400") == noon
        ```

        Raises:
            SchemaPathError: If this map is schemaless, or ``component`` is not a
                key of its key type in canonical encoding.
        """
        if cls._key is None or cls._element is None:
            raise SchemaPathError(f"{cls.__name__} is a schemaless map")
        try:
            key = cls._key.key_decode(component)
        except Exception as error:  # noqa: BLE001 — any decode failure is a bad key
            raise SchemaPathError(f"{component!r} is not a {cls._key.__name__} map key") from error
        # A lenient decoder must not admit an alias folder name that shadows an entry.
        if cls._key.key_encode(key) != component:
            raise SchemaPathError(
                f"{component!r} is not the canonical {cls._key.__name__} key encoding"
            )
        return key

    @classmethod
    def load(cls, base: Path) -> Self:
        """The map stored at ``base``, keys decoded from folder names; an unwritten tree is empty.

        Raises:
            TypeError: If this map is schemaless.
        """
        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.is_dir():
            return cls()
        entries = {
            cls._key.key_decode(stem): cls._element.load(base / stem) for stem in entry_stems(base)
        }
        mapping = cls(cast("dict[K, V]", entries))
        remember_base(mapping, base)
        return mapping

    @classmethod
    def store(cls, base: Path, value: object) -> None:
        """Write ``value`` under ``base``, one key-encoded folder per entry.

        Raises:
            TypeError: If this map is schemaless, ``value`` is not a ``Map`` or
                ``dict``, or an entry value is ``Missing`` — omitting it would
                silently drop its key.
            SchemaPathError: If a key encodes to something that cannot name a folder.
        """
        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():
            stem = safe_component(cls._key.key_encode(key))
            if is_missing(item):
                raise TypeError(f"{cls.__name__}.store: entry {stem!r} is Missing")
            cls._element.store(base / stem, 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  # pyright: ignore[reportAssignmentType] # a mutable container is unhashable

    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."""
        forget_base(self)
        self._entries[key] = value

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

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

    def __getattr__(self, name: str) -> Map[K, Any]:
        """Each entry's ``name``, keys kept; a name no value has raises, sparse is ``Missing``.

        This map's own machinery is not an entry field (``broadcastable``).
        """
        if not broadcastable(self, name):
            raise AttributeError(name)
        accessed = broadcast(
            self._entries.values(),
            f"attribute {name!r}",
            lambda value: access(value, name),
        )
        return Map(zip(self._entries, accessed, strict=True))

    def __call__(self, *args: object, **kwargs: object) -> Map[K, Any]:
        """Call each entry value with the same arguments, keys kept; a failure is ``Missing``."""

        def called(value: V) -> object:
            if not callable(value):
                raise TypeError(f"Map value {value!r} is not callable")
            return value(*args, **kwargs)

        results = broadcast(self._entries.values(), "call", called)
        return Map(zip(self._entries, results, strict=True))

    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:
            forget_base(self)
            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[R](self, function: Callable[[V], R], label: str | None = None) -> Map[K, R]:
        """Apply ``function`` to each value, keeping keys; failures become ``Missing``.

        The walk is a progress step counting the entries, shown as ``label`` or
        ``function``'s name. 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``.
        """
        results = map_broadcast(list(self._entries.values()), function, label)
        return Map(zip(self._entries, results, strict=True))

    def pmap[R](self, function: Callable[[V], R], label: str | None = None) -> Map[K, R]:
        """Like ``map``, across the process's one thread pool, sized to the CPU count.

        Keys are kept. Parallel only where ``function`` releases the GIL:
        decoding an image, NumPy, a subprocess. ``cg.parallelism(workers)`` caps
        how many elements are in flight at once. Each worker runs in a copy of
        the caller's context, so the active scratch, workspace and progress
        step carry over; a ``pmap`` inside a worker runs sequentially, the outer
        pool being the parallelism.

        ```python
        thumbnails = journal.sessions.photo.pmap(session_thumbnail)
        ```
        """
        results = pbroadcast(self._entries.values(), "pmap", function, label)
        return Map(zip(self._entries, results, strict=True))

    @overload
    def to_dict[E](self: Map[Time, E]) -> dict[datetime, E]: ...
    @overload
    def to_dict[E](self: Map[String, E]) -> dict[str, E]: ...
    @overload
    def to_dict[E](self: Map[UserID, E]) -> dict[str, E]: ...
    @overload
    def to_dict[E](self: Map[Bool, E]) -> dict[bool, E]: ...
    @overload
    def to_dict(self) -> dict[K, V]: ...
    def to_dict(self: Map[_KS, _E]) -> dict[_KD, _E]:
        """A plain ``dict`` of the entries — decoded keys for a schema-keyed map."""
        return cast("dict[_KD, _E]", dict(self._entries))