from __future__ import annotations

import torch

from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.sb3 import SB3ReplayBufferWrapper
from hippotorch.memory.store import MemoryStore


def test_sb3_query_hook_exception_falls_back_to_raw_observation():
    state_dim, action_dim = 3, 2
    input_dim = state_dim + action_dim + 1
    encoder = IdentityDualEncoder(input_dim=input_dim)
    memory = MemoryStore(embed_dim=input_dim)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=1.0)

    # Populate memory with two short episodes at different regions
    wrapper = SB3ReplayBufferWrapper(
        buffer,
        query_state_fn=lambda _obs: (_ for _ in ()).throw(RuntimeError("hook failed")),
    )

    a = torch.zeros(action_dim)
    o0, o1 = torch.full((state_dim,), 1.0), torch.full((state_dim,), 2.0)
    wrapper.add(o0, o1, a, reward=1.0, done=False)
    wrapper.add(o1, o1 + 1.0, a, reward=1.0, done=True)

    o0b, o1b = torch.full((state_dim,), -1.0), torch.full((state_dim,), -2.0)
    wrapper.add(o0b, o1b, a, reward=0.1, done=False)
    wrapper.add(o1b, o1b - 1.0, a, reward=0.1, done=True)

    # Provide a query observation; the hook will throw, wrapper should fallback to raw obs
    q = torch.full((state_dim,), 2.0)
    batch = wrapper.sample(8, query_observation=q, top_k=1)

    # Should not crash and should return a batch with standard keys
    assert set(batch.keys()) >= {"states", "actions", "rewards", "next_states", "dones"}
    # With mixture=1.0, semantic should be > 0 even if hook failed
    stats = wrapper.sampling_diagnostics()
    assert stats.get("semantic", 0.0) > 0.0
