"""Persisted ComfyUI node-output cache keys, records, restore, and write-back."""

from __future__ import annotations

import asyncio
import base64
from contextlib import AbstractContextManager
from dataclasses import dataclass
import hashlib
import inspect
import json
import logging
import math
from typing import Any, Callable, Iterable, Mapping, Sequence
import zlib

try:
    from .cloud_comfy_bootstrap import _load_nodes_module
    from .cloud_runtime_context import node_output_cache_store
    from .remote_protocol import (
        BOUNDARY_INPUT_SIGNATURES_KEY as _BOUNDARY_INPUT_SIGNATURES_KEY,
    )
    from .serialization import (
        deserialize_node_outputs,
        deserialize_value,
        serialize_node_outputs,
        serialize_value,
    )
    from .settings import get_settings
except ImportError:  # pragma: no cover - flat Modal-container import.
    from cloud_comfy_bootstrap import _load_nodes_module
    from cloud_runtime_context import node_output_cache_store
    from remote_protocol import (
        BOUNDARY_INPUT_SIGNATURES_KEY as _BOUNDARY_INPUT_SIGNATURES_KEY,
    )
    from serialization import (
        deserialize_node_outputs,
        deserialize_value,
        serialize_node_outputs,
        serialize_value,
    )
    from settings import get_settings

logger = logging.getLogger(__name__)

_NODE_OUTPUT_CACHE_KEY_PREFIX = "NC_"
_NODE_OUTPUT_CACHE_RECORD_VERSION = 1


@dataclass(frozen=True)
class CloudNodeOutputCacheHooks:
    """Callbacks supplied by the stable cloud entrypoint."""

    emit_cloud_info: Callable[..., None]
    timed_phase: Callable[..., AbstractContextManager[None]]


_NODE_OUTPUT_CACHE_HOOKS: CloudNodeOutputCacheHooks | None = None


def configure_cloud_node_output_cache_hooks(
    hooks: CloudNodeOutputCacheHooks,
) -> None:
    """Install cloud logging callbacks without importing upward into the entrypoint."""
    global _NODE_OUTPUT_CACHE_HOOKS
    _NODE_OUTPUT_CACHE_HOOKS = hooks


def _node_output_cache_hooks() -> CloudNodeOutputCacheHooks:
    """Return configured callbacks or fail with a clear import-order error."""
    if _NODE_OUTPUT_CACHE_HOOKS is None:
        raise RuntimeError("Cloud node-output cache hooks have not been configured.")
    return _NODE_OUTPUT_CACHE_HOOKS


def _emit_cloud_info(message: str, *args: Any) -> None:
    """Delegate timestamped cloud logging to the stable entrypoint."""
    _node_output_cache_hooks().emit_cloud_info(message, *args)


def _timed_phase(phase: str, **fields: Any) -> AbstractContextManager[None]:
    """Delegate phase timing to the stable entrypoint."""
    return _node_output_cache_hooks().timed_phase(phase, **fields)


@dataclass
class _PersistedNodeCacheRestoreState:
    """Track distributed cache entries restored into one prompt execution."""

    restored_node_ids: list[str]
    restored_cache_keys_by_node_id: dict[str, str]
    restore_original_method: Callable[[], None]


@dataclass(frozen=True)
class _NodeOutputCacheLookupResult:
    """Hold one distributed cache lookup result before live-cache hydration."""

    node_id: str
    cache_key: str | None
    raw_record: Any | None
    cache_entry: Any | None


def _is_link(value: Any) -> bool:
    """Return whether a prompt input value is a ComfyUI link."""
    return (
        isinstance(value, list)
        and len(value) == 2
        and all(not isinstance(item, dict) for item in value)
    )


def _normalize_link_output_index(value: Any) -> Any:
    """Unwrap a singleton list around a prompt-link output index when present."""
    while isinstance(value, list) and len(value) == 1:
        value = value[0]
    return value


def _node_output_cache_store() -> Any | None:
    """Return the shared Modal Dict used for persisted transport-safe node outputs."""
    return node_output_cache_store()


def _node_output_cache_key_preview(
    cache_key: str | None, *, max_chars: int = 32
) -> str:
    """Return a short human-readable prefix of one persisted node-cache key."""
    if cache_key is None:
        return "<none>"
    return cache_key[:max_chars]


def _node_output_cache_value_preview(value: Any, *, max_chars: int = 160) -> str:
    """Return a truncated repr for node-cache debug logging."""
    try:
        rendered = repr(value)
    except Exception as exc:  # pragma: no cover - defensive logging path.
        rendered = f"<repr failed: {type(exc).__name__}: {exc}>"
    if len(rendered) <= max_chars:
        return rendered
    return f"{rendered[:max_chars]}..."


def _tensor_cache_key_digest(value: Any) -> dict[str, Any]:
    """Return a stable digest payload for one tensor used inside a cache key."""
    from safetensors.torch import save

    tensor = value.detach().contiguous().cpu()
    tensor_bytes = save({"value": tensor})
    return {
        "kind": "tensor",
        "dtype": str(tensor.dtype),
        "shape": list(tensor.shape),
        "sha256": hashlib.sha256(tensor_bytes).hexdigest(),
    }


async def _node_output_cache_store_get(cache_store: Any, cache_key: str) -> Any:
    """Return one persisted node-cache record, preferring Modal's async Dict interface."""
    aio_get = getattr(getattr(cache_store, "get", None), "aio", None)
    if callable(aio_get):
        return await aio_get(cache_key)
    return cache_store.get(cache_key)


async def _node_output_cache_store_put(
    cache_store: Any,
    cache_key: str,
    record: dict[str, Any],
) -> None:
    """Persist one node-cache record without blocking an active event loop."""
    aio_put = getattr(getattr(cache_store, "put", None), "aio", None)
    if callable(aio_put):
        await aio_put(cache_key, record)
        return
    cache_store[cache_key] = record


def _is_input_signature_cache_key_set(cache_key_set: Any) -> bool:
    """Return whether one cache-key set uses ComfyUI input-signature semantics."""
    return all(
        hasattr(cache_key_set, attribute)
        for attribute in (
            "dynprompt",
            "is_changed_cache",
            "get_ordered_ancestry",
            "include_node_id_in_input",
        )
    )


def _include_unique_id_in_input_signature(class_type: str) -> bool:
    """Return whether ComfyUI includes the unique node id in this input signature."""
    from comfy_execution.caching import include_unique_id_in_input

    return bool(include_unique_id_in_input(class_type))


def _build_node_output_cache_immediate_signature(
    cache_key_set: Any,
    *,
    dynprompt: Any,
    node_id: str,
    ancestor_order_mapping: dict[str, int],
    is_changed_value: Any,
) -> list[Any]:
    """Return one raw ComfyUI input-signature fragment before `to_hashable()` runs."""
    if not dynprompt.has_node(node_id):
        return [float("NaN")]

    node = dynprompt.get_node(node_id)
    class_type = node["class_type"]
    class_def = _load_nodes_module().NODE_CLASS_MAPPINGS[class_type]
    signature: list[Any] = [class_type, is_changed_value]
    if (
        cache_key_set.include_node_id_in_input()
        or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT)
        or _include_unique_id_in_input_signature(class_type)
    ):
        signature.append(node_id)

    inputs = node["inputs"]
    boundary_input_signatures = node.get(_BOUNDARY_INPUT_SIGNATURES_KEY)
    for key in sorted(inputs.keys()):
        input_value = inputs[key]
        ancestor_socket = _cache_signature_link_output_index(input_value)
        ancestor_id = str(input_value[0]) if ancestor_socket is not None else None
        if ancestor_socket is not None and ancestor_id in ancestor_order_mapping:
            ancestor_index = int(ancestor_order_mapping[ancestor_id])
            signature.append((key, ("ANCESTOR", ancestor_index, ancestor_socket)))
        else:
            boundary_signature = None
            if isinstance(boundary_input_signatures, dict):
                candidate_signature = boundary_input_signatures.get(str(key))
                if candidate_signature is not None:
                    boundary_signature = candidate_signature
            if boundary_signature is not None:
                signature.append((key, ("BOUNDARY_SOURCE", boundary_signature)))
            else:
                signature.append((key, input_value))
    return signature


def _cache_signature_link_output_index(value: Any) -> int | None:
    """Return a prompt-link output index when `value` is safe to treat as graph wiring."""
    if not _is_link(value):
        return None

    output_index = _normalize_link_output_index(value[1])
    if isinstance(output_index, bool):
        return None

    try:
        return int(output_index)
    except (TypeError, ValueError):
        return None


async def _build_node_output_cache_signature_from_key_set_async(
    cache_key_set: Any,
    node_id: str,
) -> Any:
    """Return one distributed cache signature derived from a live ComfyUI cache-key set."""
    if not _is_input_signature_cache_key_set(cache_key_set):
        return cache_key_set.get_data_key(node_id)

    dynprompt = cache_key_set.dynprompt
    ancestors, order_mapping = cache_key_set.get_ordered_ancestry(dynprompt, node_id)
    signature = [
        _build_node_output_cache_immediate_signature(
            cache_key_set,
            dynprompt=dynprompt,
            node_id=node_id,
            ancestor_order_mapping=order_mapping,
            is_changed_value=await cache_key_set.is_changed_cache.get(node_id),
        )
    ]
    for ancestor_id in ancestors:
        signature.append(
            _build_node_output_cache_immediate_signature(
                cache_key_set,
                dynprompt=dynprompt,
                node_id=str(ancestor_id),
                ancestor_order_mapping=order_mapping,
                is_changed_value=await cache_key_set.is_changed_cache.get(ancestor_id),
            )
        )
    return signature


def _build_node_output_cache_signature_from_key_set_sync(
    cache_key_set: Any,
    node_id: str,
) -> Any | None:
    """Return one distributed cache signature using cached `is_changed` values only."""
    if not _is_input_signature_cache_key_set(cache_key_set):
        return cache_key_set.get_data_key(node_id)

    cached_is_changed = getattr(cache_key_set.is_changed_cache, "is_changed", None)
    if not isinstance(cached_is_changed, dict):
        return None

    dynprompt = cache_key_set.dynprompt
    ancestors, order_mapping = cache_key_set.get_ordered_ancestry(dynprompt, node_id)
    all_node_ids = [str(node_id), *[str(ancestor_id) for ancestor_id in ancestors]]
    missing_node_ids = [
        candidate for candidate in all_node_ids if candidate not in cached_is_changed
    ]
    if missing_node_ids:
        _emit_cloud_info(
            "Node output cache signature rebuild node=%s result=skip reason=missing-is-changed values=%s",
            node_id,
            missing_node_ids,
        )
        return None

    return [
        _build_node_output_cache_immediate_signature(
            cache_key_set,
            dynprompt=dynprompt,
            node_id=candidate,
            ancestor_order_mapping=order_mapping,
            is_changed_value=cached_is_changed[candidate],
        )
        for candidate in all_node_ids
    ]


def _canonicalize_node_output_cache_key_part(
    value: Any,
    *,
    path: str = "root",
) -> Any | None:
    """Return a JSON-stable representation of one CacheKeySetInputSignature fragment."""
    value_type_name = type(value).__name__
    if value_type_name == "Unhashable":
        _emit_cloud_info(
            "Node output cache canonicalization path=%s result=unhashable reason=comfy-unhashable-marker type=%s value=%s",
            path,
            value_type_name,
            _node_output_cache_value_preview(value),
        )
        return None
    if value is None or isinstance(value, bool | int | str):
        return value
    if isinstance(value, float):
        if math.isnan(value):
            return {"kind": "float", "value": "nan"}
        if math.isinf(value):
            return {"kind": "float", "value": "inf" if value > 0 else "-inf"}
        return value
    if isinstance(value, bytes):
        return {
            "kind": "bytes",
            "value": base64.b64encode(value).decode("ascii"),
        }
    try:
        import torch
    except ModuleNotFoundError:
        torch = None
    if torch is not None and isinstance(value, torch.Tensor):
        return _tensor_cache_key_digest(value)
    if isinstance(value, tuple):
        items = []
        for index, item in enumerate(value):
            child_path = f"{path}[{index}]"
            canonical_item = _canonicalize_node_output_cache_key_part(
                item,
                path=child_path,
            )
            if canonical_item is None:
                _emit_cloud_info(
                    "Node output cache canonicalization path=%s result=unhashable reason=tuple-child child_path=%s parent_type=%s parent_value=%s",
                    path,
                    child_path,
                    value_type_name,
                    _node_output_cache_value_preview(value),
                )
                return None
            items.append(canonical_item)
        return {"kind": "tuple", "items": items}
    if isinstance(value, list):
        items = []
        for index, item in enumerate(value):
            child_path = f"{path}[{index}]"
            canonical_item = _canonicalize_node_output_cache_key_part(
                item,
                path=child_path,
            )
            if canonical_item is None:
                _emit_cloud_info(
                    "Node output cache canonicalization path=%s result=unhashable reason=list-child child_path=%s parent_type=%s parent_value=%s",
                    path,
                    child_path,
                    value_type_name,
                    _node_output_cache_value_preview(value),
                )
                return None
            items.append(canonical_item)
        return {"kind": "list", "items": items}
    if isinstance(value, dict):
        items: list[dict[str, Any]] = []
        for key in sorted(value):
            rendered_key = _node_output_cache_value_preview(key, max_chars=48)
            key_path = f"{path}.key[{rendered_key}]"
            value_path = f"{path}[{rendered_key}]"
            canonical_key = _canonicalize_node_output_cache_key_part(
                key,
                path=key_path,
            )
            canonical_value = _canonicalize_node_output_cache_key_part(
                value[key],
                path=value_path,
            )
            if canonical_key is None:
                _emit_cloud_info(
                    "Node output cache canonicalization path=%s result=unhashable reason=dict-key child_path=%s parent_type=%s parent_value=%s",
                    path,
                    key_path,
                    value_type_name,
                    _node_output_cache_value_preview(value),
                )
                return None
            if canonical_value is None:
                _emit_cloud_info(
                    "Node output cache canonicalization path=%s result=unhashable reason=dict-value key=%s child_path=%s parent_type=%s parent_value=%s",
                    path,
                    rendered_key,
                    value_path,
                    value_type_name,
                    _node_output_cache_value_preview(value),
                )
                return None
            items.append({"key": canonical_key, "value": canonical_value})
        return {"kind": "dict", "items": items}
    if isinstance(value, frozenset):
        canonical_items: list[Any] = []
        for index, item in enumerate(
            sorted(
                value,
                key=lambda item: _node_output_cache_value_preview(item, max_chars=120),
            )
        ):
            child_path = f"{path}{{{index}}}"
            canonical_item = _canonicalize_node_output_cache_key_part(
                item,
                path=child_path,
            )
            if canonical_item is None:
                _emit_cloud_info(
                    "Node output cache canonicalization path=%s result=unhashable reason=frozenset-child child_path=%s parent_type=%s parent_value=%s",
                    path,
                    child_path,
                    value_type_name,
                    _node_output_cache_value_preview(value),
                )
                return None
            canonical_items.append(canonical_item)
        canonical_items.sort(
            key=lambda item: json.dumps(
                item,
                sort_keys=True,
                separators=(",", ":"),
            )
        )
        return {"kind": "frozenset", "items": canonical_items}
    _emit_cloud_info(
        "Node output cache canonicalization path=%s result=unhashable reason=unsupported-type type=%s value=%s",
        path,
        value_type_name,
        _node_output_cache_value_preview(value),
    )
    return None


def _node_output_cache_key(signature: Any) -> str | None:
    """Return the persisted Modal Dict key for one ComfyUI cache signature."""
    canonical_signature = _canonicalize_node_output_cache_key_part(signature)
    if canonical_signature is None:
        return None
    signature_payload = json.dumps(
        canonical_signature,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    signature_digest = hashlib.sha256(signature_payload).hexdigest()
    return f"{_NODE_OUTPUT_CACHE_KEY_PREFIX}{signature_digest}"


async def _node_output_cache_key_from_key_set_async(
    cache_key_set: Any,
    node_id: str,
) -> str | None:
    """Return the persisted Modal Dict key for one live ComfyUI cache-key-set node."""
    return _node_output_cache_key(
        await _build_node_output_cache_signature_from_key_set_async(
            cache_key_set,
            node_id,
        )
    )


def _node_output_cache_key_from_key_set_sync(
    cache_key_set: Any,
    node_id: str,
) -> str | None:
    """Return the persisted Modal Dict key for one executed ComfyUI cache-key-set node."""
    signature = _build_node_output_cache_signature_from_key_set_sync(
        cache_key_set, node_id
    )
    if signature is None:
        return None
    return _node_output_cache_key(signature)


def _estimate_node_output_cache_value_size_bytes(
    value: Any,
    *,
    byte_limit: int,
) -> int | None:
    """Return a best-effort raw-size estimate for one transport-safe value."""
    if byte_limit < 0:
        return None
    if value is None or isinstance(value, bool):
        return 1
    if isinstance(value, int):
        return 8
    if isinstance(value, float):
        return 8
    if isinstance(value, str):
        return len(value.encode("utf-8"))
    if isinstance(value, bytes):
        return len(value)

    try:
        import torch
    except ModuleNotFoundError:
        torch = None

    if torch is not None and isinstance(value, torch.Tensor):
        return int(value.numel()) * int(value.element_size())

    if isinstance(value, tuple | list):
        total_size = 0
        for item in value:
            item_size = _estimate_node_output_cache_value_size_bytes(
                item, byte_limit=byte_limit
            )
            if item_size is None:
                return None
            total_size += item_size
            if total_size > byte_limit:
                return total_size
        return total_size

    if isinstance(value, dict):
        total_size = 0
        for key, item in value.items():
            total_size += len(str(key).encode("utf-8"))
            if total_size > byte_limit:
                return total_size
            item_size = _estimate_node_output_cache_value_size_bytes(
                item, byte_limit=byte_limit
            )
            if item_size is None:
                return None
            total_size += item_size
            if total_size > byte_limit:
                return total_size
        return total_size

    return None


def _serialize_node_output_cache_entry(
    cache_entry: Any,
    *,
    max_bytes: int,
) -> dict[str, Any] | None:
    """Return a persisted node-cache record when the outputs are safe and small enough."""
    if max_bytes <= 0:
        return None

    outputs_size = _estimate_node_output_cache_value_size_bytes(
        list(getattr(cache_entry, "outputs", [])),
        byte_limit=max_bytes,
    )
    if outputs_size is None or outputs_size > max_bytes:
        return None

    try:
        serialized_outputs = serialize_node_outputs(
            tuple(getattr(cache_entry, "outputs", []))
        )
    except TypeError:
        return None

    ui_payload: Any | None = None
    ui_value = getattr(cache_entry, "ui", None)
    if ui_value is not None:
        ui_size = _estimate_node_output_cache_value_size_bytes(
            ui_value, byte_limit=max_bytes
        )
        if ui_size is not None and ui_size <= max_bytes:
            try:
                ui_payload = serialize_value(ui_value)
            except TypeError:
                ui_payload = None

    return {
        "version": _NODE_OUTPUT_CACHE_RECORD_VERSION,
        "outputs_zlib": zlib.compress(serialized_outputs),
        "outputs_size_bytes": outputs_size,
        "ui": ui_payload,
    }


def _deserialize_node_output_cache_entry(
    execution: Any,
    record: Any,
) -> Any | None:
    """Return a ComfyUI CacheEntry reconstructed from one persisted Modal Dict record."""
    if not isinstance(record, dict):
        return None
    if int(record.get("version", -1)) != _NODE_OUTPUT_CACHE_RECORD_VERSION:
        return None
    compressed_outputs = record.get("outputs_zlib")
    if not isinstance(compressed_outputs, (bytes, bytearray)):
        return None

    try:
        outputs = list(
            deserialize_node_outputs(zlib.decompress(bytes(compressed_outputs)))
        )
    except (TypeError, ValueError, zlib.error):
        return None

    ui_payload = record.get("ui")
    try:
        ui_value = deserialize_value(ui_payload) if ui_payload is not None else None
    except TypeError:
        ui_value = None
    return execution.CacheEntry(ui=ui_value, outputs=outputs)


async def _await_maybe(value: Any) -> Any:
    """Await an asynchronous compatibility result or return a synchronous value."""
    if inspect.isawaitable(value):
        return await value
    return value


def _prompt_executor_cache_get_sync(outputs_cache: Any, node_id: str) -> Any | None:
    """Read one prepared cache entry without leaking a coroutine into synchronous code."""
    local_getter = getattr(outputs_cache, "get_local", None)
    if callable(local_getter):
        return local_getter(node_id)

    cache_entry = outputs_cache.get(node_id)
    if not inspect.isawaitable(cache_entry):
        return cache_entry
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(cache_entry)
    if inspect.iscoroutine(cache_entry):
        cache_entry.close()
    raise RuntimeError(
        "ComfyUI exposed an asynchronous outputs cache without the synchronous get_local API."
    )


async def _restore_persisted_node_output_cache_entries(
    execution: Any,
    executor: Any,
    *,
    prompt_id: str,
    prompt: dict[str, Any],
    cache_store: Any,
    required_materialized_node_ids: Iterable[str] | None = None,
    restored_cache_keys_by_node_id: dict[str, str] | None = None,
) -> list[str]:
    """Hydrate PromptExecutor output-cache misses from the shared Modal Dict."""
    outputs_cache = executor.caches.outputs
    dynamic_prompt = execution.DynamicPrompt(prompt)
    is_changed_cache = execution.IsChangedCache(
        prompt_id, dynamic_prompt, outputs_cache
    )
    await _await_maybe(
        outputs_cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache)
    )
    outputs_cache.clean_unused()

    return await _restore_persisted_node_output_cache_entries_into_prepared_cache(
        execution,
        outputs_cache,
        prompt=prompt,
        cache_store=cache_store,
        required_materialized_node_ids=required_materialized_node_ids,
        restored_cache_keys_by_node_id=restored_cache_keys_by_node_id,
    )


def _node_output_cache_ancestor_ids(cache_key_set: Any, node_id: str) -> set[str]:
    """Return the prompt ancestor ids used by one ComfyUI input-signature cache key."""
    if not _is_input_signature_cache_key_set(cache_key_set):
        return set()
    dynprompt = cache_key_set.dynprompt
    if not dynprompt.has_node(node_id):
        return set()
    ancestors, _order_mapping = cache_key_set.get_ordered_ancestry(dynprompt, node_id)
    return {str(ancestor_id) for ancestor_id in ancestors}


async def _restore_persisted_node_output_cache_entries_into_prepared_cache(
    execution: Any,
    outputs_cache: Any,
    *,
    prompt: dict[str, Any],
    cache_store: Any,
    required_materialized_node_ids: Iterable[str] | None = None,
    restored_cache_keys_by_node_id: dict[str, str] | None = None,
) -> list[str]:
    """Hydrate one already-prepared PromptExecutor outputs cache from the shared Modal Dict."""
    restored_node_ids: list[str] = []
    pending_lookup_tasks: list[asyncio.Task[_NodeOutputCacheLookupResult]] = []
    required_node_ids = {
        str(node_id)
        for node_id in (required_materialized_node_ids or [])
        if str(node_id) in prompt
    }
    local_node_ids: set[str] = set()

    async def lookup_node(node_id: str) -> _NodeOutputCacheLookupResult:
        """Resolve one distributed cache candidate without mutating the live outputs cache."""
        cache_key = await _node_output_cache_key_from_key_set_async(
            outputs_cache.cache_key_set, node_id
        )
        if cache_key is None:
            return _NodeOutputCacheLookupResult(
                node_id=node_id,
                cache_key=None,
                raw_record=None,
                cache_entry=None,
            )
        raw_record = await _node_output_cache_store_get(cache_store, cache_key)
        return _NodeOutputCacheLookupResult(
            node_id=node_id,
            cache_key=cache_key,
            raw_record=raw_record,
            cache_entry=_deserialize_node_output_cache_entry(execution, raw_record),
        )

    for node_id in prompt:
        if await _await_maybe(outputs_cache.get(node_id)) is not None:
            _emit_cloud_info(
                "Node output cache lookup node=%s result=local-hit",
                node_id,
            )
            local_node_ids.add(str(node_id))
            continue
        pending_lookup_tasks.append(asyncio.create_task(lookup_node(str(node_id))))

    if not pending_lookup_tasks:
        return restored_node_ids

    lookup_results = await asyncio.gather(*pending_lookup_tasks)
    restorable_node_ids = {
        lookup_result.node_id
        for lookup_result in lookup_results
        if lookup_result.cache_key is not None and lookup_result.cache_entry is not None
    }
    available_node_ids = local_node_ids | restorable_node_ids

    for lookup_result in lookup_results:
        node_id = lookup_result.node_id
        cache_key = lookup_result.cache_key
        if cache_key is None:
            _emit_cloud_info(
                "Node output cache lookup node=%s key_prefix=%s result=skip reason=key-unhashable",
                node_id,
                _node_output_cache_key_preview(cache_key),
            )
            continue
        raw_record = lookup_result.raw_record
        cache_entry = lookup_result.cache_entry
        if cache_entry is None:
            result = "miss"
            if raw_record is not None:
                result = "miss-invalid"
            _emit_cloud_info(
                "Node output cache lookup node=%s key_prefix=%s result=%s",
                node_id,
                _node_output_cache_key_preview(cache_key),
                result,
            )
            continue
        missing_required_ancestors = sorted(
            required_node_ids
            & _node_output_cache_ancestor_ids(outputs_cache.cache_key_set, node_id)
            - available_node_ids
        )
        if missing_required_ancestors:
            _emit_cloud_info(
                "Node output cache lookup node=%s key_prefix=%s result=skip reason=missing-required-boundary-ancestors ancestors=%s",
                node_id,
                _node_output_cache_key_preview(cache_key),
                missing_required_ancestors,
            )
            continue
        await _await_maybe(outputs_cache.set(node_id, cache_entry))
        _emit_cloud_info(
            "Node output cache lookup node=%s key_prefix=%s result=hit",
            node_id,
            _node_output_cache_key_preview(cache_key),
        )
        if restored_cache_keys_by_node_id is not None:
            restored_cache_keys_by_node_id[str(node_id)] = cache_key
        restored_node_ids.append(str(node_id))
    return restored_node_ids


def _install_prompt_executor_persisted_cache_restore(
    execution: Any,
    executor: Any,
    *,
    component_id: str,
    prompt: dict[str, Any],
    cache_store: Any,
    required_materialized_node_ids: Iterable[str] | None = None,
) -> _PersistedNodeCacheRestoreState:
    """Patch one executor so persisted-cache restore runs after its live `set_prompt()` call."""
    restored_node_ids: list[str] = []
    restored_cache_keys_by_node_id: dict[str, str] = {}
    outputs_cache = executor.caches.outputs
    original_set_prompt = outputs_cache.set_prompt

    async def wrapped_set_prompt(
        dynprompt: Any, node_ids: Any, is_changed_cache: Any
    ) -> None:
        await original_set_prompt(dynprompt, node_ids, is_changed_cache)
        with _timed_phase("restore_persisted_node_cache", component=component_id):
            restored_cache_keys_by_node_id.clear()
            restored_node_ids[
                :
            ] = await _restore_persisted_node_output_cache_entries_into_prepared_cache(
                execution,
                outputs_cache,
                prompt=prompt,
                cache_store=cache_store,
                required_materialized_node_ids=required_materialized_node_ids,
                restored_cache_keys_by_node_id=restored_cache_keys_by_node_id,
            )

    outputs_cache.set_prompt = wrapped_set_prompt

    def restore_original_method() -> None:
        outputs_cache.set_prompt = original_set_prompt

    return _PersistedNodeCacheRestoreState(
        restored_node_ids=restored_node_ids,
        restored_cache_keys_by_node_id=restored_cache_keys_by_node_id,
        restore_original_method=restore_original_method,
    )


def _emit_restored_node_cache_events(
    status_callback: Callable[[dict[str, Any]], None] | None,
    restored_node_ids: Sequence[str],
) -> None:
    """Publish one status event per node whose outputs were restored from the persisted cache."""
    if status_callback is None:
        return

    for node_id in restored_node_ids:
        safe_node_id = str(node_id)
        status_callback(
            {
                "event_type": "node_cached",
                "node_id": safe_node_id,
                "display_node_id": safe_node_id,
                "real_node_id": safe_node_id,
            }
        )


def _boundary_output_node_ids(
    boundary_outputs: Iterable[Mapping[str, Any]]
) -> set[str]:
    """Return node ids whose local cache entries must exist for boundary collection."""
    return {
        str(boundary_output.get("node_id"))
        for boundary_output in boundary_outputs
        if boundary_output.get("node_id") is not None
    }


async def _persist_node_output_cache_entries(
    executor: Any,
    *,
    prompt: dict[str, Any],
    cache_store: Any,
    restored_cache_keys_by_node_id: dict[str, str] | None = None,
) -> list[str]:
    """Persist eligible PromptExecutor cache entries into the shared Modal Dict."""
    max_bytes = get_settings().node_output_cache_max_bytes
    if max_bytes <= 0:
        return []

    outputs_cache = executor.caches.outputs
    cache_key_set = getattr(outputs_cache, "cache_key_set", None)
    if cache_key_set is None:
        return []

    persisted_node_ids: list[str] = []
    for node_id in prompt:
        cache_entry = await _await_maybe(outputs_cache.get(node_id))
        if cache_entry is None:
            _emit_cloud_info(
                "Node output cache write node=%s result=skip reason=no-local-cache-entry",
                node_id,
            )
            continue
        cache_key = _node_output_cache_key_from_key_set_sync(
            cache_key_set, str(node_id)
        )
        if cache_key is None:
            _emit_cloud_info(
                "Node output cache write node=%s key_prefix=%s result=skip reason=key-unhashable",
                node_id,
                _node_output_cache_key_preview(cache_key),
            )
            continue
        restored_cache_key = None
        if restored_cache_keys_by_node_id is not None:
            restored_cache_key = restored_cache_keys_by_node_id.get(str(node_id))
        if restored_cache_key == cache_key:
            _emit_cloud_info(
                "Node output cache write node=%s key_prefix=%s result=skip reason=restored-hit",
                node_id,
                _node_output_cache_key_preview(cache_key),
            )
            continue
        record = _serialize_node_output_cache_entry(cache_entry, max_bytes=max_bytes)
        if record is None:
            _emit_cloud_info(
                "Node output cache write node=%s key_prefix=%s result=skip reason=ineligible-or-oversize",
                node_id,
                _node_output_cache_key_preview(cache_key),
            )
            continue
        await _node_output_cache_store_put(cache_store, cache_key, record)
        _emit_cloud_info(
            "Node output cache write node=%s key_prefix=%s result=write outputs_size_bytes=%s",
            node_id,
            _node_output_cache_key_preview(cache_key),
            record.get("outputs_size_bytes"),
        )
        persisted_node_ids.append(str(node_id))
    return persisted_node_ids