import pytest

torch = pytest.importorskip("torch")

from hippotorch import Episode, HippocampalReplayBuffer, IdentityDualEncoder, MemoryStore
from hippotorch.query_api import query_batch


def _make_episode(value: float = 0.0) -> Episode:
    states = torch.full((4, 2), value)
    actions = torch.zeros(4, 1)
    rewards = torch.full((4,), value + 1.0)
    dones = torch.zeros(4, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_query_batch_accepts_2d_embeddings_and_returns_attention():
    embed_dim = 4
    encoder = IdentityDualEncoder(input_dim=embed_dim)
    memory = MemoryStore(embed_dim=embed_dim, capacity=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.0)
    for offset in [0.0, 1.0]:
        buffer.add_episode(_make_episode(value=offset))

    # Pass pre-built embeddings [B, D]
    queries = torch.randn(2, embed_dim)
    result = query_batch(
        queries,
        buffer=buffer,
        top_k=2,
        return_scores=True,
        return_attention=True,
    )
    assert result.scores is not None and result.attention is not None
    assert result.query.shape == (2, embed_dim)
