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.hub import (LocalHubBackend, load_memory_from_hub,
                                  push_memory_to_hub)


def _make_buffer(embed_dim: int = 8) -> HippocampalReplayBuffer:
    memory = MemoryStore(embed_dim=embed_dim, capacity=8)
    encoder = DualEncoder(input_dim=5, embed_dim=embed_dim)
    return HippocampalReplayBuffer(memory=memory, encoder=encoder)


def test_save_to_hub_delegates(monkeypatch, tmp_path):
    buffer = _make_buffer()
    called = {}

    def fake_push(obj, repo_id, **kwargs):  # pragma: no cover - delegated call
        called["buffer"] = obj
        called["repo_id"] = repo_id
        called["kwargs"] = kwargs
        return tmp_path

    monkeypatch.setattr("hippotorch.utils.hub.push_memory_to_hub", fake_push)

    result = buffer.save_to_hub("user/test-repo", filename="memory.pt", private=True)

    assert result == tmp_path
    assert called["buffer"] is buffer
    assert called["repo_id"] == "user/test-repo"
    assert called["kwargs"]["filename"] == "memory.pt"
    assert called["kwargs"]["private"] is True


def test_load_memory_from_hub_replaces_store(monkeypatch):
    buffer = _make_buffer()
    loaded_memory = MemoryStore(
        embed_dim=buffer.memory.embed_dim, capacity=2, device=buffer.memory.device
    )

    monkeypatch.setattr(
        "hippotorch.utils.hub.load_memory_from_hub", lambda *_, **__: loaded_memory
    )

    result = buffer.load_memory_from_hub("user/test-repo", device=buffer.memory.device)

    assert result is loaded_memory
    assert buffer.memory is loaded_memory


def test_load_memory_from_hub_validates_embed_dim(monkeypatch):
    buffer = _make_buffer(embed_dim=4)
    incompatible_memory = MemoryStore(embed_dim=6)

    monkeypatch.setattr(
        "hippotorch.utils.hub.load_memory_from_hub",
        lambda *_, **__: incompatible_memory,
    )

    with pytest.raises(ValueError):
        buffer.load_memory_from_hub("user/test-repo")


def test_buffer_save_load_roundtrip(tmp_path):
    buffer = _make_buffer()
    episode = Episode(
        states=torch.randn(5, 3),
        actions=torch.zeros(5, 1),
        rewards=torch.linspace(0, 1, 5),
    )
    buffer.add_episode(episode)
    buffer.encoder.eval()
    query = torch.randn(1, 1, episode.states.shape[1] + episode.actions.shape[1] + 1)
    before_scores, _, before_indices = buffer.memory.retrieve(
        buffer.encoder.encode_query(query), encoder=buffer.encoder, current_step=0
    )

    target = tmp_path / "buffer.pt"
    buffer.save(target)
    restored = HippocampalReplayBuffer.load(target)
    restored.encoder.eval()

    after_scores, _, after_indices = restored.memory.retrieve(
        restored.encoder.encode_query(query),
        encoder=restored.encoder,
        current_step=0,
    )

    assert torch.allclose(before_scores, after_scores, atol=1e-3)
    assert before_indices.tolist() == after_indices.tolist()


def test_local_hub_roundtrip(tmp_path):
    buffer = _make_buffer()
    ep = Episode(
        states=torch.randn(4, 3),
        actions=torch.zeros(4, 1),
        rewards=torch.ones(4),
    )
    buffer.add_episode(ep)
    buffer.encoder.eval()
    backend = LocalHubBackend(root_dir=tmp_path)

    push_memory_to_hub(buffer, "user/local-repo", backend=backend, filename="bundle.pt")
    loaded_buffer = load_memory_from_hub(
        "user/local-repo",
        backend=backend,
        filename="bundle.pt",
        return_buffer=True,
    )
    loaded_buffer.encoder.eval()

    query = torch.randn(1, 1, ep.states.shape[1] + ep.actions.shape[1] + 1)
    before = buffer.memory.retrieve(
        buffer.encoder.encode_query(query), encoder=buffer.encoder, current_step=0
    )[0]
    after = loaded_buffer.memory.retrieve(
        loaded_buffer.encoder.encode_query(query),
        encoder=loaded_buffer.encoder,
        current_step=0,
    )[0]
    assert torch.allclose(before, after, atol=1e-3)
