import pytest

from hippotorch.curriculum.callbacks import HippotorchCurriculumCallback


class _MockEnv:
    def __init__(self, length: int = 5, max_length: int = 30):
        self.length = int(length)
        self._max_length = int(max_length)
        self.calls = []

    def set_corridor_length(self, L: int) -> None:
        value = min(int(L), self._max_length)
        self.calls.append(value)
        self.length = value


class _NoSetterEnv:
    def __init__(self, length: int = 5):
        self.length = length


class _VectorEnv:
    def __init__(self, envs):
        self.envs = envs
        self.length = envs[0].length


def test_curriculum_threshold_progression_and_window_behavior():
    env = _MockEnv(length=5, max_length=15)
    cb = HippotorchCurriculumCallback(
        min_length=5, max_length=15, step=5, window=4, threshold=0.75
    )

    # Window not yet full → no length change despite rewards
    for _ in range(3):
        cb.on_episode_end(env, reward=1.0)
    assert env.length == 5

    # Fourth success fills the window and steps to 10
    cb.on_episode_end(env, reward=1.0)
    assert env.length == 10
    assert env.calls[0] == 5  # initialization call
    assert env.calls[-1] == 10

    # Below-threshold window leaves length unchanged
    for reward in [1.0, 0.0, 1.0, 0.0]:
        cb.on_episode_end(env, reward=reward)
    assert env.length == 10

    # Subsequent successes reach the max_length and stay there
    for _ in range(4):
        cb.on_episode_end(env, reward=1.0)
    assert env.length == 15
    for _ in range(4):
        cb.on_episode_end(env, reward=1.0)
    assert env.length == 15


def test_curriculum_warns_when_env_lacks_setter():
    env = _NoSetterEnv(length=4)
    cb = HippotorchCurriculumCallback(min_length=4, max_length=6, window=2, step=1)

    with pytest.warns(
        UserWarning, match="NoSetterEnv.*set_corridor_length\\(\\)"
    ) as captured:
        cb.on_episode_end(env, reward=1.0)
        cb.on_episode_end(env, reward=1.0)

    assert len(captured) == 1
    assert env.length == 4  # unchanged without setter


def test_curriculum_updates_vectorized_envs():
    envs = [_MockEnv(length=5, max_length=12), _MockEnv(length=5, max_length=12)]
    vec_env = _VectorEnv(envs)
    cb = HippotorchCurriculumCallback(
        min_length=5, max_length=12, step=2, window=3, threshold=0.8
    )

    for _ in range(3):
        cb.on_episode_end(vec_env, reward=1.0)

    assert envs[0].length == 7
    assert envs[1].length == 7
    assert all(7 in e.calls for e in envs)
