"""Tests for Gymnasium environment registration."""
import gymnasium as gym
import pytest

import peptidegym  # noqa: F401


def test_amp_v0_registered():
    """PeptideGym/AMP-v0 should be registered."""
    spec = gym.spec("PeptideGym/AMP-v0")
    assert spec is not None
    assert spec.id == "PeptideGym/AMP-v0"


def test_cyclic_v0_registered():
    """PeptideGym/CyclicPeptide-v0 should be registered."""
    spec = gym.spec("PeptideGym/CyclicPeptide-v0")
    assert spec is not None
    assert spec.id == "PeptideGym/CyclicPeptide-v0"


def test_epitope_v0_registered():
    """PeptideGym/Epitope-v0 should be registered."""
    spec = gym.spec("PeptideGym/Epitope-v0")
    assert spec is not None
    assert spec.id == "PeptideGym/Epitope-v0"


def test_all_nine_configs_registered():
    """All 9 benchmark environment configs should be registered."""
    expected_ids = [
        "PeptideGym/AMP-Easy-v0",
        "PeptideGym/AMP-v0",
        "PeptideGym/AMP-Hard-v0",
        "PeptideGym/CyclicPeptide-Easy-v0",
        "PeptideGym/CyclicPeptide-v0",
        "PeptideGym/CyclicPeptide-Hard-v0",
        "PeptideGym/Epitope-Easy-v0",
        "PeptideGym/Epitope-v0",
        "PeptideGym/Epitope-Hard-v0",
    ]
    for env_id in expected_ids:
        spec = gym.spec(env_id)
        assert spec is not None, f"{env_id} not registered"


def test_make_and_reset_all():
    """Every registered environment should be make-able and reset-able."""
    env_ids = [
        "PeptideGym/AMP-Easy-v0",
        "PeptideGym/AMP-v0",
        "PeptideGym/AMP-Hard-v0",
        "PeptideGym/CyclicPeptide-Easy-v0",
        "PeptideGym/CyclicPeptide-v0",
        "PeptideGym/CyclicPeptide-Hard-v0",
        "PeptideGym/Epitope-Easy-v0",
        "PeptideGym/Epitope-v0",
        "PeptideGym/Epitope-Hard-v0",
    ]
    for env_id in env_ids:
        env = gym.make(env_id)
        obs, info = env.reset(seed=42)
        assert isinstance(obs, dict), f"{env_id} reset did not return dict obs"
        assert isinstance(info, dict), f"{env_id} reset did not return dict info"
        env.close()
