"""Tests for RLHF template generation."""

from bilgeai.generator.codegen import CodeGenerator
from bilgeai.optimizer.models import TrainingConfig


class TestRLHFTemplate:
    def test_generate_rlhf_basic(self):
        gen = CodeGenerator()
        config = TrainingConfig(model_name="meta-llama/Llama-3-8B")
        script = gen.generate_rlhf(config)

        assert "PPO" in script or "ppo" in script.lower()
        assert "meta-llama/Llama-3-8B" in script
        assert "SEED = 42" in script

    def test_generate_rlhf_custom_reward_model(self):
        gen = CodeGenerator()
        config = TrainingConfig()
        script = gen.generate_rlhf(
            config,
            reward_model_name="OpenAssistant/reward-model-deberta-v3-large",
        )
        assert "OpenAssistant/reward-model-deberta-v3-large" in script

    def test_generate_rlhf_with_quantization(self):
        gen = CodeGenerator()
        config = TrainingConfig(quantization="4bit", bf16=True)
        script = gen.generate_rlhf(config)
        assert "load_in_4bit" in script
        assert "bfloat16" in script

    def test_generate_rlhf_with_lora(self):
        gen = CodeGenerator()
        config = TrainingConfig(lora_enabled=True, lora_rank=32)
        script = gen.generate_rlhf(config)
        assert "LoraConfig" in script
        assert "r=32" in script

    def test_generate_rlhf_seed(self):
        gen = CodeGenerator()
        config = TrainingConfig(seed=123)
        script = gen.generate_rlhf(config)
        assert "SEED = 123" in script

    def test_generate_rlhf_mps_support(self):
        gen = CodeGenerator()
        config = TrainingConfig()
        script = gen.generate_rlhf(config)
        assert "mps" in script
