import pytest

torch = pytest.importorskip("torch")

from hippotorch import Episode, MemoryStore


def _episode(value: float = 0.0) -> Episode:
    states = torch.full((3, 2), value)
    actions = torch.zeros(3, 1)
    rewards = torch.ones(3) * value
    dones = torch.zeros(3, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_apply_key_snapshot_rejects_stale_slots():
    embed_dim = 4
    store = MemoryStore(embed_dim=embed_dim, capacity=4)
    keys = torch.randn(2, embed_dim)
    store.write(keys, [_episode(0.0), _episode(1.0)], encoder_step=0)

    metadata = store.snapshot_metadata()
    timestamps = list(store.key_timestamps)
    new_keys = store.keys + 1.0

    store.slot_versions[0] += 1  # simulate concurrent refresh
    before_key = store.keys[0].clone()

    applied, dropped = store.apply_key_snapshot(
        keys=new_keys,
        key_timestamps=timestamps,
        episode_ids=metadata["episode_ids"],
        slot_versions=metadata["slot_versions"],
    )
    assert applied == 1
    assert dropped == 1
    assert torch.equal(store.keys[0], before_key)
    assert torch.allclose(store.keys[1], new_keys[1])
