"""Integration tests for the GRPO training pipeline flow."""

from __future__ import annotations

from sql_env.models import SQLObservation
from sql_env.training import rollout as rollout_module
from sql_env.training.config import GRPOConfig
from sql_env.training.rewards import (
    reward_correctness,
    reward_operational,
    reward_progress,
)
from sql_env.training.rollout import rollout_func


class _Tokenizer:
    def apply_chat_template(
        self,
        messages: list[dict[str, str]],
        tokenize: bool = False,
        add_generation_prompt: bool = True,
    ) -> str:
        del messages
        del tokenize
        del add_generation_prompt
        return "prompt"


class _Model:
    def __init__(self) -> None:
        self.calls = 0

    def generate(self, prompt: str, max_new_tokens: int) -> str:
        del prompt
        del max_new_tokens
        self.calls += 1
        if self.calls == 1:
            return "hello world random text"
        return "ANSWER: 42"


class _Environment:
    def __init__(self, step_budget: int) -> None:
        self.step_budget = step_budget
        self.step_count = 0
        self.state = type("State", (), {"episode_id": "ep-integration"})()

    def reset(self, *, seed: int | None = None) -> SQLObservation:
        del seed
        self.step_count = 0
        return self._observation(done=False, result="")

    def step(self, action) -> SQLObservation:
        self.step_count += 1
        if (
            action.action_type == "QUERY"
            and action.argument == "hello world random text"
        ):
            return self._observation(done=False, result="", error="unparseable action")
        if action.action_type == "ANSWER":
            return self._observation(
                done=True, result="Answer submitted: correct.", reward=1.0
            )
        return self._observation(done=False, result="ok", reward=0.1)

    def _observation(
        self,
        *,
        done: bool,
        result: str,
        error: str = "",
        reward: float | None = 0.0,
    ) -> SQLObservation:
        return SQLObservation(
            question="How many rows?",
            schema_info="Available tables:\n- t",
            result=result,
            error=error,
            step_count=self.step_count,
            budget_remaining=max(0, self.step_budget - self.step_count),
            action_history=[],
            done=done,
            reward=reward,
        )


def test_training_pipeline_flow_with_reward_functions(monkeypatch) -> None:
    """Rollout output can be consumed by all reward callables."""

    config = GRPOConfig(
        questions_path="data/questions/questions_train.json",
        db_dir="data/databases",
        output_dir="outputs/grpo_test",
        step_budget=3,
    )
    tokenizer = _Tokenizer()
    model = _Model()
    fake_env = _Environment(step_budget=3)

    monkeypatch.setattr(rollout_module, "_build_environment", lambda *_: fake_env)

    rollouts = rollout_func(["Count rows"], model, tokenizer, config)
    assert len(rollouts) == 1

    metadata = [item["metadata"] for item in rollouts]
    completions = [
        [{"role": "assistant", "content": item["content"]}] for item in rollouts
    ]

    assert reward_correctness(completions, metadata=metadata) == [1.0]
    progress = reward_progress(completions, metadata=metadata)
    operational = reward_operational(completions, metadata=metadata)
    assert len(progress) == 1
    assert 0.0 <= progress[0] <= 1.0
    assert len(operational) == 1


def test_unparseable_action_recovers_and_episode_continues(monkeypatch) -> None:
    """Unparseable model output falls back to QUERY and does not abort episode."""

    config = GRPOConfig(
        questions_path="data/questions/questions_train.json",
        db_dir="data/databases",
        output_dir="outputs/grpo_test",
        step_budget=3,
    )
    tokenizer = _Tokenizer()
    model = _Model()
    fake_env = _Environment(step_budget=3)

    monkeypatch.setattr(rollout_module, "_build_environment", lambda *_: fake_env)

    rollout = rollout_func(["Count rows"], model, tokenizer, config)[0]
    assert rollout["metadata"]["step_count"] >= 2
    assert rollout["metadata"]["done"] is True