import pytest
from environment.env import VoiceAuthenticityEnv, TASKS

def test_reset_returns_observation():
    env = VoiceAuthenticityEnv("clean_detection")
    obs = env.reset()
    assert obs is not None
    assert obs.step_number == 0
    assert obs.task_name == "clean_detection"
    assert "hint" in obs.dict()

def test_step_returns_reward_in_range():
    env = VoiceAuthenticityEnv("clean_detection")
    env.reset()
    obs, reward, done, info = env.step({"action_type": "request_temporal_features"})
    assert 0.05 <= reward <= 0.95
    assert not done

def test_five_actions_complete_episode():
    env = VoiceAuthenticityEnv("clean_detection")
    env.reset()
    actions = [
        "request_temporal_features",
        "request_spectral_features",
        "request_comparison",
        "analyze_evidence",
        "final_classify"
    ]
    
    for i, act in enumerate(actions):
        obs, reward, done, info = env.step({
            "action_type": act,
            "label": 0,
            "confidence": 0.8,
            "reasoning": "Test reasoning"
        })
        if i < len(actions) - 1:
            assert not done
        else:
            assert done

def test_reward_never_zero_or_one():
    env = VoiceAuthenticityEnv("clean_detection")
    env.reset()
    # Test an action that could get penalties or rewards
    obs, reward, done, info = env.step({"action_type": "request_temporal_features"})
    assert reward != 0.0
    assert reward != 1.0

def test_all_tasks_load():
    for task in TASKS:
        env = VoiceAuthenticityEnv(task)
        assert env.task_name == task
        obs = env.reset()
        assert obs.task_name == task


def test_realtime_classify_after_step_2():
    """Realtime task: agent can classify after 2 steps with time penalty."""
    env = VoiceAuthenticityEnv("realtime_detection")
    env.reset(seed=42)

    # Step 1: gather temporal
    obs, r1, done, info = env.step({"action_type": "request_temporal_features"})
    assert not done
    # final_classify should NOT be available yet (only 1 step taken)
    assert "final_classify" not in obs.available_actions

    # Step 2: gather spectral
    obs, r2, done, info = env.step({"action_type": "request_spectral_features"})
    assert not done
    # final_classify SHOULD be available now (2 steps taken)
    assert "final_classify" in obs.available_actions

    # Step 3: classify immediately (1 extra step beyond step 2 = -0.03 penalty)
    obs, r3, done, info = env.step({
        "action_type": "final_classify",
        "label": 0,
        "confidence": 0.75,
        "reasoning": "Natural jitter and shimmer suggest real human speech"
    })
    assert done
    assert 0.05 <= r3 <= 0.95
    # Should have realtime penalty info
    assert "realtime_time_penalty" in info
    assert info["realtime_extra_steps"] == 1  # step 3 is 1 extra beyond step 2


def test_realtime_no_penalty_at_step_2():
    """Classifying exactly at step 2 should have 0 extra steps penalty."""
    env = VoiceAuthenticityEnv("realtime_detection")
    env.reset(seed=42)

    # Step 1: gather temporal
    env.step({"action_type": "request_temporal_features"})

    # Step 2: gather spectral
    env.step({"action_type": "request_spectral_features"})

    # The penalty math: step_number=2, extra = 2 - 2 = 0, penalty = 0
    # But we need step 3 for classify, so minimum penalty is 0.03
    # Actually step_number increments on step(), so at classify it becomes 3
    # extra = 3 - 2 = 1, penalty = 0.03
    # This is by design: the minimum cost for classifying is 1 extra step