"""
Simulated intervention environment for CompanionGuard-RL.

Wraps the pre-processed dataset as a Gymnasium-compatible offline RL environment.
Each episode = one dataset sample (single-step MDP).

Observation:
  d_score(1) | l_risk_onehot(5) | c_primary_probs(10) |
  e_H_pool(H) | e_P_pool(H) | t_norm(1)

Action: Discrete(5) → {PASS, WARN, REWRITE, REJECT, CRISIS}
Reward: multi-objective safety reward from src.rl.reward
"""

import numpy as np
import gymnasium as gym
from gymnasium import spaces
from typing import Dict, Tuple, Optional, Any, List

from src.rl.reward import compute_reward
from src.utils.taxonomy import NUM_ACTIONS, NUM_PRIMARY, NUM_RISK_LEVELS
from src.utils.preprocessing import build_obs_vector


class CompanionEnv(gym.Env):
    """
    Offline simulated environment built from a pre-processed detector-annotated dataset.

    Since each sample is a one-step MDP (the intervention is decided once per AI response),
    every call to step() terminates the episode immediately (terminated=True).
    The collect_rollout loop in PPOTrainer handles auto-resets correctly.
    """

    metadata = {"render_modes": []}

    def __init__(
        self,
        samples: List[Dict],
        detector_hidden: int = 1024,
        reward_weights: Optional[Dict] = None,
        max_turns: int = 20,
    ):
        super().__init__()
        self.samples = samples
        self.detector_hidden = detector_hidden
        self.reward_weights = reward_weights
        self.max_turns = max_turns

        obs_dim = 1 + NUM_RISK_LEVELS + NUM_PRIMARY + detector_hidden * 2 + 1
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32
        )
        self.action_space = spaces.Discrete(NUM_ACTIONS)

        self._current_sample: Optional[Dict] = None

    def _get_obs(self) -> np.ndarray:
        return build_obs_vector(self._current_sample, max_turns=self.max_turns)

    def reset(
        self,
        *,
        seed: Optional[int] = None,
        options: Optional[Dict] = None,
    ) -> Tuple[np.ndarray, Dict]:
        super().reset(seed=seed)
        idx = self.np_random.integers(0, len(self.samples))
        self._current_sample = self.samples[idx]
        return self._get_obs(), {}

    def step(self, action: int) -> Tuple[np.ndarray, float, bool, bool, Dict]:
        assert self._current_sample is not None, "Call reset() before step()"

        sample = self._current_sample

        from src.utils.taxonomy import ACTION_NAME_TO_ID, category_to_index, PRIMARY_CATEGORY_LIST

        # Convert a_recommend string → int for alignment signal
        a_rec_str = sample.get("a_recommend", "PASS")
        a_rec_int = ACTION_NAME_TO_ID.get(a_rec_str, 0)

        # Use ground-truth c_primary for reward (training time has GT; deployment uses det)
        gt_category = sample.get("c_primary", "None")
        if gt_category in PRIMARY_CATEGORY_LIST:
            reward_cat_idx = category_to_index(gt_category)
        else:
            reward_cat_idx = int(sample.get("c_primary_idx", 0))

        reward = compute_reward(
            action=int(action),
            y_risk=int(sample["y_risk"]),
            l_risk=int(sample["l_risk"]),
            c_primary_idx=reward_cat_idx,
            a_recommend=a_rec_int,
        )

        # One-step MDP: always terminate
        terminated = True
        truncated = False
        info = {
            "y_risk": int(sample["y_risk"]),
            "l_risk": int(sample["l_risk"]),
            "a_recommend": sample.get("a_recommend", "PASS"),
            "action_taken": action,
        }

        # Return current obs (episode is over, but Gymnasium requires a valid obs)
        obs = self._get_obs()
        return obs, float(reward), terminated, truncated, info

    def render(self):
        pass


class BatchCompanionEnv:
    """
    Simple vectorized wrapper around multiple CompanionEnv instances.
    Used for faster rollout collection in PPO.
    """

    def __init__(
        self,
        samples: List[Dict],
        n_envs: int = 8,
        detector_hidden: int = 1024,
        reward_weights: Optional[Dict] = None,
        max_turns: int = 20,
    ):
        self.n_envs = n_envs
        self.envs = [
            CompanionEnv(
                samples=samples,
                detector_hidden=detector_hidden,
                reward_weights=reward_weights,
                max_turns=max_turns,
            )
            for _ in range(n_envs)
        ]

    def reset(self) -> np.ndarray:
        obs_list = [env.reset()[0] for env in self.envs]
        return np.stack(obs_list)

    def step(self, actions: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List]:
        obs_list, reward_list, done_list, info_list = [], [], [], []

        for env, action in zip(self.envs, actions):
            obs, reward, terminated, truncated, info = env.step(int(action))
            done = terminated or truncated

            if done:
                # Auto-reset
                obs, _ = env.reset()

            obs_list.append(obs)
            reward_list.append(reward)
            done_list.append(done)
            info_list.append(info)

        return (
            np.stack(obs_list),
            np.array(reward_list, dtype=np.float32),
            np.array(done_list, dtype=bool),
            info_list,
        )
