from __future__ import annotations

from collections import deque
from typing import Any, Callable, Deque, Dict, Optional, Tuple

import cv2
import gym
import gym_super_mario_bros
import numpy as np
from gym.spaces import Box
from gym_super_mario_bros.actions import RIGHT_ONLY, SIMPLE_MOVEMENT
from nes_py.wrappers import JoypadSpace


def reset_compat(env: gym.Env, seed: Optional[int] = None) -> Tuple[np.ndarray, Dict[str, Any]]:
    try:
        result = env.reset(seed=seed)
    except TypeError:
        result = env.reset()
        if seed is not None and hasattr(env, "seed"):
            env.seed(seed)

    if isinstance(result, tuple) and len(result) == 2:
        obs, info = result
        return obs, info
    return result, {}


def step_compat(env: gym.Env, action: Any) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
    result = env.step(action)
    if isinstance(result, tuple) and len(result) == 5:
        obs, reward, terminated, truncated, info = result
        return obs, float(reward), bool(terminated), bool(truncated), info
    if isinstance(result, tuple) and len(result) == 4:
        obs, reward, done, info = result
        return obs, float(reward), bool(done), False, info
    raise RuntimeError(f"Unexpected step return format: {type(result)} / {result}")


class SkipFrame(gym.Wrapper):
    def __init__(self, env: gym.Env, skip: int = 4):
        super().__init__(env)
        self._skip = skip

    def step(self, action: Any):
        total_reward = 0.0
        terminated = False
        truncated = False
        info: Dict[str, Any] = {}
        obs = None
        for _ in range(self._skip):
            obs, reward, terminated, truncated, info = step_compat(self.env, action)
            total_reward += reward
            if terminated or truncated:
                break
        return obs, total_reward, terminated, truncated, info


class PreprocessFrame(gym.ObservationWrapper):
    """Convert RGB frame to grayscale 84x84 uint8."""

    def __init__(self, env: gym.Env, width: int = 84, height: int = 84):
        super().__init__(env)
        self.width = width
        self.height = height
        self.observation_space = Box(low=0, high=255, shape=(height, width, 1), dtype=np.uint8)

    def observation(self, observation: np.ndarray) -> np.ndarray:
        if observation.ndim == 3 and observation.shape[2] == 3:
            gray = cv2.cvtColor(observation, cv2.COLOR_RGB2GRAY)
        elif observation.ndim == 2:
            gray = observation
        else:
            gray = np.squeeze(observation)
        resized = cv2.resize(gray, (self.width, self.height), interpolation=cv2.INTER_AREA)
        return resized[:, :, None].astype(np.uint8)


class ChannelLastFrameStack(gym.Wrapper):
    """Stack frames on the channel axis: (H, W, C*num_stack)."""

    def __init__(self, env: gym.Env, num_stack: int = 4):
        super().__init__(env)
        self.num_stack = num_stack
        self.frames: Deque[np.ndarray] = deque(maxlen=num_stack)

        obs_space = env.observation_space
        assert isinstance(obs_space, Box), "Frame stack requires Box observation space."
        h, w, c = obs_space.shape
        self.observation_space = Box(
            low=0,
            high=255,
            shape=(h, w, c * num_stack),
            dtype=np.uint8,
        )

    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
        del options
        obs, info = reset_compat(self.env, seed=seed)
        self.frames.clear()
        for _ in range(self.num_stack):
            self.frames.append(obs)
        return self._get_observation(), info

    def step(self, action: Any):
        obs, reward, terminated, truncated, info = step_compat(self.env, action)
        self.frames.append(obs)
        return self._get_observation(), reward, terminated, truncated, info

    def _get_observation(self) -> np.ndarray:
        assert len(self.frames) == self.num_stack
        return np.concatenate(list(self.frames), axis=2)


class TransposeObservation(gym.ObservationWrapper):
    """Convert observation from HWC to CHW for CNN policy."""

    def __init__(self, env: gym.Env):
        super().__init__(env)
        obs_space = env.observation_space
        assert isinstance(obs_space, Box), "TransposeObservation requires Box observation space."
        h, w, c = obs_space.shape
        self.observation_space = Box(low=0, high=255, shape=(c, h, w), dtype=obs_space.dtype)

    def observation(self, observation: np.ndarray) -> np.ndarray:
        return np.transpose(observation, (2, 0, 1)).astype(np.uint8)


class ClipRewardEnv(gym.RewardWrapper):
    def reward(self, reward):
        return float(np.sign(reward))


class ProgressRewardEnv(gym.Wrapper):
    """Reward shaping focused on moving right and avoiding local traps."""

    def __init__(
        self,
        env: gym.Env,
        progress_scale: float = 0.02,
        death_penalty: float = -50.0,
        flag_bonus: float = 100.0,
        stall_penalty: float = 0.05,
        stall_steps: int = 40,
        backward_penalty_scale: float = 0.01,
        milestone_interval: int = 32,
        milestone_bonus: float = 1.0,
        no_progress_terminate_steps: int = 120,
        no_progress_terminate_penalty: float = 20.0,
    ):
        super().__init__(env)
        self.progress_scale = progress_scale
        self.death_penalty = death_penalty
        self.flag_bonus = flag_bonus
        self.stall_penalty = stall_penalty
        self.stall_steps = stall_steps
        self.backward_penalty_scale = backward_penalty_scale
        self.milestone_interval = milestone_interval
        self.milestone_bonus = milestone_bonus
        self.no_progress_terminate_steps = no_progress_terminate_steps
        self.no_progress_terminate_penalty = no_progress_terminate_penalty
        self._last_x_pos: Optional[float] = None
        self._best_x_pos = 0.0
        self._stall_count = 0
        self._next_milestone_x = float(milestone_interval)

    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
        del options
        obs, info = reset_compat(self.env, seed=seed)
        self._last_x_pos = float(info.get("x_pos", 0.0))
        self._best_x_pos = self._last_x_pos
        self._stall_count = 0
        if self.milestone_interval > 0:
            k = int(self._best_x_pos // self.milestone_interval) + 1
            self._next_milestone_x = float(k * self.milestone_interval)
        return obs, info

    def step(self, action: Any):
        obs, reward, terminated, truncated, info = step_compat(self.env, action)
        x_pos = float(info.get("x_pos", 0.0))

        if self._last_x_pos is None:
            delta_x = 0.0
        else:
            delta_x = x_pos - self._last_x_pos
        self._last_x_pos = x_pos

        shaped_reward = float(reward)
        if delta_x > 0:
            shaped_reward += self.progress_scale * delta_x
            self._stall_count = 0
            if x_pos > self._best_x_pos:
                self._best_x_pos = x_pos
            if self.milestone_interval > 0 and self.milestone_bonus != 0.0:
                while x_pos >= self._next_milestone_x:
                    shaped_reward += self.milestone_bonus
                    self._next_milestone_x += self.milestone_interval
        else:
            self._stall_count += 1
            if delta_x < 0:
                shaped_reward -= self.backward_penalty_scale * abs(delta_x)

        if self._stall_count >= self.stall_steps:
            shaped_reward -= self.stall_penalty

        if (
            self.no_progress_terminate_steps > 0
            and not terminated
            and not truncated
            and self._stall_count >= self.no_progress_terminate_steps
        ):
            truncated = True
            shaped_reward -= self.no_progress_terminate_penalty
            info["terminated_by_stall"] = True
            info["done_reason"] = "no_progress"

        if terminated or truncated:
            if bool(info.get("flag_get", False)):
                shaped_reward += self.flag_bonus
            elif terminated:
                shaped_reward += self.death_penalty

        return obs, shaped_reward, terminated, truncated, info


class TimePenaltyHardStuckEnv(gym.Wrapper):
    """Optional living cost and hard-stuck truncation based on x_pos movement."""

    def __init__(
        self,
        env: gym.Env,
        time_penalty: float = 0.0,
        hard_stuck_steps: int = 0,
        hard_stuck_epsilon: float = 1.0,
        hard_stuck_penalty: float = 5.0,
    ):
        super().__init__(env)
        if time_penalty > 0.0:
            raise ValueError(f"time_penalty must be <= 0.0, got {time_penalty}")
        if hard_stuck_steps < 0:
            raise ValueError(f"hard_stuck_steps must be >= 0, got {hard_stuck_steps}")
        if hard_stuck_epsilon < 0.0:
            raise ValueError(f"hard_stuck_epsilon must be >= 0.0, got {hard_stuck_epsilon}")
        if hard_stuck_penalty < 0.0:
            raise ValueError(f"hard_stuck_penalty must be >= 0.0, got {hard_stuck_penalty}")

        self.time_penalty = float(time_penalty)
        self.hard_stuck_steps = int(hard_stuck_steps)
        self.hard_stuck_epsilon = float(hard_stuck_epsilon)
        self.hard_stuck_penalty = float(hard_stuck_penalty)
        self._last_x_pos: Optional[float] = None
        self._hard_stuck_count = 0

    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
        del options
        obs, info = reset_compat(self.env, seed=seed)
        self._last_x_pos = float(info.get("x_pos", 0.0))
        self._hard_stuck_count = 0
        return obs, info

    def step(self, action: Any):
        obs, reward, terminated, truncated, info = step_compat(self.env, action)
        x_pos = float(info.get("x_pos", 0.0))
        if self._last_x_pos is None:
            delta_x = 0.0
        else:
            delta_x = x_pos - self._last_x_pos
        self._last_x_pos = x_pos

        shaped_reward = float(reward)
        if self.time_penalty != 0.0:
            shaped_reward += self.time_penalty

        if self.hard_stuck_steps > 0 and not terminated and not truncated:
            if abs(delta_x) < self.hard_stuck_epsilon:
                self._hard_stuck_count += 1
            else:
                self._hard_stuck_count = 0

            if self._hard_stuck_count >= self.hard_stuck_steps:
                shaped_reward -= self.hard_stuck_penalty
                truncated = True
                info["terminated_by_hard_stuck"] = True
                info["done_reason"] = "hard_stuck"
        else:
            self._hard_stuck_count = 0

        return obs, shaped_reward, terminated, truncated, info


def get_action_set(name: str):
    name = name.lower().strip()
    if name == "simple":
        return SIMPLE_MOVEMENT
    if name == "right_only":
        return RIGHT_ONLY
    raise ValueError(f"Unsupported movement='{name}'. Use one of: right_only, simple")


def make_mario_env(
    env_id: str = "SuperMarioBros-1-1-v0",
    seed: int = 0,
    movement: str = "right_only",
    reward_mode: str = "raw",
    clip_reward: bool = False,
    frame_skip: int = 4,
    render_mode: Optional[str] = None,
    progress_scale: float = 0.02,
    death_penalty: float = -50.0,
    flag_bonus: float = 100.0,
    stall_penalty: float = 0.05,
    stall_steps: int = 40,
    backward_penalty_scale: float = 0.01,
    milestone_interval: int = 32,
    milestone_bonus: float = 1.0,
    no_progress_terminate_steps: int = 120,
    no_progress_terminate_penalty: float = 20.0,
    time_penalty: float = 0.0,
    hard_stuck_steps: int = 0,
    hard_stuck_epsilon: float = 1.0,
    hard_stuck_penalty: float = 5.0,
) -> gym.Env:
    kwargs: Dict[str, Any] = {}
    if render_mode is not None:
        kwargs["render_mode"] = render_mode

    try:
        env = gym_super_mario_bros.make(env_id, apply_api_compatibility=True, **kwargs)
    except TypeError:
        env = gym_super_mario_bros.make(env_id, **kwargs)

    env = JoypadSpace(env, get_action_set(movement))
    env = SkipFrame(env, skip=frame_skip)

    mode = reward_mode.lower().strip()
    if clip_reward and mode == "raw":
        mode = "clip"
    if mode == "clip":
        env = ClipRewardEnv(env)
    elif mode == "progress":
        env = ProgressRewardEnv(
            env=env,
            progress_scale=progress_scale,
            death_penalty=death_penalty,
            flag_bonus=flag_bonus,
            stall_penalty=stall_penalty,
            stall_steps=stall_steps,
            backward_penalty_scale=backward_penalty_scale,
            milestone_interval=milestone_interval,
            milestone_bonus=milestone_bonus,
            no_progress_terminate_steps=no_progress_terminate_steps,
            no_progress_terminate_penalty=no_progress_terminate_penalty,
        )
    elif mode != "raw":
        raise ValueError(f"Unsupported reward_mode='{reward_mode}'. Use one of: raw, clip, progress")

    if time_penalty != 0.0 or hard_stuck_steps > 0:
        env = TimePenaltyHardStuckEnv(
            env=env,
            time_penalty=time_penalty,
            hard_stuck_steps=hard_stuck_steps,
            hard_stuck_epsilon=hard_stuck_epsilon,
            hard_stuck_penalty=hard_stuck_penalty,
        )

    env = PreprocessFrame(env, width=84, height=84)
    env = ChannelLastFrameStack(env, num_stack=4)
    env = TransposeObservation(env)

    # Seed once so each env subprocess/dummy env has deterministic startup.
    reset_compat(env, seed=seed)
    if hasattr(env.action_space, "seed"):
        env.action_space.seed(seed)
    return env


def make_env_fn(
    env_id: str,
    seed: int,
    movement: str,
    reward_mode: str,
    clip_reward: bool,
    frame_skip: int,
    progress_scale: float,
    death_penalty: float,
    flag_bonus: float,
    stall_penalty: float,
    stall_steps: int,
    backward_penalty_scale: float,
    milestone_interval: int,
    milestone_bonus: float,
    no_progress_terminate_steps: int,
    no_progress_terminate_penalty: float,
    time_penalty: float,
    hard_stuck_steps: int,
    hard_stuck_epsilon: float,
    hard_stuck_penalty: float,
) -> Callable[[], gym.Env]:
    def _thunk() -> gym.Env:
        return make_mario_env(
            env_id=env_id,
            seed=seed,
            movement=movement,
            reward_mode=reward_mode,
            clip_reward=clip_reward,
            frame_skip=frame_skip,
            render_mode=None,
            progress_scale=progress_scale,
            death_penalty=death_penalty,
            flag_bonus=flag_bonus,
            stall_penalty=stall_penalty,
            stall_steps=stall_steps,
            backward_penalty_scale=backward_penalty_scale,
            milestone_interval=milestone_interval,
            milestone_bonus=milestone_bonus,
            no_progress_terminate_steps=no_progress_terminate_steps,
            no_progress_terminate_penalty=no_progress_terminate_penalty,
            time_penalty=time_penalty,
            hard_stuck_steps=hard_stuck_steps,
            hard_stuck_epsilon=hard_stuck_epsilon,
            hard_stuck_penalty=hard_stuck_penalty,
        )

    return _thunk
