from concurrent.futures import ThreadPoolExecutor

import pytest
import 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
from hippotorch.utils.ann import (_FAISS_AVAILABLE,
                                  ApproximateNearestNeighborIndex)


def _episode(value: float, length: int = 4) -> Episode:
    states = torch.full((length, 1), value, dtype=torch.float32)
    actions = torch.zeros(length, 1)
    rewards = torch.ones(length, dtype=torch.float32) * value
    return Episode(states=states, actions=actions, rewards=rewards)


def test_needle_in_haystack_precision():
    memory = MemoryStore(embed_dim=8, capacity=8)
    encoder = DualEncoder(input_dim=3, embed_dim=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder)
    for idx in range(4):
        buffer.add_episode(_episode(float(idx)))

    needle = torch.zeros(1, encoder.projector[0].out_features)
    needle[:, 2] = 1.0
    scores, _, indices = buffer.memory.retrieve(needle, top_k=1)
    assert indices[0, 0].item() >= 0
    assert torch.isfinite(scores).all()


@pytest.mark.skipif(not _FAISS_AVAILABLE, reason="FAISS not installed")
def test_retrieval_matches_faiss():
    dim = 32
    keys = torch.randn(256, dim)
    queries = torch.randn(8, dim)

    faiss_index = ApproximateNearestNeighborIndex(use_faiss=True)
    torch_index = ApproximateNearestNeighborIndex(use_faiss=False)

    faiss_index.rebuild(keys)
    torch_index.rebuild(keys)

    scores_faiss, idx_faiss = faiss_index.search(queries, top_k=5)
    scores_torch, idx_torch = torch_index.search(queries, top_k=5)

    assert torch.allclose(scores_faiss, scores_torch, atol=1e-5)
    assert torch.equal(idx_faiss, idx_torch)


def test_eviction_and_lazy_refresh_interaction():
    memory = MemoryStore(embed_dim=16, capacity=4)
    encoder = DualEncoder(input_dim=3, embed_dim=16, refresh_interval=1)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder)
    for val in [0.0, 1.0, 2.0]:
        buffer.add_episode(_episode(val))

    memory.key_timestamps = [0 for _ in range(len(memory))]
    stale_key = memory.keys[0].clone()
    memory.keys[0] = torch.randn_like(memory.keys[0])

    query = buffer.encoder.encode_query(torch.zeros(1, 1, 3))
    memory.retrieve(
        query,
        top_k=2,
        encoder=buffer.encoder,
        current_step=10,
        lazy_refresh=True,
    )
    assert not torch.allclose(memory.keys[0], stale_key)

    evicted_id = memory.episode_ids[0]
    memory.prune_indices([0])
    _, _, indices = memory.retrieve(query, top_k=2)
    current_ids = []
    for idx in indices.flatten().tolist():
        if 0 <= idx < len(memory.episode_ids):
            current_ids.append(memory.episode_ids[idx])
    assert evicted_id not in current_ids


@pytest.mark.slow
def test_multithreaded_refresh_is_stable():
    memory = MemoryStore(embed_dim=8, capacity=16)
    encoder = DualEncoder(input_dim=3, embed_dim=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder)
    for val in range(5):
        buffer.add_episode(_episode(float(val)))

    query = buffer.encoder.encode_query(torch.zeros(1, 1, 3))

    def refresh_worker():
        memory.refresh_keys(range(len(memory)), encoder)

    def retrieval_worker():
        for _ in range(5):
            memory.retrieve(query, top_k=2, encoder=encoder)

    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = [
            executor.submit(refresh_worker),
            executor.submit(retrieval_worker),
        ]
        for fut in futures:
            fut.result()
