from __future__ import annotations

import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode
from hippotorch.encoders.dual import DualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.store import MemoryStore


def _make_episode(T: int, state_dim: int, action_dim: int) -> Episode:
    states = torch.randn(T, state_dim)
    actions = torch.randn(T, action_dim)
    rewards = torch.randn(T)
    dones = torch.zeros(T, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_prepare_query_pads_and_truncates_to_encoder_input_dim(monkeypatch):
    # Encoder expects input_dim=7
    input_dim = 7
    embed_dim = 8
    encoder = DualEncoder(input_dim=input_dim, embed_dim=embed_dim, momentum=0.0)
    memory = MemoryStore(embed_dim=embed_dim, capacity=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=1.0)

    # Add one episode so sampling path is enabled if needed
    ep = _make_episode(3, state_dim=3, action_dim=3)
    buffer.add_episode(ep, encoder_step=0)

    captured = {}

    def _spy_encode_query(x: torch.Tensor, mask=None):  # type: ignore[no-redef]
        captured["shape"] = tuple(x.shape)
        # Return a valid embedding tensor
        return torch.zeros(x.size(0), x.size(1), embed_dim)

    # Patch encode_query to observe the input to the encoder
    monkeypatch.setattr(encoder, "encode_query", _spy_encode_query)

    # Short query (dim=5) should be padded up to 7
    short = torch.randn(5)
    _ = buffer._prepare_query(short)  # type: ignore[attr-defined]
    assert captured["shape"][-1] == input_dim

    # Long query (dim=10) should be truncated down to 7
    long = torch.randn(10)
    _ = buffer._prepare_query(long)  # type: ignore[attr-defined]
    assert captured["shape"][-1] == input_dim
