from __future__ import annotations

import numpy as np
import pytest

gym = pytest.importorskip("gymnasium")

torch = pytest.importorskip("torch")

from gymnasium import spaces

from hippotorch.core.episode import Episode
from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.store import MemoryStore
from hippotorch.memory.wrapper import HippotorchMemoryWrapper


class DummyEnv(gym.Env):  # type: ignore[misc]
    metadata = {"render_modes": []}

    def __init__(self, obs_dim: int = 3, horizon: int = 3):
        super().__init__()
        self.obs_dim = obs_dim
        self.horizon = horizon
        self.t = 0
        high = np.ones((obs_dim,), dtype=np.float32)
        self.observation_space = spaces.Box(-high, high, dtype=np.float32)
        self.action_space = spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32)

    def reset(self, *, seed: int | None = None, options: dict | None = None):
        super().reset(seed=seed)
        self.t = 0
        obs = np.zeros((self.obs_dim,), dtype=np.float32)
        return obs, {}

    def step(self, action):
        self.t += 1
        obs = np.full((self.obs_dim,), float(self.t), dtype=np.float32)
        reward = 0.0
        terminated = self.t >= self.horizon
        truncated = False
        info = {}
        return obs, reward, terminated, truncated, info


def _make_buffer(state_dim: int, action_dim: int, top_k: int = 3) -> HippocampalReplayBuffer:
    input_dim = state_dim + action_dim + 1
    encoder = IdentityDualEncoder(input_dim=input_dim)
    memory = MemoryStore(embed_dim=input_dim, capacity=16)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.0)

    # Add a couple episodes to ensure query returns something
    for v in [0.0, 1.0]:
        T = 4
        states = torch.full((T, state_dim), v)
        actions = torch.zeros(T, action_dim)
        rewards = torch.zeros(T)
        dones = torch.zeros(T, dtype=torch.bool)
        buffer.add_episode(Episode(states=states, actions=actions, rewards=rewards, dones=dones))
    return buffer


def test_wrapper_single_observation_shapes():
    obs_dim, act_dim, top_k = 4, 1, 3
    buffer = _make_buffer(obs_dim, act_dim, top_k)
    env = DummyEnv(obs_dim)
    wrapped = HippotorchMemoryWrapper(env, buffer, top_k=top_k)

    obs, info = wrapped.reset()
    assert isinstance(obs, dict)
    assert obs["obs"].shape == (obs_dim,)
    # memory vector: embed_dim + top_k
    assert obs["memory"].shape == (buffer.memory.embed_dim + top_k,)

    obs2, reward, terminated, truncated, info2 = wrapped.step(np.zeros((1,), dtype=np.float32))
    assert isinstance(obs2, dict)
    assert obs2["obs"].shape == (obs_dim,)
    assert obs2["memory"].shape == (buffer.memory.embed_dim + top_k,)


def test_wrapper_batched_observation_shapes():
    obs_dim, act_dim, top_k = 3, 1, 2
    buffer = _make_buffer(obs_dim, act_dim, top_k)

    # Create a dummy vectorized env that returns batched obs
    class BatchedEnv(DummyEnv):  # type: ignore[misc]
        def __init__(self, n_envs: int, obs_dim: int):
            super().__init__(obs_dim=obs_dim)
            self.num_envs = n_envs

        def reset(self, *, seed: int | None = None, options: dict | None = None):
            obs, info = super().reset(seed=seed, options=options)
            return np.stack([obs for _ in range(self.num_envs)], axis=0), {}

        def step(self, action):
            obs, reward, terminated, truncated, info = super().step(action)
            obs_b = np.stack([obs for _ in range(self.num_envs)], axis=0)
            return obs_b, reward, terminated, truncated, info

    env = BatchedEnv(n_envs=4, obs_dim=obs_dim)
    wrapped = HippotorchMemoryWrapper(env, buffer, top_k=top_k)
    obs, info = wrapped.reset()
    assert isinstance(obs, dict)
    assert obs["obs"].shape == (env.num_envs, obs_dim)
    assert obs["memory"].shape == (env.num_envs, buffer.memory.embed_dim + top_k)
