from __future__ import annotations

from typing import Callable, Optional

import numpy as np
import torch

try:  # pragma: no cover - optional dependency
    import gymnasium as gym
    from gymnasium import spaces
except Exception:  # pragma: no cover - handled lazily
    gym = None
    spaces = None

from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.query_api import QueryResult, query

if gym is None or spaces is None:  # pragma: no cover - optional dependency

    class HippotorchMemoryWrapper:  # type: ignore[misc]
        def __init__(self, *_, **__):
            raise ImportError(
                "gymnasium is required for HippotorchMemoryWrapper. "
                "Install hippotorch[envs] to pull in the dependency."
            )

        def reset(self, **kwargs):
            raise RuntimeError("gymnasium is required to use HippotorchMemoryWrapper")

        def step(self, action):
            raise RuntimeError("gymnasium is required to use HippotorchMemoryWrapper")

else:

    class HippotorchMemoryWrapper(gym.Wrapper):  # type: ignore[misc]
        """Observation wrapper that injects memory retrieval features."""

        def __init__(
            self,
            env: "gym.Env",
            buffer: HippocampalReplayBuffer,
            *,
            query_state_fn: Optional[Callable[[np.ndarray], np.ndarray]] = None,
            top_k: int = 5,
        ) -> None:
            super().__init__(env)
            self.buffer = buffer
            self.query_state_fn = query_state_fn
            self.top_k = top_k
            self._embed_dim = buffer.memory.embed_dim
            self._memory_dim = self._embed_dim + top_k
            low = np.full((self._memory_dim,), -np.inf, dtype=np.float32)
            high = np.full((self._memory_dim,), np.inf, dtype=np.float32)
            self.observation_space = spaces.Dict(
                {
                    "obs": env.observation_space,
                    "memory": spaces.Box(low=low, high=high, dtype=np.float32),
                }
            )
            self._last_memory = np.zeros(self._memory_dim, dtype=np.float32)
            self._last_query: Optional[QueryResult] = None

        def reset(self, **kwargs):
            obs, info = self.env.reset(**kwargs)
            mem = self._compute_memory(obs)
            return {"obs": obs, "memory": mem}, info

        def step(self, action):
            obs, reward, terminated, truncated, info = self.env.step(action)
            mem = self._compute_memory(obs)
            return {"obs": obs, "memory": mem}, reward, terminated, truncated, info

        def _compute_memory(self, obs) -> np.ndarray:
            result = query(
                observation=obs,
                buffer=self.buffer,
                query_state_fn=self._wrap_query_fn(),
                top_k=self.top_k,
                return_scores=True,
                return_attention=False,
            )
            self._last_query = result
            vec = np.zeros(self._memory_dim, dtype=np.float32)
            if result.query.numel():
                q = result.query.view(-1).cpu().numpy().astype(np.float32)
                take = min(q.size, self._embed_dim)
                vec[:take] = q[:take]
            if result.scores is not None and result.scores.numel():
                scores = result.scores.view(-1).cpu().numpy().astype(np.float32)
                take = min(scores.size, self.top_k)
                vec[self._embed_dim : self._embed_dim + take] = scores[:take]
            self._last_memory = vec
            return vec

        def _wrap_query_fn(self):
            if self.query_state_fn is None:
                return None

            def _builder(obs_tensor: torch.Tensor) -> torch.Tensor:
                arr = np.asarray(obs_tensor.cpu().numpy(), dtype=np.float32)
                built = self.query_state_fn(arr)
                return torch.as_tensor(
                    built, dtype=torch.float32, device=obs_tensor.device
                )

            return _builder


__all__ = ["HippotorchMemoryWrapper"]
