import dataclasses
import abc
from typing import Any, Generic, TypeVar

import numpy as np
import torch

from src.core.env import BaseEnv


T = TypeVar("T")

@dataclasses.dataclass
class BaseRunnerConfig:
    num_envs: int = 1
    device: str = "cpu"
    history_length: int = 0  # 0 = plain obs, >0 = append (obs, action) history

    # ── Domain randomization (sim-to-real) ─────────────────────────
    # Empty dict = disabled (every field below is a no-op).  Supported keys:
    #   qpos_noise_std:      float  — Gaussian sensor noise on joint angles (rad)
    #   qvel_noise_std:      float  — Gaussian sensor noise on joint velocities (rad/s)
    #   action_delay_steps:  [lo, hi] — per-env integer control-step latency
    #   friction_scale:      [lo, hi] — per-env multiplier on Coulomb friction
    #   damping_scale:       [lo, hi] — per-env multiplier on viscous damping
    #   torque_scale:        [lo, hi] — per-env multiplier on applied motor torque
    # With history_length > 0 the policy can implicitly infer the sampled
    # dynamics from the recent (obs, action) window — end-to-end adaptation.
    domain_rand: dict = dataclasses.field(default_factory=dict)

class BaseRunner(abc.ABC, Generic[T]):
    def __init__(self, env: BaseEnv, config: T) -> None:
        self.env = env
        self.config = config

        # Resolve "auto" device before anything uses it
        if getattr(self.config, "device", None) == "auto":
            self.config.device = "cuda" if torch.cuda.is_available() else "cpu"

        self._last_actions: torch.Tensor | None = None

        self._sim_initialize(config)

        self.observation_space = self.env.observation_space
        self.action_space = self.env.action_space
        self.num_agents: int = 1  # single-agent RL (required by skrl)

        self.step_counts = torch.zeros(
            self.config.num_envs, dtype=torch.long, device=self.config.device
        )

        # ── Domain randomization (latency / sensor noise / dynamics) ─
        self._setup_domain_rand()

        # ── History buffer (implicit adaptation input) ────────────
        self._history_len: int = getattr(self.config, "history_length", 0)
        if self._history_len > 0:
            obs_dim = self.observation_space.shape[0]
            act_dim = self.action_space.shape[0]
            self._history_step_dim = obs_dim + act_dim  # each step stores (obs, action)
            # Ring buffer: (num_envs, history_length, obs_dim + act_dim)
            self._history_buf = torch.zeros(
                self.config.num_envs, self._history_len, self._history_step_dim,
                device=self.config.device,
            )

            # Policy obs = [raw_obs, history_flat]
            from gymnasium import spaces
            aug_dim = obs_dim + self._history_len * self._history_step_dim
            self.observation_space = spaces.Box(
                low=-torch.inf, high=torch.inf, shape=(aug_dim,),
            )

    @property
    @abc.abstractmethod
    def num_envs(self) -> int:
        ...

    @property
    @abc.abstractmethod
    def device(self) -> torch.device:
        ...

    @abc.abstractmethod
    def _sim_initialize(self, config: T) -> None:
        ...

    @abc.abstractmethod
    def _sim_step(self, actions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        ...

    @abc.abstractmethod
    def _sim_reset(self, env_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Reset the given envs; return FULL-batch (num_envs, nq/nv) state.

        Returning the full batch (not just the reset envs) lets GPU
        backends hand back zero-copy views without host synchronisation —
        the caller indexes the reset rows itself.
        """
        ...

    def _sim_close(self) -> None:
        """Release simulator resources. Override for extra cleanup."""
        if hasattr(self, "_offscreen_renderer") and self._offscreen_renderer is not None:
            self._offscreen_renderer.close()

    # ── Domain randomization ─────────────────────────────────────

    _SCALE_FIELDS = ("friction_scale", "damping_scale", "torque_scale")

    def _setup_domain_rand(self) -> None:
        """Parse the domain_rand config into per-env buffers.

        All buffers are no-ops when ``domain_rand`` is empty: scales are 1.0,
        delay is 0 and noise std is 0.
        """
        dr = dict(getattr(self.config, "domain_rand", {}) or {})
        n = self.config.num_envs
        dev = self.config.device

        # Fixed (not per-env) Gaussian sensor noise.
        self._qpos_noise_std = float(dr.get("qpos_noise_std", 0.0))
        self._qvel_noise_std = float(dr.get("qvel_noise_std", 0.0))

        # Per-env multiplicative dynamics scales (applied by the sim runner).
        self._dr_scales: dict[str, torch.Tensor] = {
            f: torch.ones(n, device=dev) for f in self._SCALE_FIELDS
        }
        self._dr_scale_ranges: dict[str, tuple[float, float]] = {}
        for f in self._SCALE_FIELDS:
            rng = dr.get(f)
            if rng:
                self._dr_scale_ranges[f] = (float(rng[0]), float(rng[1]))

        # Per-env integer action delay (in control steps).
        self._dr_delay = torch.zeros(n, dtype=torch.long, device=dev)
        delay_range = dr.get("action_delay_steps")
        if delay_range:
            self._delay_range = (int(delay_range[0]), int(delay_range[1]))
            self._max_delay = int(delay_range[1])
        else:
            self._delay_range = (0, 0)
            self._max_delay = 0

        # Action-delay ring buffer: (num_envs, max_delay + 1, act_dim).
        if self._max_delay > 0:
            act_dim = self.env.action_space.shape[0]
            self._action_buf = torch.zeros(
                n, self._max_delay + 1, act_dim, device=dev,
            )

    def _resample_domain_rand(self, env_ids: torch.Tensor) -> None:
        """Sample fresh per-env DR factors (call on every (re)set)."""
        if env_ids.numel() == 0:
            return
        dev = self.config.device
        for name, (lo, hi) in self._dr_scale_ranges.items():
            vals = torch.rand(env_ids.numel(), device=dev) * (hi - lo) + lo
            self._dr_scales[name][env_ids] = vals
        if self._max_delay > 0:
            self._dr_delay[env_ids] = torch.randint(
                self._delay_range[0], self._delay_range[1] + 1,
                (env_ids.numel(),), device=dev,
            )

    def _reset_action_buffer(self, env_ids: torch.Tensor) -> None:
        if self._max_delay > 0:
            self._action_buf[env_ids] = 0.0

    def _apply_action_delay(self, actions: torch.Tensor) -> torch.Tensor:
        """Return the per-env delayed action that the simulator should apply.

        The policy's commanded action is what gets stored in history; only
        the action handed to ``_sim_step`` is delayed.
        """
        if self._max_delay <= 0:
            return actions
        self._action_buf = torch.roll(self._action_buf, 1, dims=1)
        self._action_buf[:, 0] = actions
        idx = torch.arange(self.num_envs, device=self.device)
        return self._action_buf[idx, self._dr_delay]

    def _add_sensor_noise(
        self, qpos: torch.Tensor, qvel: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if self._qpos_noise_std > 0:
            qpos = qpos + torch.randn_like(qpos) * self._qpos_noise_std
        if self._qvel_noise_std > 0:
            qvel = qvel + torch.randn_like(qvel) * self._qvel_noise_std
        return qpos, qvel

    def _compute_obs(self, qpos: torch.Tensor, qvel: torch.Tensor) -> torch.Tensor:
        """Observation the policy sees — built from noisy (sensor) state."""
        nqpos, nqvel = self._add_sensor_noise(qpos, qvel)
        return self.env.compute_observations(self.env.build_state(nqpos, nqvel))

    # ── Observation augmentation ─────────────────────────────────

    def _augment_obs(self, obs: torch.Tensor) -> torch.Tensor:
        """Append the flattened (obs, action) history when enabled."""
        if self._history_len <= 0:
            return obs
        hist_flat = self._history_buf.reshape(obs.shape[0], -1)
        return torch.cat([obs, hist_flat], dim=-1)

    def _push_history(self, obs: torch.Tensor, actions: torch.Tensor,
                      env_ids: torch.Tensor | None = None) -> None:
        """Push (obs, action) into the ring buffer (shift left, append right)."""
        if self._history_len <= 0:
            return
        step = torch.cat([obs, actions.reshape(obs.shape[0], -1)], dim=-1)
        if env_ids is None:
            # All envs.
            self._history_buf = torch.roll(self._history_buf, -1, dims=1)
            self._history_buf[:, -1] = step
        else:
            self._history_buf[env_ids] = torch.roll(
                self._history_buf[env_ids], -1, dims=1
            )
            self._history_buf[env_ids, -1] = step[env_ids]

    def _reset_history(self, env_ids: torch.Tensor) -> None:
        """Zero the history buffer for reset envs."""
        if self._history_len > 0:
            self._history_buf[env_ids] = 0.0

    def reset(self) -> tuple[torch.Tensor, dict[str, Any]]:
        all_ids = torch.arange(self.num_envs, device=self.device)
        self._resample_domain_rand(all_ids)
        self._reset_action_buffer(all_ids)
        qpos, qvel = self._sim_reset(all_ids)
        self.step_counts.zero_()
        self._reset_history(all_ids)

        obs = self._compute_obs(qpos, qvel)
        return self._augment_obs(obs), {}

    def step(self, actions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
        prev_actions = (
            self._last_actions
            if self._last_actions is not None
            else torch.zeros_like(actions)
        )
        self._last_actions = actions
        # Latency: the simulator applies a (per-env) delayed action.
        sim_actions = self._apply_action_delay(actions)
        qpos, qvel = self._sim_step(sim_actions)
        self.step_counts += 1

        # Reward / termination use the TRUE state (no sensor noise) so the
        # learning signal and safety checks stay clean.
        clean_state = self.env.build_state(qpos, qvel)
        rewards = self.env.compute_rewards(clean_state, actions, prev_actions)
        terminated = self.env.compute_terminations(clean_state)
        truncated = self.env.compute_truncations(self.step_counts)

        # The observation the policy sees is built from the NOISY sensor state.
        obs = self._compute_obs(qpos, qvel)

        # Push current (obs, action) into history before augmenting.
        self._push_history(obs, actions)

        info: dict[str, Any] = {}

        done = terminated | truncated
        done_ids = done.nonzero(as_tuple=False).squeeze(-1)

        if done_ids.numel() > 0:
            info["final_observations"] = self._augment_obs(obs)[done_ids].clone()
            info["final_env_ids"] = done_ids.clone()

            # New episode → fresh dynamics + cleared latency buffer.
            self._resample_domain_rand(done_ids)
            self._reset_action_buffer(done_ids)
            full_qpos, full_qvel = self._sim_reset(done_ids)
            self.step_counts[done_ids] = 0
            self._reset_history(done_ids)

            # _sim_reset returns the full batch — index the reset rows here.
            obs[done_ids] = self._compute_obs(
                full_qpos[done_ids], full_qvel[done_ids],
            )

        # skrl expects (num_envs, 1) for rewards/terminated/truncated
        return self._augment_obs(obs), rewards.unsqueeze(-1), terminated.unsqueeze(-1), truncated.unsqueeze(-1), info
        
    def _render_frame(self, env_idx: int = 0) -> np.ndarray:
        """Return a raw RGB frame. Override in subclass."""
        raise NotImplementedError("Render not implemented for this runner.")

    def render(self, env_idx: int = 0) -> np.ndarray:
        """Render frame with action overlay."""
        frame = self._render_frame(env_idx)
        if self._last_actions is not None:
            ctrl = float(self._last_actions[env_idx, 0].clamp(-1.0, 1.0))
            _draw_action_overlay(frame, ctrl)
        return frame

    def close(self) -> None:
        self._sim_close()


def _draw_action_overlay(frame: np.ndarray, action: float) -> None:
    """Draw an action bar on a rendered frame (no OpenCV needed).

    Bar is centered horizontally: green to the right (+), red to the left (-).
    """
    h, w = frame.shape[:2]

    bar_y = h - 30
    bar_h = 16
    bar_x_center = w // 2
    bar_half_w = w // 4
    bar_x_left = bar_x_center - bar_half_w
    bar_x_right = bar_x_center + bar_half_w

    # Background (dark grey)
    frame[bar_y:bar_y + bar_h, bar_x_left:bar_x_right] = [40, 40, 40]

    # Filled bar
    fill_len = int(abs(action) * bar_half_w)
    if action > 0:
        color = [60, 200, 60]  # green
        x0 = bar_x_center
        x1 = min(bar_x_center + fill_len, bar_x_right)
    else:
        color = [200, 60, 60]  # red
        x1 = bar_x_center
        x0 = max(bar_x_center - fill_len, bar_x_left)
    frame[bar_y:bar_y + bar_h, x0:x1] = color

    # Center tick mark (white)
    frame[bar_y:bar_y + bar_h, bar_x_center - 1:bar_x_center + 1] = [255, 255, 255]