"""Tests for the UreaClearing-v0 Gymnasium environment."""

import gymnasium as gym
import numpy as np
import pytest

import nephrosim  # noqa: F401
from nephrosim.envs.urea_clearing import UreaClearingEnv


@pytest.fixture
def env():
    e = gym.make("nephrosim/UreaClearing-v0")
    yield e
    e.close()


class TestUreaClearingEnv:
    def test_observation_space_shape(self, env):
        obs, info = env.reset(seed=42)
        assert obs.shape == (7,)

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

    def test_action_space_shape(self, env):
        assert env.action_space.shape == (2,)

    def test_step_returns_five_tuple(self, env):
        env.reset(seed=42)
        action = env.action_space.sample()
        result = env.step(action)
        assert len(result) == 5
        obs, reward, terminated, truncated, info = result
        assert isinstance(reward, float)
        assert isinstance(terminated, bool)
        assert isinstance(truncated, bool)
        assert isinstance(info, dict)

    def test_episode_terminates_at_max_steps(self, env):
        env.reset(seed=42)
        done = False
        steps = 0
        while not done:
            obs, reward, terminated, truncated, info = env.step(
                np.array([0.0, 0.0], dtype=np.float32)
            )
            done = terminated or truncated
            steps += 1
        assert steps == 240  # default session_duration / dt

    def test_kt_v_in_info(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(env.action_space.sample())
        assert "kt_v" in info
        assert info["kt_v"] >= 0

    def test_urea_decreases_with_zero_action(self, env):
        obs, _ = env.reset(seed=42)
        initial_urea = obs[0]
        for _ in range(60):
            obs, _, _, _, _ = env.step(np.array([0.0, 0.0], dtype=np.float32))
        assert obs[0] < initial_urea

    def test_clearance_in_info(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(env.action_space.sample())
        assert "clearance" in info
        assert info["clearance"] > 0

    def test_total_urea_removed_increases(self, env):
        env.reset(seed=42)
        _, _, _, _, info1 = env.step(np.array([0.0, 0.0], dtype=np.float32))
        _, _, _, _, info2 = env.step(np.array([0.0, 0.0], dtype=np.float32))
        assert info2["total_urea_removed"] > info1["total_urea_removed"]

    def test_obs_stays_in_bounds_during_episode(self, env):
        env.reset(seed=42)
        for _ in range(100):
            obs, _, _, _, _ = env.step(env.action_space.sample())
            assert env.observation_space.contains(obs)

    def test_different_seeds_give_different_episodes(self):
        env1 = gym.make("nephrosim/UreaClearing-v0")
        env2 = gym.make("nephrosim/UreaClearing-v0")
        obs1, _ = env1.reset(seed=1)
        obs2, _ = env2.reset(seed=99)
        env1.close()
        env2.close()
        # Easy patients are deterministic, so test with medium
        env3 = UreaClearingEnv(difficulty="medium")
        env4 = UreaClearingEnv(difficulty="medium")
        obs3, _ = env3.reset(seed=1)
        obs4, _ = env4.reset(seed=99)
        assert not np.allclose(obs3, obs4)

    def test_difficulty_easy(self):
        env = UreaClearingEnv(difficulty="easy")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (7,)
        env.close()

    def test_difficulty_hard(self):
        env = UreaClearingEnv(difficulty="hard")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (7,)
        env.close()
