# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import ast
import base64
import dataclasses
import hashlib
import inspect
import sys
import types
from typing import TYPE_CHECKING, Any, NamedTuple

from marimo._ast.transformers import DeprivateVisitor
from marimo._ast.variables import (
    get_cell_from_local,
    if_local_then_mangle,
    unmangle_local,
)
from marimo._ast.visitor import ImportData, Name, ScopedVisitor
from marimo._plugins.ui._core.ui_element import UIElement
from marimo._runtime.context import ContextNotInitializedError, get_context
from marimo._runtime.dataflow import induced_subgraph
from marimo._runtime.dataflow.topology import GraphTopology
from marimo._runtime.primitives import (
    CLONE_PRIMITIVES,
    FN_CACHE_TYPE,
    build_ref_predicate_for_primitives,
    is_data_primitive,
    is_data_primitive_container,
    is_primitive,
    is_pure_function,
)
from marimo._runtime.side_effect import CellHash, SideEffect
from marimo._runtime.state import SetFunctor, State
from marimo._runtime.watch._path import PathState
from marimo._save.cache import Cache, CacheType, HashMemoCleanup
from marimo._save.encode import (
    attempt_signed_bytes,
    common_container_to_bytes,
    data_to_buffer,
    deterministic_dumps,
    primitive_to_bytes,
    type_sign,
)
from marimo._save.stubs import maybe_get_custom_stub
from marimo._types.ids import CellId_t

if TYPE_CHECKING:
    from collections.abc import Callable
    from types import CodeType

    from marimo._ast.cell import CellImpl
    from marimo._runtime.context.types import RuntimeContext
    from marimo._runtime.dataflow import DirectedGraph
    from marimo._save.loaders import Loader

# Default hash type is generally inconsequential, there may be implications of
# malicious hash collision or performance. Malicious hash collision can be
# mitigated with a signed cache, and performance is neligible compared to the
# rest of the hashing mechanism.
DEFAULT_HASH = "sha256"


# NamedTuple over dataclass for unpacking.
class SerialRefs(NamedTuple):
    refs: set[Name]
    content_serialization: dict[Name, bytes]
    stateful_refs: set[Name]


def hash_module(code: CodeType | None, hash_type: str = DEFAULT_HASH) -> bytes:
    hash_alg = hashlib.new(hash_type, usedforsecurity=False)
    if not code:
        # Hash of zeros, in the case of no code object as a recognizable noop.
        # Artifact of typing for mypy, but reasonable fallback.
        return b"0" * len(hash_alg.digest())

    def process(code_obj: CodeType) -> None:
        # Recursively hash the constants that are also code objects
        for const in code_obj.co_consts:
            if isinstance(const, types.CodeType):
                process(const)
            elif isinstance(const, frozenset) and len(const) > 1:
                # Set literals fold to frozensets whose str()/iteration order is
                # PYTHONHASHSEED-dependent once there are 2+ elements.
                # Sort the element reprs for a deterministic order.
                hash_alg.update(
                    ",".join(sorted(map(repr, const))).encode("utf8")
                )
            else:
                hash_alg.update(str(const).encode("utf8"))
        # Concatenate the names and bytecode of the current code object
        # Will cause invalidation of variable naming at the top level

        names = [unmangle_local(name).name for name in code_obj.co_names]
        hash_alg.update(bytes("|".join(names), "utf8"))
        hash_alg.update(code_obj.co_code)

    process(code)
    return hash_alg.digest()


def hash_wrapped_functions(
    wrapped: Callable[..., Any], hash_type: str = DEFAULT_HASH
) -> bytes:
    seen = set()

    # there is a chance for a circular reference
    # likely manually created, but easy to guard against.
    def process_function(fn: Callable[..., Any]) -> bytes:
        if not inspect.isbuiltin(fn):
            fn_hash = hash_module(fn.__code__, hash_type)
        else:
            # Builtin functions are not hashable, so we use their name.
            fn_hash = type_sign(bytes(fn.__name__, "utf-8"), "builtin")
        if fn_hash not in seen and hasattr(fn, "__wrapped__"):
            child_hash = hash_wrapped_functions(fn.__wrapped__, hash_type)
            return child_hash + fn_hash
        seen.add(fn_hash)
        return fn_hash

    return process_function(wrapped)


def hash_raw_module(
    module: ast.Module, hash_type: str = DEFAULT_HASH
) -> bytes:
    # AST has to be compiled to code object prior to process.
    return hash_module(
        compile(
            DeprivateVisitor().visit(module),
            "<hash>",
            mode="exec",
            flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
        ),
        hash_type,
    )


def hash_cell_impl(cell: CellImpl, hash_type: str = DEFAULT_HASH) -> bytes:
    return hash_module(cell.body, hash_type) + hash_module(
        cell.last_expr, hash_type
    )


def hash_cell_group(
    cell_ids: set[CellId_t],
    graph: GraphTopology,
    hash_type: str = DEFAULT_HASH,
) -> bytes:
    hash_alg = hashlib.new(hash_type, usedforsecurity=False)
    hashes = []
    for cell_id in cell_ids:
        cell_impl = graph.cells[cell_id]
        hashes.append(hash_cell_impl(cell_impl, hash_alg.name))

    # Sort results post hash, to ensure deterministic ordering.
    for hashed_cell in sorted(hashes):
        hash_alg.update(hashed_cell)
    return hash_alg.digest()


def hash_cell_execution(
    cell_id: CellId_t, graph: GraphTopology, hash_type: str = DEFAULT_HASH
) -> bytes:
    ancestors = graph.ancestors(cell_id)
    return hash_cell_group(ancestors, graph, hash_type)


def get_and_update_context_from_scope(
    scope: dict[str, Any],
    scope_refs: set[Name] | None = None,
) -> RuntimeContext | None:
    """Get stateful registers"""

    # Remove non-global references
    ctx_scope = set(scope)
    if scope_refs is None:
        scope_refs = set()
    for ref in scope_refs:
        ctx_scope.discard(ref)

    # This is typically done in post execution hook, but it will not be
    # called in script mode.
    # TODO: Strip this out to allow for hash based look up. Name based
    # lookup fails for anonymous instances of state and UI Elements.
    try:
        ctx = get_context()
        ctx.ui_element_registry.register_scope(scope)
        ctx.state_registry.register_scope(scope)
        return ctx
    except ContextNotInitializedError:
        return None


@dataclasses.dataclass
class HashKey:
    hash: str
    cache_type: CacheType


class BlockHasher:
    def __init__(
        self,
        module: ast.Module,
        graph: DirectedGraph,
        cell_id: CellId_t,
        scope: dict[str, Any],
        *,
        context: ast.Module | None = None,
        pin_modules: bool = False,
        hash_type: str = DEFAULT_HASH,
        apply_content_hash: bool = True,
        scoped_refs: set[Name] | None = None,
        external: bool = False,
    ) -> None:
        """Hash the context of the module, and return a cache object.

        Hashing uses 3 combined methods: pure hashing, content addressed, and execution path:

        1) "Pure" hashing is used when a block has no references. The hash is computed from the code
           itself.

        2) "Content Addressed" hashing is used when all references are known and are shown to be
           primitive types (like a "pure" function).

        3) "Execution Path" hashing is when objects may contain state or other hidden values that are
           difficult to hash deterministically. For this, the code used to produce the object is used
           as the basis of the hash. It follows that code which does not change, will produce the same
           output. This draws inspiration from hashing methods in Nix. One notable difference between
           these methods is that Nix sandboxes all execution, preventing external file access, and
           internet. Sources of non-determinism are not accounted for in this implementation, and are
           left to the user.
           NB. The ContextExecutionPath is an extended case of ExecutionPath hashing, just utilizing
           additional context.

        In these cases, as long as the module is deterministic, the output will be deterministic.
        A side effects api is implemented to account for "uncontrollable external state" when
        utilizing builtin functions like `mo.watch.file`, but the onus is put on the user currently
        to manage external state.

        For optimization, the content hash is performed after the execution cache- however the content
        references are collected first. This deferred content hash is useful in cases like repeated
        calls to a cached function.

        Args:
            module: The code content to create a hash for (e.g. for persistent_cache, the body of the
                `With` statement).
            graph: The dataflow graph of the notebook.
            cell_id: The cell id attached to the module.
            scope: The definitions of (globals) available in execution context.
            context: The "context" of the module, is a module corresponding additional execution
                context for the cell. For instance, in persistent_cache case, this applies to the code
                prior to invocation, but still in the same cell.
            pin_modules: If True, then the module will be pinned to the version
            hash_type: The type of hash to use.
            apply_content_hash: If True, then the content hash will be attempted, otherwise only use
                execution path hash.
            scoped_refs: A set of references that cannot be traced via execution path, and must be
                accounted for via content hashing.
            external: If True, then the object was imported as a module. As such, the context should
                not be respected, and ignored.
        """

        # Hash should not be pinned to cell id
        scope = {unmangle_local(k, cell_id).name: v for k, v in scope.items()}
        self.module = DeprivateVisitor().visit(module)

        if not scoped_refs:
            scoped_refs = set()
        else:
            assert not apply_content_hash, (
                "scoped_refs should only be used with deferred hashing."
            )

        self._hash: str | None = None
        self._exe_hash: str | None = None
        self.graph = graph
        self.cell_id = cell_id
        self.pin_modules = pin_modules
        self.fn_cache: FN_CACHE_TYPE = {}

        # Empty name, so we can match and fill in cell context on load.
        self.visitor = ScopedVisitor("", ignore_local=True)
        self.visitor.visit(module)
        # Determine immediate references
        refs = set(self.visitor.refs)
        self.defs = self.visitor.defs

        # Deferred hashing (i.e. instantiation without applying content hash),
        # may yield missing references.
        self.missing: set[Name] = set()
        if not apply_content_hash:
            refs, self.missing = self.extract_missing_ref(refs, scope)

        ctx = None
        if not external:
            ctx = get_and_update_context_from_scope(scope)
        refs, _, stateful_refs = self.extract_ref_state_and_normalize_scope(
            refs, scope, ctx
        )
        self.stateful_refs = stateful_refs

        # usedforsecurity=False used to satisfy some static analysis tools.
        self.hash_alg = hashlib.new(hash_type, usedforsecurity=False)

        # Hold on to each ref type
        self.content_refs = set(refs)
        self.execution_refs = set(refs)
        self.context_refs = set(refs)

        # Default type, means that there are no references at all.
        cache_type: CacheType = "Pure"

        # Collect references that will be utilized for a content hash.
        content_serialization: dict[Name, bytes] = {}
        if refs:
            cache_type = "ContentAddressed"
            refs, content_serialization, stateful_refs = (
                self.collect_for_content_hash(
                    refs, scope, ctx, scoped_refs, apply_hash=False
                )
            )
            self.stateful_refs |= stateful_refs
        self.content_refs -= refs

        # If there are still unaccounted for references, then fallback on
        # execution hashing.
        if refs:
            cache_type = "ExecutionPath"
            to_hash = self.cells_from_refs(refs)
            if ctx:
                self.collect_and_hash_side_effects(to_hash, ctx)
            subgraph, children = induced_subgraph(
                graph, to_hash | {self.cell_id}
            )
            assert not children.get(self.cell_id, None), (
                "No children expected, please report this issue. "
                "You may have a cycle which marimo should have previously "
                "detected."
            )
            refs = self.hash_and_dequeue_execution_refs(
                refs, subgraph[self.cell_id], scope, ctx
            )
        self.execution_refs -= refs | self.content_refs

        # Remove values that should be provided by external scope.
        refs -= scoped_refs

        # If there are remaining references, they should be part of the
        # provided context.
        if refs:
            cache_type = "ContextExecutionPath"
            if ctx:
                self.collect_and_hash_side_effects({self.cell_id}, ctx)
            self.hash_and_verify_context_refs(refs, context)
        self.context_refs -= refs | self.content_refs | self.execution_refs

        # Hold on to execution_path refs before content update.
        self.exe_alg = self.hash_alg.copy()

        # Now run the content hash on the content refs.
        if apply_content_hash:
            self._apply_content_hash(content_serialization)
        elif self.missing:
            cache_type = "Deferred"

        # Finally, utilize the unrun block itself, and clean up.
        self.cache_type = cache_type
        module_hash = hash_raw_module(module, hash_type)
        self.hash_alg.update(module_hash)
        # Update execution path without content data.
        self.exe_alg.update(module_hash)

    @staticmethod
    def from_parent(
        parent: BlockHasher,
    ) -> BlockHasher:
        # Use a previous block as the basis of a new block.
        block = BlockHasher.__new__(BlockHasher)
        block.module = parent.module
        block.graph = parent.graph
        block.cell_id = parent.cell_id
        block.pin_modules = parent.pin_modules
        block.fn_cache = {}
        if parent.fn_cache is not None:
            block.fn_cache = dict(parent.fn_cache)
        block.visitor = parent.visitor
        block.defs = set(parent.defs)
        block.stateful_refs = set(parent.stateful_refs)
        block.hash_alg = parent.hash_alg.copy()
        block.exe_alg = parent.exe_alg.copy()
        block._hash = None
        block._exe_hash = parent._exe_hash  # The exe hash is the same
        block.cache_type = parent.cache_type
        block.content_refs = set(parent.content_refs)
        block.execution_refs = set(parent.execution_refs)
        block.context_refs = set(parent.context_refs)
        return block

    @property
    def raw_hash(self) -> bytes:
        assert self.hash_alg is not None, "Hash algorithm not initialized."
        return self.hash_alg.digest()

    @property
    def hash(self) -> str:
        if self._hash is None:
            assert self.hash_alg is not None, "Hash algorithm not initialized."
            self._hash = (
                base64.urlsafe_b64encode(self.raw_hash)
                .decode("utf-8")
                .strip("=")
            )
        return self._hash

    @property
    def exe_hash(self) -> str:
        if self._exe_hash is None:
            assert self.exe_alg is not None, "Hash algorithm not initialized."
            self._exe_hash = (
                base64.urlsafe_b64encode(self.exe_alg.digest())
                .decode("utf-8")
                .strip("=")
            )
        return self._exe_hash

    def __hash__(self) -> int:
        return hash(self.hash)

    @property
    def key(self) -> HashKey:
        return HashKey(
            hash=self.hash,
            cache_type=self.cache_type,
        )

    def _is_memoizable(
        self,
        local_ref: Name,
        value: Any,
        ctx: RuntimeContext,
    ) -> bool:
        """Check if a variable's content hash can be memoized."""
        defs = self.graph.definitions.get(local_ref, set())
        return (
            ctx.cache.is_memoizable(value)
            and local_ref not in self.stateful_refs
            and bool(defs)
            and self.cell_id not in defs
        )

    def _apply_content_hash(
        self, content_serialization: dict[Name, bytes]
    ) -> None:
        self._hash = None
        for ref in sorted(content_serialization):
            self.hash_alg.update(content_serialization[ref])

    def collect_for_content_hash(
        self,
        refs: set[Name],
        scope: dict[str, Any],
        ctx: RuntimeContext | None,
        scoped_refs: set[Name],
        apply_hash: bool = True,
    ) -> SerialRefs:
        self._hash = None
        refs, content_serialization, _ = (
            self.serialize_and_dequeue_content_refs(refs, scope, ctx)
        )
        # If scoped refs are present, then they are unhashable
        # and we should fallback to normal hash or fail.
        if unhashable := (refs & scoped_refs) - self.execution_refs:
            # pickle is a python default
            import pickle

            failed = []
            exceptions = []
            # By rights, could just fail here - but this final attempt should
            # provide better user experience.
            #
            # Get a transitive closure over the object, and attempt to pickle
            # each dependent object.
            closure = self.graph.get_transitive_references(
                unhashable,
                predicate=build_ref_predicate_for_primitives(
                    scope, CLONE_PRIMITIVES
                ),
            )
            closure -= set(content_serialization.keys()) | self.execution_refs
            unhashable_closure, relevant_serialization, _ = (
                self.serialize_and_dequeue_content_refs(
                    closure - unhashable, scope, ctx
                )
            )
            unhashable |= unhashable_closure
            content_serialization.update(relevant_serialization)
            refs |= unhashable_closure

            for ref in unhashable:
                try:
                    _hashed = deterministic_dumps(
                        scope[ref], self.hash_alg.name
                    )
                    content_serialization[ref] = type_sign(_hashed, "pickle")
                    refs.remove(ref)
                except (pickle.PicklingError, TypeError) as e:
                    exceptions.append(e)
                    failed.append(ref)
            if failed:
                # Ruff didn't like a lambda here
                def get_type(ref: Name) -> str:
                    return str(type(scope[ref])) if ref in scope else "missing"

                ref_list = ", ".join(
                    [
                        f"{ref}: {get_type(ref)} ({e!s})"
                        for ref, e in zip(failed, exceptions, strict=False)
                    ]
                )
                # Note ExceptionGroup nicest here, but only available in 3.11
                # ExceptionGroup(msg, exceptions)
                raise TypeError(
                    "Content addressed hash could not be utilized. "
                    "Try defining the dependent sections in a separate cell. "
                    "The unhashable arguments/ references are: " + ref_list
                )

        # Given an active thread, extract state based variables that
        # influence the graph, and hash them accordingly.
        if ctx:
            (
                refs,
                content_serialization_tmp,
                stateful_refs,
            ) = self.serialize_and_dequeue_stateful_content_refs(
                refs, scope, ctx
            )
            content_serialization.update(content_serialization_tmp)
        else:
            stateful_refs = set()

        if apply_hash:
            self._apply_content_hash(content_serialization)
        return SerialRefs(refs, content_serialization, stateful_refs)

    def extract_missing_ref(
        self,
        refs: set[Name],
        scope: dict[str, Any],
    ) -> tuple[set[Name], set[Name]]:
        _refs = set(refs)
        missing = set()
        for ref in refs:
            # The block will likely throw a NameError, so remove and defer to
            # execution.
            if ref in scope.get("__builtins__", ()):
                continue
            if ref not in scope:
                _refs.remove(ref)
                missing.add(ref)
        return _refs, missing

    def extract_ref_state_and_normalize_scope(
        self,
        refs: set[Name],
        scope: dict[str, Any],
        ctx: RuntimeContext | None = None,
    ) -> SerialRefs:
        """
        Preprocess the scope and references, and extract state references.

        This method performs the following operations:
        1. Removes references that are not present in the scope.
        2. Identifies and returns stateful references.
        3. Adjusts the scope, replacing UI elements and state setters with
           their corresponding values.

        Args:
            refs: A set of reference names.
            scope: A dictionary representing the current scope.
            ctx: An optional runtime context for stateful lookup.

        Returns:
            SerialRefs tuple containing the following elements:
                - The filtered references.
                - _
                - The stateful references.
        """
        refs = set(refs)
        stateful_refs = set()

        for ref in set(refs):
            if ref in scope.get("__builtins__", ()):
                refs.remove(ref)
                continue

            # Clean up the scope, and extract missing references.
            refs, _ = self.extract_missing_ref(refs, scope)

            # State relevant to the context, should be dependent on it's value-
            # not the object.
            value: State[Any] | None = None
            # Prefer actual object over reference.
            # Skip if the reference has already been subbed in, or if it is
            # a shadowed reference.
            if ref in scope and isinstance(scope[ref], State):
                value = scope[ref]
            elif ctx:
                value = ctx.state_registry.lookup(ref)

            if value is not None and (
                ref not in scope or isinstance(scope[ref], State)
            ):
                if isinstance(value, PathState):
                    # Path state only contains the path as value, it should
                    # contain the path contents.
                    scope[ref] = attempt_signed_bytes(
                        repr(value).encode(), "pathstate"
                    )
                else:
                    scope[ref] = attempt_signed_bytes(value(), "state")
                if ctx:
                    for state_name in ctx.state_registry.bound_names(value):
                        scope[state_name] = scope[ref]

            # Likewise, UI objects should be dependent on their value.
            ui: UIElement[Any, Any] | None = None
            if ref in scope and isinstance(scope[ref], UIElement):
                ui = scope[ref]
            elif ctx:
                ui = ctx.ui_element_registry.lookup(ref)
            if ui is not None and (
                ref not in scope or isinstance(scope[ref], UIElement)
            ):
                scope[ref] = attempt_signed_bytes(ui.value, "ui")
                if ctx:
                    for ui_name in ctx.ui_element_registry.bound_names(ui._id):
                        scope[ui_name] = scope[ref]
                # If the UI is directly consumed, then hold on to the
                # reference for proper cache update.
                stateful_refs.add(ref)

        # State Setters that are not directly consumed, are not needed.
        for ref in self.visitor.refs:
            # If the setter is consumed, let the hash be tied to the state
            # value.
            if ref in scope and isinstance(scope[ref], SetFunctor):
                stateful_refs.add(ref)
                scope[ref] = scope[ref]._state

        return SerialRefs(refs, {}, stateful_refs)

    def serialize_and_dequeue_content_refs(
        self,
        refs: set[Name],
        scope: dict[Name, Any],
        ctx: RuntimeContext | None = None,
    ) -> SerialRefs:
        """Use hashable references to update the hash object and dequeue them.

        NB. "Hashable" types are primitives, data primitives, and pure
        functions. With modules being "hashed" by version number, or ignored.

        The current "hashables" are:
         - module (with pin by version)
         - primitive (bytes, str, numbers.Number, type(None))
         - data primitive (e.g. numpy array, torch tensor)
         - external module definitions (imported anything)
         - pure functions (no state, no external dependencies)
         - pure containers of the above (list, dict, set, tuple)
         - custom types defined in CUSTOM_STUBS

        Args:
            refs: A set of reference names unaccounted for.
            scope: A dictionary representing the current scope.
            ctx: Optional runtime context for memoization.

        Returns a filtered list of remaining references that were not utilized
        in updating the hash, and a dictionary of the content serialization.
        """
        self._hash = None

        content_serialization = {}
        refs = set(refs)
        defining_cells: set[CellId_t] = set()
        # Content addressed hash is valid if every reference is accounted for
        # and can be shown to be a primitive value.
        imports = get_imports(scope)
        for local_ref in sorted(refs):
            ref = if_local_then_mangle(local_ref, self.cell_id)
            # An underscore import (e.g. `import marimo as _private`) is
            # mangled in the cell, but a cached function that references
            # it can surface the raw, unmangled name as a ref (its body
            # is hashed in the function scope). Depending on the path the
            # import may therefore appear in `imports` under either the
            # mangled `ref` or the raw `local_ref`. Accept either so the
            # module is treated as an import (and not pickled, which
            # would raise `cannot pickle 'module' object`).
            import_key: Name | None = (
                ref
                if ref in imports
                else local_ref
                if local_ref in imports
                else None
            )
            if import_key is not None:
                # TODO: There may be a way to tie this in with module watching.
                # e.g. module watcher could mutate the version number based
                # last updated timestamp.
                version = ""
                module = None
                if self.pin_modules:
                    # Fall back to the in-scope value (which may be a module
                    # stub) so its replayed `__version__` reproduces the pinned
                    # hash. The scope may hold the module under either the
                    # mangled or the raw name, mirroring `import_key`.
                    module = (
                        sys.modules.get(imports[import_key].module)
                        or scope.get(ref)
                        or scope.get(local_ref)
                    )
                    version = getattr(module, "__version__", "") or ""
                    if not version:
                        module = sys.modules.get(
                            imports[import_key].namespace
                        ) or scope.get(imports[import_key].namespace)
                        version = getattr(module, "__version__", "") or ""

                content_serialization[ref] = type_sign(
                    bytes(f"module:{ref}:{version}", "utf-8"), "module"
                )
                # No need to watch the module otherwise. If the block depends
                # on it then it should be caught when hashing the block.
                refs.remove(local_ref)
                continue
            if local_ref not in scope:
                # ref is somehow not defined, because of execution path
                # so do not utilize content hash in this case.
                continue
            value = scope[local_ref]

            # Check memo before expensive serialization
            if (
                ctx is not None
                and self._is_memoizable(local_ref, value, ctx)
                and local_ref in ctx.cache.hash_memo
            ):
                content_serialization[ref] = ctx.cache.hash_memo[local_ref]
                refs.remove(local_ref)
                continue

            serial_value = None
            if is_primitive(value):
                serial_value = primitive_to_bytes(value)
            elif is_data_primitive(value):
                serial_value = data_to_buffer(value)
            elif is_data_primitive_container(value):
                serial_value = common_container_to_bytes(value)
            elif is_pure_function(
                local_ref, value, scope, self.fn_cache, self.graph
            ):
                serial_value = hash_wrapped_functions(
                    value, self.hash_alg.name
                )
            # A restored stub — a placeholder left in scope for a value we
            # could not materialize here (e.g. a tensor whose optional dep is
            # absent in this environment). The `__marimo_unhashable__` protocol
            # marker avoids importing the stub class into the hasher.
            elif getattr(type(value), "__marimo_unhashable__", False):
                digest = getattr(value, "content_hash", "")
                type_name = getattr(value, "type_name", "") or ""
                stub_module = (
                    type_name.rsplit(".", 1)[0] if "." in type_name else ""
                )
                is_marimo_stub = stub_module.startswith("marimo")
                if digest:
                    # The stub carries the value's persisted content digest;
                    # replay it so the key matches the native content hash
                    # without recomputing `data_to_buffer(value)`.
                    content_serialization[ref] = bytes.fromhex(digest)
                elif is_marimo_stub:
                    # No digest, but the stub stands in for a marimo-owned value
                    # (e.g. a `persistent_cache` wrapper `_cache_call`). Native
                    # dequeues those with zero contribution: the external-module
                    # guard below is False for marimo modules (`not
                    # __module__.startswith("marimo")`), so on the real value it
                    # falls through to `refs.remove`. Mirror that here so a
                    # restored marimo stub does not flip the consumer's key from
                    # content- to execution-addressed.
                    pass
                elif local_ref in self.graph.definitions:
                    # No digest: the stub stands in for a graph-defined
                    # function/class/lambda, which the native run routes
                    # through the execution path (via the `__main__` branch
                    # below on the real value). Route by graph provenance so
                    # this environment reproduces the same key.
                    continue
                # A digest-less stub that is neither marimo-owned nor
                # graph-defined is dequeued with no content, matching native's
                # fall-through behaviour.
                refs.remove(local_ref)
                continue
            # An external module variable is assumed to be pure, with module
            # pinning being the mechanism for invalidation.
            elif getattr(value, "__module__", "__main__") == "__main__":
                continue
            elif stub := maybe_get_custom_stub(value):
                serial_value = stub.to_bytes()
            # External module that is not a class or function, may be some
            # container we don't know how to hash.
            # Note, function cases care caught by is_pure_function
            # And we assume all marimo cases are caught
            elif not inspect.isclass(
                value
            ) and not value.__module__.startswith("marimo"):
                continue

            if serial_value is not None:
                _hash = hashlib.new(self.hash_alg.name, usedforsecurity=False)
                _hash.update(serial_value)
                hash_digest = _hash.digest()
                content_serialization[ref] = hash_digest
                if ctx is not None and self._is_memoizable(
                    local_ref, value, ctx
                ):
                    # Local ref as a key is sufficient, because values defined in
                    # cell are not memoizable and rebinds will clean up the
                    # cache via cell life cycle.
                    ctx.cache.hash_memo[local_ref] = hash_digest
                    defining_cells |= self.graph.definitions.get(
                        local_ref, set()
                    )
            # Fall through means that the references should be dequeued.
            refs.remove(local_ref)

        # Register lifecycle cleanup — deduped, at most one per cell
        if ctx:
            for def_cell in defining_cells:
                ctx.cell_lifecycle_registry.inject(
                    def_cell,
                    HashMemoCleanup(),
                )

        return SerialRefs(refs, content_serialization, set())

    def serialize_and_dequeue_stateful_content_refs(
        self,
        refs: set[Name],
        scope: dict[str, Any],
        ctx: RuntimeContext,
    ) -> SerialRefs:
        """Determines and uses stateful references that impact the code block.

        Args:
            refs: A set of reference names.
            scope: A dictionary representing the current scope.
            ctx: Runtime context for stateful lookup.

        Returns:
            tuple of:
                - The updated references.
                - A dictionary of the content serialization.
                - additional stateful references.
        """
        refs = set(refs)
        # Determine _all_ additional relevant references
        transitive_state_refs = self.graph.get_transitive_references(
            refs, inclusive=False
        )

        # Filter for relevant stateful cases.
        refs |= set(
            filter(
                lambda ref: (
                    ctx.state_registry.lookup(ref) is not None
                    or ctx.ui_element_registry.lookup(ref) is not None
                ),
                transitive_state_refs,
            )
        )

        # Need to run extract again for the expanded ref set.
        refs, _, stateful_refs = self.extract_ref_state_and_normalize_scope(
            refs, scope, ctx
        )
        # Attempt content hash again on the extracted stateful refs.
        # Do not pass ctx — stateful values can change between cells,
        # so memoization is unsafe here.
        refs, content_serialization, _ = (
            self.serialize_and_dequeue_content_refs(refs, scope)
        )
        return SerialRefs(refs, content_serialization, stateful_refs)

    def hash_and_dequeue_execution_refs_fallback(
        self, refs: set[Name], to_hash: set[CellId_t]
    ) -> set[Name]:
        """Fallback for execution refs, when the context is not available."""
        self._hash = None
        refs = set(refs)
        for ancestor_id in to_hash:
            cell_impl = self.graph.cells[ancestor_id]
            for ref in cell_impl.defs:
                # Look for both, since mangle reference depends on the context
                # of the definition.
                if ref in refs:
                    refs.remove(ref)
                unmangled_ref, _ = unmangle_local(ref)
                if unmangled_ref in refs:
                    refs.remove(unmangled_ref)
        self.hash_alg.update(
            hash_cell_group(to_hash, self.graph, self.hash_alg.name)
        )
        return refs

    def hash_and_dequeue_execution_refs(
        self,
        refs: set[Name],
        parents: set[CellId_t],
        scope: dict[str, Any],
        ctx: RuntimeContext | None = None,
    ) -> set[Name]:
        """Determines and uses the hash of refs' cells to update the hash.

        Args:
          refs: List of references to account for in cell lookup.
          parents: A set of parent cell ids to hash.
          scope: The definitions of (globals) available in execution context.
          ctx: The runtime context to use for stateful lookups.

        Returns a list of references that were not utilized in updating the
        hash. This should only be possible in the case where a cell context is
        provided, as those references should be accounted for in that context.
        """
        if ctx is None:
            return self.hash_and_dequeue_execution_refs_fallback(refs, parents)

        self._hash = None
        # We want to do a DFS graph traversal, pruning when we either get to a
        # content, or pure hashed cell. For each cell, we either use the known
        # hash, present in CellLifeCycle OR we compute and register it.
        hashes = []
        # Recursive, but in a very round about way.
        # As such, only utilize the top layer.
        for cell_id in parents:
            cell = self.graph.cells[cell_id]

            if cell_id in ctx.cell_lifecycle_registry.registry:
                registry = ctx.cell_lifecycle_registry.registry[cell_id]
                cell_hash = [
                    obj.key for obj in registry if isinstance(obj, CellHash)
                ]
                if cell_hash:
                    assert len(cell_hash) == 1, (
                        "Cell Hash registered multiple times, this is unexpected. "
                        "Please report this issue to marimo-team/marimo"
                    )
                    hashes.append(cell_hash[0])
                    refs -= cell.defs
                    continue

            attempt = BlockHasher(
                cell.mod,
                self.graph,
                cell.cell_id,
                scope=scope,
                pin_modules=self.pin_modules,
                hash_type=self.hash_alg.name,
            )
            # Register the hash in the lifecycle registry
            ctx.cell_lifecycle_registry.inject(
                cell_id, CellHash(attempt.raw_hash)
            )
            refs -= cell.defs
            hashes.append(attempt.raw_hash)

        self.hash_alg.update(b"".join(sorted(hashes)))
        return refs

    def hash_and_verify_context_refs(
        self, refs: set[Name], context: ast.Module | None
    ) -> None:
        """Utilizes the provided context to update the hash with sanity check.

        If there are remaining references, they must be part of the provided
        context. This ensures this is the case, and updates the hash.

        Args:
          refs: List of references to account for in cell lookup.
          context: The context of the module, is a module corresponding
            additional execution context for the cell. For instance, in
            persistent_cache case, this applies to the code prior to
            invocation, but still in the same cell.
        """
        self._hash = None

        # Native save won't pass down context, so if we are here,
        # then something is wrong with the remaining references.
        assert context is not None, (
            "Execution path could not be resolved. "
            "There may be cyclic definitions in the code. "
            f"The unresolved references are: {refs}. "
            "This is unexpected, please report this issue to "
            "https://github.com/marimo-team/marimo/issues"
        )

        ref_cells = set().union(
            *[self.graph.definitions.get(ref, set()) for ref in refs]
        )
        ref_cells |= {
            cell
            for ref in refs
            if (cell := get_cell_from_local(ref, self.cell_id))
        }
        assert len(ref_cells) == 1, (
            "Inconsistent references, cannot determine execution path. "
            f"Got {ref_cells} expected set({self.cell_id}). "
            "This is unexpected, please report this issue to "
            "https://github.com/marimo-team/marimo/issues"
        )
        assert ref_cells == {self.cell_id}, (
            "Unexpected execution cell residual "
            f"{ref_cells.pop()} expected {self.cell_id}. "
            "This is unexpected, please report this issue to "
            "https://github.com/marimo-team/marimo/issues"
        )
        self.hash_alg.update(hash_raw_module(context, self.hash_alg.name))
        # refs have been accounted for at this point. Nothing to return

    def collect_and_hash_side_effects(
        self, cell_ids: set[CellId_t], ctx: RuntimeContext
    ) -> None:
        """Collect side effects from the context.

        Args:
            cell_ids: A set of cell ids to collect side effects from.
            ctx: The runtime context to use for collection.
        """
        side_effects = []
        for cell_id in cell_ids:
            if cell_id in ctx.cell_lifecycle_registry.registry:
                registry = ctx.cell_lifecycle_registry.registry[cell_id]
                side_effects.extend(
                    [
                        obj.hash
                        for obj in registry
                        if isinstance(obj, SideEffect)
                    ]
                )

        # Hash the side effects
        self.hash_alg.update(b"".join(sorted(side_effects)))

    def cells_from_refs(self, refs: set[Name]) -> set[CellId_t]:
        refs = set(refs)
        # Execution path works by just analyzing the input cells to hash.
        ancestors = self.graph.ancestors(self.cell_id)
        # Prune to only the ancestors that are tied to the references.
        ref_cells = set().union(
            *[self.graph.definitions.get(ref, set()) for ref in refs]
        )
        cell_basis = ancestors & ref_cells
        return (
            set().union(
                *[self.graph.ancestors(cell_id) for cell_id in cell_basis]
            )
            | cell_basis
        )


def get_imports(scope: dict[str, Any]) -> dict[Name, ImportData]:
    """Get the imports from the scope.

    Args:
        scope: The scope to get the imports from.

    Returns:
        A dictionary of imports.
    """
    # In cases without context, we must build the imports
    # implicitly from scope.
    imports = {
        name: ImportData(
            module=obj.__name__,
            definition=name,
        )
        for name, obj in scope.items()
        if inspect.ismodule(obj)
    }
    return imports


def cache_attempt_from_hash(
    module: ast.Module,
    graph: DirectedGraph,
    cell_id: CellId_t,
    scope: dict[str, Any],
    *,
    context: ast.Module | None = None,
    pin_modules: bool = False,
    hash_type: str = DEFAULT_HASH,
    scoped_refs: set[Name] | None = None,
    loader: Loader,
    as_fn: bool = False,
) -> Cache:
    """Hash a code block with context from the same cell, and return a cache
    object.

    Extra args
          - loader: The loader to use for cache operations.
          - as_fn: If True, then the block is treated as a function

    Returns:
      - A cache object that may, or may not be fully populated.
    """

    hasher = BlockHasher(
        module=module,
        graph=graph,
        cell_id=cell_id,
        scope=scope,
        context=context,
        pin_modules=pin_modules,
        hash_type=hash_type,
        scoped_refs=scoped_refs,
    )

    if as_fn:
        hasher.defs.clear()

    return loader.cache_attempt(
        hasher.defs,
        hasher.key,
        hasher.stateful_refs,
        glbls=scope,
    )


def content_cache_attempt_from_base(
    previous_block: BlockHasher,
    scope: dict[str, Any],
    loader: Loader,
    scoped_refs: set[Name] | None = None,
    required_refs: set[Name] | None = None,
    *,
    as_fn: bool = False,
    sensitive: bool = False,
) -> Cache:
    """Hash a code block with context from the same cell, and attempt a cache lookup.

    Args:
        previous_block: The block to base the new block on.
        scope: The scope of the new block.
        loader: The loader to use for cache operations.
        scoped_refs: A set of references that cannot be traced via execution path, and must be
            accounted for via content hashing.
        required_refs: A set of references that must be present in the scope.
        as_fn: If True, then the block is treated as a function, and will not cache definitions
            in scope.
        sensitive: If True, then the cache hash will to rehash references resolved with path
            execution. This will invalidate the cache more frequently.

    Returns:
        A cache object that may, or may not be fully populated.
    """
    if scoped_refs is None:
        scoped_refs = set()

    if required_refs is None:
        required_refs = set()

    scope = {
        unmangle_local(k, previous_block.cell_id).name: v
        for k, v in scope.items()
    }

    # Manually add back missing refs, which should now be in scope.
    scoped_refs |= previous_block.missing
    scoped_refs |= required_refs

    # refine to values present
    refs = scoped_refs & previous_block.visitor.refs
    # Required refs are made explicit in case the examined block does not
    # specify them e.g.
    # @cache
    # def foo(x):
    #    return random.random()
    # assert foo(0) != foo(1)
    refs |= required_refs
    # Assume all execution refs could be content refs
    # but only if sensitive is set.
    if sensitive:
        refs |= previous_block.execution_refs
    refs |= previous_block.content_refs
    refs |= previous_block.context_refs

    hasher = BlockHasher.from_parent(previous_block)
    ctx = get_and_update_context_from_scope(scope, required_refs)
    refs, _, stateful_refs = hasher.extract_ref_state_and_normalize_scope(
        refs, scope, ctx
    )

    refs, _content, tmp_stateful_refs = hasher.collect_for_content_hash(
        refs, scope, ctx, scoped_refs, apply_hash=True
    )
    # If the execution block covers this variable, then that's OK
    refs -= previous_block.execution_refs

    stateful_refs |= tmp_stateful_refs

    assert not refs, (
        "Content addressed hash could not be resolved. "
        "Try defining the cached block in a separate cell. "
        f"The unresolved references are: {refs}. "
    )

    if as_fn:
        hasher.defs.clear()

    return loader.cache_attempt(
        hasher.defs,
        hasher.key,
        stateful_refs,
        glbls=scope,
    )
