"""Tests for the code generator."""

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


class TestCodeGenerator:
    def test_generate_sft_basic(self):
        config = TrainingConfig(
            model_name="meta-llama/Llama-3-8B",
            quantization="4bit",
            batch_size=4,
            num_epochs=3,
            max_seq_length=512,
            lora_enabled=True,
            lora_rank=16,
            lora_alpha=32,
            lora_target_modules=["q_proj", "v_proj"],
            bf16=True,
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert "meta-llama/Llama-3-8B" in script
        assert "load_in_4bit=True" in script
        assert "r=16" in script
        assert "lora_alpha=32" in script
        assert "num_train_epochs=3" in script
        assert "per_device_train_batch_size=4" in script
        assert "bf16=True" in script

    def test_generate_sft_no_quantization(self):
        config = TrainingConfig(
            quantization="none",
            bf16=True,
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert "load_in_4bit" not in script
        assert "load_in_8bit" not in script
        assert "bnb_config = None" in script

    def test_generate_sft_8bit(self):
        config = TrainingConfig(quantization="8bit")
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert "load_in_8bit=True" in script

    def test_generate_sft_alpaca_format(self):
        config = TrainingConfig(quantization="4bit")
        gen = CodeGenerator()
        script = gen.generate_sft(config, dataset_format="alpaca")

        assert "format_alpaca" in script
        assert "### Instruction:" in script

    def test_generate_sft_sharegpt_format(self):
        config = TrainingConfig(quantization="4bit")
        gen = CodeGenerator()
        script = gen.generate_sft(config, dataset_format="sharegpt")

        assert "format_sharegpt" in script
        assert "conversations" in script

    def test_generate_sft_gradient_checkpointing(self):
        config = TrainingConfig(
            quantization="4bit",
            use_gradient_checkpointing=True,
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert "gradient_checkpointing_enable" in script
        assert "gradient_checkpointing=True" in script

    def test_generate_sft_ddp(self):
        config = TrainingConfig(
            quantization="4bit",
            num_gpus=2,
            distributed_strategy="ddp",
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert "ddp_find_unused_parameters=False" in script

    def test_generate_sft_fsdp(self):
        config = TrainingConfig(
            quantization="none",
            num_gpus=8,
            distributed_strategy="fsdp",
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert 'fsdp="full_shard auto_wrap"' in script

    def test_generate_and_save(self, tmp_path):
        config = TrainingConfig(quantization="4bit")
        gen = CodeGenerator()
        output = tmp_path / "train.py"
        result = gen.generate_and_save(config, output)

        assert result == output
        assert output.exists()
        content = output.read_text()
        assert "SFTTrainer" in content

    def test_generate_sft_deepspeed(self):
        config = TrainingConfig(
            quantization="4bit",
            num_gpus=4,
            distributed_strategy="deepspeed",
            deepspeed_stage=2,
            deepspeed_offload_optimizer=True,
            bf16=True,
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert "deepspeed" in script
        assert "zero_optimization" in script

    def test_generate_sft_wandb_tracking(self):
        config = TrainingConfig(
            quantization="4bit",
            report_to="wandb",
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert 'report_to="wandb"' in script

    def test_generate_sft_no_tracking(self):
        config = TrainingConfig(
            quantization="4bit",
            report_to="none",
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config)

        assert 'report_to="none"' in script

    def test_generate_dpo_basic(self):
        config = TrainingConfig(
            model_name="meta-llama/Llama-3-8B",
            quantization="4bit",
            batch_size=2,
            num_epochs=1,
            lora_enabled=True,
            lora_rank=16,
            lora_alpha=32,
            lora_target_modules=["q_proj", "v_proj"],
            bf16=True,
        )
        gen = CodeGenerator()
        script = gen.generate_dpo(config)

        assert "DPOTrainer" in script
        assert "ref_model" in script
        assert "meta-llama/Llama-3-8B" in script
        assert "load_in_4bit=True" in script

    def test_generate_dpo_valid_python(self):
        config = TrainingConfig(
            quantization="4bit",
            lora_enabled=True,
            lora_rank=16,
            lora_alpha=32,
            lora_target_modules=["q_proj", "v_proj"],
            bf16=True,
        )
        gen = CodeGenerator()
        script = gen.generate_dpo(config)
        compile(script, "<generated_dpo>", "exec")

    def test_script_is_valid_python(self):
        config = TrainingConfig(
            quantization="4bit",
            lora_enabled=True,
            lora_rank=16,
            lora_alpha=32,
            lora_target_modules=["q_proj", "v_proj"],
            bf16=True,
        )
        gen = CodeGenerator()
        script = gen.generate_sft(config, dataset_format="alpaca")

        # Should be valid Python syntax
        compile(script, "<generated>", "exec")
