"""Integration tests for OncoSim package."""

import gymnasium as gym
import numpy as np
import pytest

import oncosim


class TestPackageIntegration:
    def test_import_oncosim(self):
        assert hasattr(oncosim, "__version__")

    def test_version_attribute(self):
        assert oncosim.__version__ == "0.1.0"

    def test_env_registration(self):
        env_ids = [
            "oncosim/BeamSelection-v0",
            "oncosim/DoseFractionation-v0",
            "oncosim/AdaptiveRT-v0",
        ]
        for env_id in env_ids:
            spec = gym.spec(env_id)
            assert spec is not None

    def test_make_all_envs(self):
        env_ids = [
            "oncosim/BeamSelection-v0",
            "oncosim/DoseFractionation-v0",
            "oncosim/AdaptiveRT-v0",
        ]
        for env_id in env_ids:
            env = gym.make(env_id)
            assert env is not None
            env.close()

    def test_full_episode_beam(self):
        env = gym.make("oncosim/BeamSelection-v0")
        obs, _ = env.reset(seed=42)
        total_reward = 0.0
        done = False
        while not done:
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            done = terminated or truncated
        assert np.isfinite(total_reward)
        env.close()

    def test_full_episode_dose(self):
        env = gym.make("oncosim/DoseFractionation-v0")
        obs, _ = env.reset(seed=42)
        total_reward = 0.0
        done = False
        while not done:
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            done = terminated or truncated
        assert np.isfinite(total_reward)
        env.close()

    def test_full_episode_adaptive(self):
        env = gym.make("oncosim/AdaptiveRT-v0")
        obs, _ = env.reset(seed=42)
        total_reward = 0.0
        done = False
        while not done:
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            done = terminated or truncated
        assert np.isfinite(total_reward)
        env.close()

    def test_all_envs_gymnasium_compatible(self):
        """All envs should work with standard Gymnasium API."""
        env_ids = [
            "oncosim/BeamSelection-v0",
            "oncosim/DoseFractionation-v0",
            "oncosim/AdaptiveRT-v0",
        ]
        for env_id in env_ids:
            env = gym.make(env_id)
            obs, info = env.reset(seed=0)
            assert isinstance(info, dict)
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            assert isinstance(info, dict)
            env.close()

    def test_physics_imports(self):
        from oncosim.physics import (
            calculate_beam_dose,
            calculate_plan_dose,
            surviving_fraction,
            tcp,
            ntcp_lkb,
            bed,
            generate_beam_profile,
            create_tumor_mask,
            compute_dvh,
        )
        assert callable(calculate_beam_dose)
        assert callable(tcp)
        assert callable(compute_dvh)
