"""
Lexical environment for s-expression evaluation.

Environments form a parent chain so inner scopes shadow outer ones
while still allowing lookup of free variables.
"""

from __future__ import annotations

from typing import Any


class Env:
    """A lexical scope mapping names → values with an optional parent."""

    __slots__ = ("_bindings", "_parent")

    def __init__(
        self,
        bindings: dict[str, Any] | None = None,
        parent: Env | None = None,
    ):
        self._bindings: dict[str, Any] = {} if bindings is None else bindings
        self._parent = parent

    # -- lookup -------------------------------------------------------------

    def lookup(self, name: str) -> Any:
        """Resolve *name*, walking the parent chain.

        Raises ``KeyError`` if not found.
        """
        if name in self._bindings:
            return self._bindings[name]
        if self._parent is not None:
            return self._parent.lookup(name)
        raise KeyError(name)

    def __contains__(self, name: str) -> bool:
        if name in self._bindings:
            return True
        if self._parent is not None:
            return name in self._parent
        return False

    def __getitem__(self, name: str) -> Any:
        return self.lookup(name)

    def __setitem__(self, name: str, value: Any) -> None:
        """Set *name* in the **current** scope (like ``define``)."""
        self._bindings[name] = value

    def get(self, name: str, default: Any = None) -> Any:
        try:
            return self.lookup(name)
        except KeyError:
            return default

    def update(self, other: dict[str, Any] | Env) -> None:
        """Merge *other*'s bindings into the **current** scope."""
        if isinstance(other, Env):
            self._bindings.update(other._bindings)
        else:
            self._bindings.update(other)

    def keys(self):
        """All keys visible from this scope (current + parents)."""
        return self.to_dict().keys()

    def __iter__(self):
        return iter(self.to_dict())

    # -- mutation -----------------------------------------------------------

    def define(self, name: str, value: Any) -> None:
        """Bind *name* in the **current** scope."""
        self._bindings[name] = value

    def set(self, name: str, value: Any) -> None:
        """Update *name* in the **nearest enclosing** scope that contains it.

        Raises ``KeyError`` if the name is not bound anywhere.
        """
        if name in self._bindings:
            self._bindings[name] = value
        elif self._parent is not None:
            self._parent.set(name, value)
        else:
            raise KeyError(f"Cannot set! undefined variable: {name}")

    # -- construction -------------------------------------------------------

    def extend(self, bindings: dict[str, Any] | None = None) -> Env:
        """Return a child environment."""
        return Env({} if bindings is None else bindings, parent=self)

    # -- conversion ---------------------------------------------------------

    def to_dict(self) -> dict[str, Any]:
        """Flatten the full chain into a single dict (parent first)."""
        if self._parent is not None:
            d = self._parent.to_dict()
        else:
            d = {}
        d.update(self._bindings)
        return d

    def __repr__(self) -> str:
        keys = list(self._bindings.keys())
        depth = 0
        p = self._parent
        while p:
            depth += 1
            p = p._parent
        return f"<Env depth={depth} keys={keys}>"


class MergedEnv(Env):
    """Env with two parent chains: primary (closure) and secondary (caller).

    Reads walk: local bindings → primary chain → secondary chain.
    set! walks: local bindings → primary chain (skips secondary).
    This allows set! to modify variables in the defining scope (closure)
    without being confused by overlay copies from the calling scope.
    """

    __slots__ = ("_secondary",)

    def __init__(
        self,
        bindings: dict[str, Any] | None = None,
        primary: Env | None = None,
        secondary: Env | None = None,
    ):
        super().__init__(bindings, parent=primary)
        self._secondary = secondary

    def lookup(self, name: str) -> Any:
        try:
            return super().lookup(name)
        except KeyError:
            if self._secondary is not None:
                return self._secondary.lookup(name)
            raise

    def __contains__(self, name: str) -> bool:
        if super().__contains__(name):
            return True
        if self._secondary is not None:
            return name in self._secondary
        return False

    def get(self, name: str, default: Any = None) -> Any:
        try:
            return self.lookup(name)
        except KeyError:
            return default

    def to_dict(self) -> dict[str, Any]:
        if self._secondary is not None:
            d = self._secondary.to_dict()
        else:
            d = {}
        if self._parent is not None:
            d.update(self._parent.to_dict())
        d.update(self._bindings)
        return d

    def extend(self, bindings: dict[str, Any] | None = None) -> Env:
        return Env(bindings or {}, parent=self)