from __future__ import annotations

from typing import Callable, Dict, List, Optional

import torch
from torch import Tensor

from hippotorch.core.episode import Episode
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.segmenters.base import EpisodeSegmenter, TerminalSegmenter


class SB3ReplayBufferWrapper:
    """Adapter that exposes a stable-baselines3 style API.

    The wrapper buffers transitions into trajectories using the provided
    ``EpisodeSegmenter`` and forwards finalized episodes to the underlying
    ``HippocampalReplayBuffer``. Sampling returns tensors keyed to align with
    SB3's replay buffer outputs (observations, actions, rewards, next_observations, dones).
    """

    def __init__(
        self,
        buffer: HippocampalReplayBuffer,
        segmenter: Optional[EpisodeSegmenter] = None,
        *,
        device: torch.device | str = "cpu",
        query_state_fn: Optional[Callable[[Tensor], Tensor]] = None,
        default_top_k: int = 5,
    ) -> None:
        self.buffer = buffer
        self.segmenter = segmenter or TerminalSegmenter()
        self.device = torch.device(device)
        self._trajectory: Dict[str, List[Tensor]] = self._empty_trajectory()
        self._last_obs: Optional[Tensor] = None
        self._query_state_fn = query_state_fn
        self._default_top_k = int(default_top_k)

    @property
    def size(self) -> int:
        """Number of episodes currently stored."""

        return self.buffer.size

    def add(
        self,
        obs: Tensor | float | int,
        next_obs: Tensor | float | int,
        action: Tensor | float | int,
        reward: float,
        done: bool,
        *,
        truncated: Optional[bool] = None,
        info: Optional[dict] = None,
    ) -> None:
        """Record a transition and flush completed trajectories into memory.

        Args:
            obs: Current observation.
            next_obs: Next observation after action.
            action: Action taken.
            reward: Scalar reward.
            done: Episode terminal flag (environment termination).
            truncated: Optional truncation flag (time-limit or external cutoff).
            info: Optional info dict (unused; kept for API compatibility).
        """

        _ = info  # unused placeholder for API compatibility
        state_t = self._to_tensor(obs)
        action_t = self._to_tensor(action)
        reward_t = torch.as_tensor(reward, device=self.device, dtype=torch.float32).view(1)
        done_flag = bool(done) or bool(truncated) if truncated is not None else bool(done)
        done_t = torch.as_tensor(done_flag, device=self.device)

        self._trajectory["states"].append(state_t)
        self._trajectory["actions"].append(action_t)
        self._trajectory["rewards"].append(reward_t)
        self._trajectory["dones"].append(done_t)

        # Track the most recent observation for semantic queries at sample-time
        self._last_obs = self._to_tensor(next_obs)

        if done_flag:
            episode = self._trajectory_to_episode(self._trajectory)
            for segment in self.segmenter.segment(episode):
                self.buffer.add_episode(segment)
            self._trajectory = self._empty_trajectory()
            self._last_obs = None

    def sample(
        self,
        batch_size: int,
        *,
        query_observation: Optional[Tensor] = None,
        top_k: Optional[int] = None,
    ) -> Dict[str, Tensor]:
        """Draw a batch of transitions using hybrid sampling.

        If ``query_observation`` is provided or a recent observation is available
        from ``add()``, a semantic query is constructed and passed to the
        underlying buffer. When no query is available, the buffer may fall back
        to uniform or prioritized sampling depending on configuration.
        """

        q = query_observation if query_observation is not None else self._last_obs
        if q is not None and self._query_state_fn is not None:
            try:
                q = self._query_state_fn(q)
            except Exception:
                # Best-effort: if hook fails, fall back to raw observation.
                pass
        return self.buffer.sample(
            batch_size=batch_size,
            query_state=q,
            top_k=int(top_k) if top_k is not None else self._default_top_k,
        )

    def sampling_diagnostics(self) -> Dict[str, float]:
        """Expose the latest sampling diagnostics from the underlying buffer."""
        return self.buffer.sampling_stats()

    def consolidate(self, **kwargs) -> Dict[str, float]:
        """Proxy to the underlying buffer consolidation method."""

        return self.buffer.consolidate(**kwargs)

    def _trajectory_to_episode(self, trajectory: Dict[str, List[Tensor]]) -> Episode:
        states = torch.stack(trajectory["states"]).to(self.device)
        actions = torch.stack(trajectory["actions"]).to(self.device)
        rewards = torch.stack(trajectory["rewards"]).squeeze(-1).to(self.device)
        dones = torch.stack(trajectory["dones"]).to(self.device)
        return Episode(states=states, actions=actions, rewards=rewards, dones=dones)

    def _empty_trajectory(self) -> Dict[str, List[Tensor]]:
        return {"states": [], "actions": [], "rewards": [], "dones": []}

    def _to_tensor(self, array_like) -> Tensor:
        tensor = array_like if isinstance(array_like, torch.Tensor) else torch.as_tensor(array_like)
        if tensor.dim() == 0:
            tensor = tensor.unsqueeze(0)
        return tensor.to(self.device)


__all__ = ["SB3ReplayBufferWrapper"]
