"""
Tests for the SQL Debug Environment lifecycle.
Run with: pytest tests/test_environment.py -v
"""
import sys
sys.path.insert(0, ".")

import pytest
from models import Action
from server.sql_environment import SQLDebugEnvironment
from tasks import TASKS, TASK_GROUPS, get_task


@pytest.fixture
def env():
    return SQLDebugEnvironment()


def test_reset_returns_observation(env):
    obs = env.reset()
    assert obs.schema_sql
    assert obs.buggy_query
    assert obs.task_description
    assert obs.attempts_remaining == 5
    assert not obs.done
    assert not obs.success


def test_reset_all_task_groups(env):
    for group in TASK_GROUPS:
        obs = env.reset(task_group=group)
        assert obs.schema_sql
        assert obs.buggy_query


def test_reset_specific_task(env):
    obs = env.reset(task_id="easy_missing_from")
    assert obs.attempts_remaining == 5
    state = env.state()
    assert state.task_id == "easy_missing_from"


def test_step_decrements_attempts(env):
    env.reset(task_id="easy_missing_from")
    obs = env.step(Action(query="SELECT 1"))
    assert obs.attempts_remaining == 4


def test_correct_answer_succeeds(env):
    env.reset(task_id="easy_missing_from")
    task = get_task("easy_missing_from")
    obs = env.step(Action(query=task["correct_query"]))
    assert obs.success
    assert obs.done
    assert obs.reward >= 1.0


def test_max_steps_terminates(env):
    env.reset(task_id="easy_missing_from")
    for _ in range(5):
        obs = env.step(Action(query="SELECT 1"))
        if obs.done:
            break
    assert obs.done


def test_step_after_done_is_safe(env):
    env.reset(task_id="easy_missing_from")
    task = get_task("easy_missing_from")
    env.step(Action(query=task["correct_query"]))
    obs = env.step(Action(query="SELECT 1"))  # after done
    assert obs.done  # should still be done


def test_state_step_count(env):
    env.reset()
    assert env.state().step_count == 0
    env.step(Action(query="SELECT 1"))
    assert env.state().step_count == 1
    env.step(Action(query="SELECT 1"))
    assert env.state().step_count == 2


def test_seed_reproducibility(env):
    obs1 = env.reset(task_group="easy_syntax_fix", seed=42)
    obs2 = env.reset(task_group="easy_syntax_fix", seed=42)
    assert obs1.buggy_query == obs2.buggy_query


def test_partial_reward_on_progress(env):
    env.reset(task_id="easy_missing_from")
    obs = env.step(Action(
        query="SELECT first_name, last_name FROM employees ORDER BY last_name"
    ))
    assert obs.reward >= 0.0
    assert not obs.success


def test_cumulative_reward_accumulates(env):
    env.reset(task_id="easy_missing_from")
    env.step(Action(query="SELECT 1"))
    env.step(Action(query="SELECT 2"))
    state = env.state()
    assert state.cumulative_reward is not None