import collections
from collections.abc import Iterator, Mapping, MutableMapping, Sequence
from typing import Self


class OverlayChainMap[K, V](MutableMapping[K, V]):
    """
    ChainMap implementation that protects its member mappings from mutation, while still allowing mutations on the chain map
    itself.

    See collections.ChainMap for more details. Does not implement the `maps` and `parents` properties,
    because they lose much of their meaning due to the overlay mapping sitting in between the user-supplied ones.
    """
    def __init__(self, *maps: Mapping[K, V]) -> None:
        # construct backing ChainMap with owned, mutable dict as first map
        self._chainmap: collections.ChainMap[K, V] = collections.ChainMap(
            {},
            # This is a type violation due to a limitation of the ChainMap constructor (python/typeshed#8430).
            # This workaround builds on the (documented) knowledge that ChainMap only ever modifies the first mapping,
            # and the fact that the constructed chain map is never exposed to the end user.
            # Therefore the type violation is a safe, if imperfect, workaround to the typing limitation.
            *maps,  # type: ignore
        )

    def new_child(self, m: Mapping[K, V] | None = None, **kwargs: V) -> Self:
        return self.__class__(
            *self._chainmap.new_child(
                m,  # type: ignore
            ).new_child({}, **kwargs).maps
        )

    def __getitem__(self, key: K) -> V:
        return self._chainmap[key]

    def __setitem__(self, key: K, value: V) -> None:
        self._chainmap[key] = value

    def __delitem__(self, key: K) -> None:
        del self._chainmap[key]

    def __iter__(self) -> Iterator[K]:
        return iter(self._chainmap)

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


class ImmutableChainMap[K, V](Mapping[K, V]):
    """
    Immutable ChainMap implementation. See collections.ChainMap for more details.
    """
    def __init__(self, *maps: Mapping[K, V]) -> None:
        # construct backing ChainMap with owned, mutable dict as first map
        self._chainmap: collections.ChainMap[K, V] = collections.ChainMap(
            # This is a type violation due to a limitation of the ChainMap constructor (python/typeshed#8430).
            # This workaround builds on the fact that we own the backing chain map, we never mutate it, and we only
            # expose it or its attributes as immutable properties.
            *maps  # type: ignore
        )

    def new_child(self, m: Mapping[K, V] | None = None) -> Self:
        return self.__class__(
            *self._chainmap.new_child(
                m,  # type: ignore
            ).maps
        )

    @property
    def maps(self) -> Sequence[Mapping[K, V]]:
        return self._chainmap.maps

    @property
    def parents(self) -> Self:
        return self.__class__(*self._chainmap.parents.maps)

    def __getitem__(self, key: K) -> V:
        return self._chainmap[key]

    def __iter__(self) -> Iterator[K]:
        return iter(self._chainmap)

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



d1: MutableMapping[str, int] = {"one": 1, "shared": 1}
d1_copy: Mapping[str, int] = {**d1}
d2: MutableMapping[str, int] = {"two": 2, "shared": 2}
d2_copy: Mapping[str, int] = {**d2}

copied: Mapping[str, int] = {**d1, **d2}
chained: MutableMapping[str, int] = collections.ChainMap(d2, d1)
overlayed: MutableMapping[str, int] = OverlayChainMap(d2, d1)
immutably_chained: Mapping[str, int] = ImmutableChainMap(d2, d1)

assert (
    list(chained.items())
    == list(overlayed.items())
    == list(immutably_chained.items())
    == list(copied.items())
    == [("one", 1), ("shared", 2), ("two", 2)]
)

overlayed["three"] = 3  # adds a new item
overlayed["one"] = 42  # changes the value in overlayed, not in d1
overlayed["two"] = 42  # changes the value in overlayed, not in d2
assert list(overlayed.items()) == [("one", 42), ("shared", 2), ("two", 42), ("three", 3)]
assert d1 == d1_copy
assert d2 == d2_copy
chained["three"] = 3  # adds a new item
chained["one"] = 42  # changes the value in chained, not in d1
chained["two"] = 42  # changes the value in both chained and d2
assert list(chained.items()) == [("one", 42), ("shared", 2), ("two", 42), ("three", 3)]
assert d1 == d1_copy
assert d2 != d2_copy