"""Context type definitions and internals helpers.

The engine uses a flat ``dict[str, Any]`` for context. Keys starting with ``__``
are reserved for engine internals (seed, trace, internal flags). User-defined
variables use plain identifier names.
"""

from __future__ import annotations

from typing import Any, TypedDict


class ContextInternals(TypedDict, total=False):
    """Documents the reserved internal keys. Not enforced at runtime."""

    __wp_node_seed__: int
    __wp_internal_flags__: dict[str, bool]
    __wp_trace__: list[dict[str, Any]]


Context = dict[str, Any]


#: Engine tables the text resolver reads by KEY rather than as variables.
#: `strip_internals` still drops them — its result is the variable surface, and
#: anything enumerating it (the assembler's variables panel, debug views) would
#: list `$__wp_axes__` as a user variable. Callers that need accessor syntax
#: re-attach them to their own local dict via `with_resolver_tables`.
RESOLVER_TABLES = ("__wp_axes__", "__wp_picks__", "__wp_negatives__")


def with_resolver_tables(
    stripped: dict[str, Any], source: dict[str, Any]
) -> dict[str, Any]:
    """Re-attach the accessor tables to an already-stripped ctx.

    `$var.0` and `$var.AXIS` resolve against `__wp_picks__` / `__wp_axes__`,
    which `strip_internals` removes along with every other `__` key. Without
    them the assembler rendered both accessors as literal text while the same
    template resolved one node upstream. The returned dict is for a resolver
    call only — never for anything that enumerates variables.
    """
    out = dict(stripped)
    for key in RESOLVER_TABLES:
        if key in source:
            out[key] = source[key]
    return out


def strip_internals(ctx: dict[str, Any]) -> dict[str, Any]:
    """Return a copy of ``ctx`` with engine-only + user-flagged-internal keys removed.

    Use at the *prompt render* boundary (PromptAssembler) — both engine
    bookkeeping AND user-marked-internal vars are filtered:

      1. ``__``-prefixed keys: engine bookkeeping (``__wp_node_seed__``,
         ``__wp_trace__``, ``__wp_internal_flags__``, …).
      2. Keys whose name appears with ``True`` in the
         ``__wp_internal_flags__`` map: user-marked-internal bindings
         from modules with ``instance.internal == True``. These vars
         drive composition (other modules can read them) but must not
         appear as standalone ``$var`` substitutions in the rendered
         prompt — that's the "internal" UX promise.

    For the *socket* boundary (PIPELINE_CONTEXT carrying state between
    Context nodes), use ``strip_engine_internals`` instead — it drops
    only the ``__``-prefixed engine keys so user-flagged internal vars
    continue to propagate through downstream Context / Combine / etc.
    until they reach a PromptAssembler that filters them at render
    time.
    """
    flags = ctx.get("__wp_internal_flags__")
    internal_names: set[str] = set()
    if isinstance(flags, dict):
        internal_names = {k for k, v in flags.items() if v}
    return {
        k: v for k, v in ctx.items()
        if not k.startswith("__") and k not in internal_names
    }


def strip_engine_internals(ctx: dict[str, Any]) -> dict[str, Any]:
    """Return a copy of ``ctx`` with only the ``__``-prefixed engine keys removed.

    User-flagged internal vars stay in the result — they need to
    propagate across Context-node boundaries so downstream Combine /
    Derivation / Constraint modules can read them. Only the final
    PromptAssembler filters them out at render time (via
    ``strip_internals``) so they don't appear as standalone ``$var``
    substitutions in the prompt text.
    """
    return {k: v for k, v in ctx.items() if not k.startswith("__")}