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

import gymnasium as gym
import numpy as np
import pytest

import nephrosim  # noqa: F401
from nephrosim.envs.full_session import FullDialysisSessionEnv


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


class TestFullDialysisSessionEnv:
    def test_observation_space_shape(self, env):
        obs, _ = env.reset(seed=42)
        assert obs.shape == (11,)

    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 == (3,)

    def test_episode_terminates(self, env):
        env.reset(seed=42)
        done = False
        steps = 0
        while not done:
            _, _, terminated, truncated, _ = env.step(
                np.array([0.0, 0.0, 0.0], dtype=np.float32)
            )
            done = terminated or truncated
            steps += 1
        assert steps == 240

    def test_info_has_kt_v(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(env.action_space.sample())
        assert "kt_v" in info

    def test_info_has_bp(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(env.action_space.sample())
        assert "systolic_bp" in info

    def test_info_has_fluid_removal(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(env.action_space.sample())
        assert "total_removed_L" in info
        assert "completion_fraction" in info

    def test_always_hard_difficulty(self):
        env = FullDialysisSessionEnv()
        env.reset(seed=42)
        assert env.patient is not None
        assert env.patient.is_diabetic
        env.close()

    def test_obs_stays_in_bounds(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_hypotension_tracked(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(env.action_space.sample())
        assert "hypotension_events" in info

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