"""Tests for AdaptiveRT-v0 environment."""

import gymnasium as gym
import numpy as np
import pytest

import oncosim  # noqa: F401


class TestAdaptiveRTEnv:
    @pytest.fixture
    def env(self):
        e = gym.make("oncosim/AdaptiveRT-v0")
        yield e
        e.close()

    def test_env_creation(self, env):
        assert env is not None

    def test_reset_valid(self, env):
        obs, info = env.reset(seed=42)
        assert isinstance(obs, dict)
        assert "treatment_progress" in obs
        assert "plan_quality" in obs

    def test_observation_space(self, env):
        obs, _ = env.reset(seed=42)
        assert env.observation_space.contains(obs)

    def test_action_space_multidiscrete(self, env):
        assert isinstance(env.action_space, gym.spaces.MultiDiscrete)
        assert list(env.action_space.nvec) == [2, 5]

    def test_step_valid(self, env):
        env.reset(seed=42)
        action = np.array([0, 2], dtype=np.int64)
        obs, reward, terminated, truncated, info = env.step(action)
        assert isinstance(reward, float)
        assert np.isfinite(reward)

    def test_replan_action(self, env):
        env.reset(seed=42)
        # Run a few sessions first to let quality degrade
        for _ in range(5):
            env.step(np.array([0, 2], dtype=np.int64))
        obs_before, _, _, _, _ = env.step(np.array([0, 2], dtype=np.int64))
        quality_before = obs_before["plan_quality"][0]
        obs_after, _, _, _, info = env.step(np.array([1, 2], dtype=np.int64))
        if info.get("replanned", False):
            assert obs_after["plan_quality"][0] >= quality_before - 0.1

    def test_no_replan_action(self, env):
        env.reset(seed=42)
        obs, _, _, _, info = env.step(np.array([0, 2], dtype=np.int64))
        assert not info.get("replanned", False)

    def test_replans_decrement(self, env):
        obs, _ = env.reset(seed=42)
        initial_replans = obs["replans_remaining"][0]
        obs, _, _, _, _ = env.step(np.array([1, 2], dtype=np.int64))
        assert obs["replans_remaining"][0] == initial_replans - 1

    def test_episode_completion(self, env):
        env.reset(seed=42)
        done = False
        steps = 0
        while not done and steps < 50:
            action = env.action_space.sample()
            _, _, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            steps += 1
        assert done
        assert steps <= 35

    def test_all_difficulty_tiers(self):
        for diff in ["stable", "moderate", "variable"]:
            e = gym.make("oncosim/AdaptiveRT-v0", difficulty=diff)
            obs, _ = e.reset(seed=42)
            assert e.observation_space.contains(obs)
            e.close()

    def test_seed_reproducibility(self, env):
        obs1, _ = env.reset(seed=42)
        obs2, _ = env.reset(seed=42)
        assert np.array_equal(obs1["plan_quality"], obs2["plan_quality"])

    def test_treatment_progress_increases(self, env):
        env.reset(seed=42)
        for i in range(5):
            obs, _, _, _, _ = env.step(np.array([0, 2], dtype=np.int64))
            expected_progress = (i + 1) / 30
            assert obs["treatment_progress"][0] == pytest.approx(expected_progress, abs=0.01)

    def test_info_dict_keys(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([0, 2], dtype=np.int64))
        assert "session" in info
        assert "plan_quality" in info
        assert "tumor_response" in info
        assert "replans_remaining" in info

    def test_plan_quality_range(self, env):
        env.reset(seed=42)
        for _ in range(30):
            obs, _, terminated, _, _ = env.step(env.action_space.sample())
            assert 0 <= obs["plan_quality"][0] <= 1
            if terminated:
                break
