import os
from pathlib import Path

import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode
from hippotorch.encoders.dual import DualEncoder
from hippotorch.memory.store import MemoryStore, TieredMemoryStore


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


def _encode_keys(
    encoder: DualEncoder, memory: MemoryStore, episodes: list[Episode]
) -> torch.Tensor:
    batch = memory.episodes_to_tensor(episodes)
    return encoder.encode_key(batch)


def test_capacity_eviction_and_prune_mask(tmp_path: Path):
    enc = DualEncoder(input_dim=4, embed_dim=8)
    mem = MemoryStore(embed_dim=8, capacity=3)
    eps = [_ep(value=float(i)) for i in range(5)]
    keys = _encode_keys(enc, mem, eps)
    # Write more than capacity to trigger evictions
    mem.write(keys, eps, encoder_step=0)
    assert len(mem) == 3
    # Prune with indices including out-of-range
    before = len(mem)
    mem.prune_indices([1, 99, -1])
    assert len(mem) == before - 1


def test_from_to_checkpoint_roundtrip_tiered(tmp_path: Path):
    enc = DualEncoder(input_dim=4, embed_dim=8)
    cold_dir = tmp_path / "cache"
    mem = TieredMemoryStore(embed_dim=8, capacity=4, cold_dir=cold_dir, hot_fraction=0.5)
    eps = [_ep(value=float(i)) for i in range(4)]
    keys = _encode_keys(enc, mem, eps)
    mem.write(keys, eps, encoder_step=0)
    # Verify cold files exist
    assert any(p.name.endswith(".pt") for p in os.scandir(cold_dir))
    # Allowlist Episode for safe loading in recent PyTorch versions
    try:
        import torch.serialization as _ts  # type: ignore[attr-defined]

        if hasattr(_ts, "add_safe_globals"):
            _ts.add_safe_globals([Episode])  # type: ignore[attr-defined]
    except Exception:
        pass
    ckpt = tmp_path / "mem.pt"
    mem.save(ckpt)
    restored = MemoryStore.load(ckpt)
    assert len(restored) == len(mem)
    # Accessing a cold episode triggers lazy load and optional cache fill
    cold_idx = 0
    _ = restored._get_episode(cold_idx)  # type: ignore[attr-defined]
    # Cloning both store types should succeed
    _ = mem.clone()
    _ = restored.clone()


def test_tiered_eviction_and_prune(tmp_path: Path):
    enc = DualEncoder(input_dim=4, embed_dim=8)
    cold_dir = tmp_path / "cache2"
    mem = TieredMemoryStore(embed_dim=8, capacity=3, cold_dir=cold_dir, hot_fraction=0.33)
    eps = [_ep(value=float(i)) for i in range(5)]
    keys = _encode_keys(enc, mem, eps)
    mem.write(keys, eps, encoder_step=0)
    # Capacity enforcement should keep only 3 episodes and maintain files
    assert len(mem) == 3
    # Prune a middle index to exercise file deletion path
    mem.prune_indices([1])
    assert len(mem) == 2


def test_episodes_to_tensor_and_encoder_hook_errors():
    mem = MemoryStore(embed_dim=4, capacity=2)
    with pytest.raises(ValueError):
        _ = mem.episodes_to_tensor([])

    def bad_encoder(_ep: Episode):  # noqa: D401 - test stub
        return "not-a-tensor"  # type: ignore[return-value]

    ep = _ep()
    with pytest.raises(TypeError):
        _ = mem._episode_to_tensor(ep, episode_encoder=bad_encoder)  # type: ignore[arg-type]


def test_memory_store_clone_copies_state():
    enc = DualEncoder(input_dim=4, embed_dim=8)
    mem = MemoryStore(embed_dim=8, capacity=4)
    eps = [_ep(value=0.0), _ep(value=1.0)]
    keys = _encode_keys(enc, mem, eps)
    mem.write(keys, eps, encoder_step=42)
    clone = mem.clone(device="cpu")
    assert clone.embed_dim == mem.embed_dim
    assert len(clone) == len(mem)
    assert torch.allclose(clone.keys, mem.keys)


def test_prune_with_mask_validates_alignment():
    mem = MemoryStore(embed_dim=4, capacity=2)
    # Add one episode to make lengths mismatch
    ep = _ep()
    enc = DualEncoder(input_dim=4, embed_dim=4)
    key = _encode_keys(enc, mem, [ep])
    mem.write(key, [ep], encoder_step=0)
    with pytest.raises(ValueError):
        _ = mem._prune_with_mask(torch.ones(2, dtype=torch.bool))  # type: ignore[attr-defined]


def test_lazy_refresh_on_retrieve_updates_timestamps():
    enc = DualEncoder(input_dim=4, embed_dim=8, refresh_interval=1)
    mem = MemoryStore(embed_dim=8, capacity=4)
    eps = [_ep(value=0.0), _ep(value=1.0)]
    keys = _encode_keys(enc, mem, eps)
    mem.write(keys, eps, encoder_step=0)
    # Force staleness
    mem.key_timestamps = [0 for _ in mem.key_timestamps]
    # Build a simple normalized query to hit path
    q = torch.randn(1, 8)
    scores, episodes, indices = mem.retrieve(
        q, top_k=2, encoder=enc, current_step=10, lazy_refresh=True, allow_refresh=True
    )
    assert indices.numel() <= 2
    # timestamps should be updated to current_step for retrieved rows
    assert any(ts == 10 for ts in mem.key_timestamps)


def test_apply_key_snapshot_validates_shapes():
    mem = MemoryStore(embed_dim=4, capacity=2)
    ep = _ep()
    enc = DualEncoder(input_dim=4, embed_dim=4)
    key = _encode_keys(enc, mem, [ep])
    mem.write(key, [ep], encoder_step=0)
    with pytest.raises(ValueError):
        _ = mem.apply_key_snapshot(
            keys=mem.keys,
            key_timestamps=[0],
            episode_ids=[123, 456],
            slot_versions=[0],
        )
