"""DiscoRL-compatible CartPole environment wrapper.

This module:
1. Wraps standard Gym CartPole in DiscoRL's Environment interface
2. CartPole naturally has discrete actions (0 or 1)
3. Provides flexible observation/action preprocessing

The design supports:
- Simple batch handling (Python-level, non-JAX)
- Discrete action space (required by DiscoRL Agent)
- Standard Gym interface (reset/step)
"""

from typing import Any, Dict, Tuple, Optional
import numpy as np
import jax
import jax.numpy as jnp
import gymnasium as gym

from disco_rl.environments import base
from disco_rl import types

try:
    from dm_env import StepType
except ImportError:
    # Fallback with correct mapping
    class StepType:
        FIRST = 0
        MID = 1
        LAST = 2


class DiscoCartPoleEnv(base.Environment):
    """DiscoRL-compatible batched CartPole environment.
    
    CartPole already has discrete actions (0, 1), so no discretization needed.
    This adapter simply wraps Gym CartPole to provide DiscoRL's Environment interface
    with types.EnvironmentTimestep.
    """
    
    def __init__(
        self,
        batch_size: int = 1,
        max_steps: int = 500,
    ):
        self.batch_size = batch_size
        self.max_steps = max_steps
        self._step_counts = np.zeros(batch_size, dtype=np.int32)
        self._episode_done = np.zeros(batch_size, dtype=bool)
        
        # Create env instances
        self._envs = [gym.make('CartPole-v1') for _ in range(batch_size)]
        
        # Build specs from first env
        base_env = self._envs[0]
        
        try:
            from dm_env import specs as dm_specs
            # CartPole has action space Discrete(2), so actions are {0, 1}
            self._single_action_spec = dm_specs.BoundedArray(
                shape=(), dtype=np.int32, minimum=0, maximum=1
            )
            obs_shape = base_env.observation_space.shape
            obs_dtype = base_env.observation_space.dtype
            self._single_observation_spec = {
                'observation': dm_specs.Array(shape=obs_shape, dtype=obs_dtype)
            }
        except Exception:
            self._single_action_spec = type('ActionSpec', (), {
                'shape': (),
                'dtype': np.int32,
                'low': 0,
                'high': 1,
            })
            self._single_observation_spec = {
                'observation': base_env.observation_space
            }
        
        self._last_obs = [None] * batch_size
        self._last_info = [{}] * batch_size
    
    def single_action_spec(self):
        return self._single_action_spec
    
    def single_observation_spec(self):
        return self._single_observation_spec
    
    def step(
        self, state_unused: Any, actions: np.ndarray
    ) -> Tuple[Any, types.EnvironmentTimestep]:
        """Step all envs.
        
        Args:
            state_unused: unused (kept for DiscoRL interface compatibility)
            actions: array of shape (batch_size,) with discrete action indices (0 or 1)
        
        Returns:
            (state, timestep) where timestep is a batched EnvironmentTimestep
        """
        # Convert actions to list if needed
        if isinstance(actions, (np.ndarray, jnp.ndarray)):
            actions_list = [int(a) for a in np.asarray(actions)]
        else:
            actions_list = list(actions)
        
        obs_batch = []
        reward_batch = []
        done_batch = []
        
        for i, env in enumerate(self._envs):
            action = actions_list[i] if i < len(actions_list) else 0
            # Action should be 0 or 1 for CartPole
            action = int(action) % 2
            
            # ✅ FIX: Never auto-reset. Always step the environment.
            # If episode is done, it should have been reset by the caller.
            # This ensures correct reward propagation.
            obs, reward, terminated, truncated, info = env.step(action)
            done = bool(terminated or truncated)
            
            # Increment step counter; mark as done on terminal or max steps
            self._step_counts[i] += 1
            if done or self._step_counts[i] >= self.max_steps:
                self._episode_done[i] = True
            
            self._last_obs[i] = obs
            self._last_info[i] = info
            obs_batch.append(jnp.asarray(obs, dtype=jnp.float32))
            reward_batch.append(float(reward))
            done_batch.append(done)
        
        # Stack into batched timestep
        obs_map = {'observation': jnp.stack(obs_batch)}
        rewards = jnp.asarray(reward_batch, dtype=jnp.float32)
        is_terminal = jnp.asarray(done_batch, dtype=jnp.bool_)
        # Use LAST (2) for terminal, MID (1) for non-terminal
        step_type = jnp.where(is_terminal, StepType.LAST, StepType.MID)
        
        timestep = types.EnvironmentTimestep(
            observation=obs_map,
            step_type=step_type,
            reward=rewards,
        )
        return None, timestep
    
    def reset(self, rng_key: Optional[Any] = None) -> Tuple[Any, types.EnvironmentTimestep]:
        """Reset all envs.
        
        Args:
            rng_key: optional JAX RNG (unused here)
        
        Returns:
            (state, timestep)
        """
        obs_batch = []
        reward_batch = []
        done_batch = []
        
        for i, env in enumerate(self._envs):
            obs, info = env.reset()
            self._last_obs[i] = obs
            self._last_info[i] = info
            self._step_counts[i] = 0
            self._episode_done[i] = False
            obs_batch.append(jnp.asarray(obs, dtype=jnp.float32))
            reward_batch.append(0.0)
            done_batch.append(False)
        
        obs_map = {'observation': jnp.stack(obs_batch)}
        rewards = jnp.asarray(reward_batch, dtype=jnp.float32)
        is_terminal = jnp.asarray(done_batch, dtype=jnp.bool_)
        # Use FIRST (0) for reset, since this is the first step of a new episode
        step_type = jnp.full((len(self._envs),), StepType.FIRST, dtype=jnp.int32)
        
        timestep = types.EnvironmentTimestep(
            observation=obs_map,
            step_type=step_type,
            reward=rewards,
        )
        return None, timestep