from __future__ import annotations

from typing import Iterator, Optional, Protocol, runtime_checkable

from torch import Tensor
from torch.nn import Parameter


@runtime_checkable
class MemoryEncoder(Protocol):
    """Structural protocol for encoders used by Hippotorch.

    Implementations are expected to provide query/key encoders and a small
    lifecycle for momentum/EMA updates and key refresh bookkeeping.

    Only the minimal surface is specified here to keep the protocol
    implementation-agnostic; concrete encoders may expose additional
    attributes (e.g., ``projector``, ``online``) that some components
    optionally use when present.
    """

    def encode_query(
        self, x: Tensor, mask: Optional[Tensor] = None
    ) -> Tensor: ...  # pragma: no cover - protocol stub

    def encode_key(
        self, x: Tensor, mask: Optional[Tensor] = None
    ) -> Tensor: ...  # pragma: no cover - protocol stub

    def update_target(self) -> None: ...  # pragma: no cover - protocol stub

    def should_refresh_keys(self) -> bool: ...  # pragma: no cover - protocol stub

    def mark_refreshed(self) -> None: ...  # pragma: no cover - protocol stub

    # Minimal attributes used by consolidation/memory paths
    refresh_interval: int

    def parameters(self) -> Iterator[Parameter]: ...  # pragma: no cover - protocol stub


__all__ = ["MemoryEncoder"]
