from __future__ import annotations

from pathlib import Path

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


def _make_buffer() -> HippocampalReplayBuffer:
    input_dim, embed_dim = 6, 8
    enc = DualEncoder(input_dim=input_dim, embed_dim=embed_dim)
    mem = MemoryStore(embed_dim=embed_dim, capacity=16)
    buf = HippocampalReplayBuffer(memory=mem, encoder=enc, mixture_ratio=0.0)
    # Add a tiny episode to populate keys and prove roundtrip
    s = torch.randn(3, 2)
    a = torch.randn(3, 3)
    r = torch.randn(3)
    ep = Episode(states=s, actions=a, rewards=r)
    buf.add_episode(ep)
    return buf


def _with_module_prefix(sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Return a copy of a state dict with '.module.' prefixes injected."""
    out: dict[str, torch.Tensor] = {}
    for k, v in sd.items():
        # Insert a single '.module.' after the first segment to mimic DataParallel
        parts = k.split(".")
        if len(parts) > 1:
            new_key = parts[0] + ".module." + ".".join(parts[1:])
        else:
            new_key = "module." + k
        out[new_key] = v
    return out


def test_buffer_load_strips_data_parallel_prefixes(tmp_path: Path) -> None:
    # Build a normal snapshot first
    buf = _make_buffer()
    snap = buf._to_snapshot()  # type: ignore[attr-defined]
    enc = snap["encoder"]  # type: ignore[assignment]
    assert isinstance(enc, dict)

    # Inject '.module.' prefixes to emulate a DP-saved state dict
    sd = enc.get("state_dict", {})
    assert isinstance(sd, dict)
    enc["state_dict"] = _with_module_prefix(sd)  # type: ignore[assignment]

    # Save mutated snapshot and load via classmethod path
    path = tmp_path / "dp_snapshot.pt"
    torch.save(snap, path)
    loaded = HippocampalReplayBuffer.load(path)

    # Sanity check that encoder is usable after load
    # Build a tiny query consistent with the input dim of the encoder
    q = torch.randn(1, 1, 6)
    _ = loaded.encoder.encode_query(q)  # should not raise
