"""
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] = bindings or {}
        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 get(self, name: str, default: Any = None) -> Any:
        try:
            return self.lookup(name)
        except KeyError:
            return default

    # -- 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(bindings or {}, 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}>"