"""Tests for the guardrail engine and optimizer models."""

from bilgeai.optimizer.engine import GuardrailEngine
from bilgeai.optimizer.models import Explanation, TrainingConfig


class TestTrainingConfig:
    def test_default_creation(self):
        config = TrainingConfig()
        assert config.batch_size == 4
        assert config.learning_rate == 2e-4

    def test_add_explanation(self):
        config = TrainingConfig()
        config.add_explanation(
            decision="batch_size = 4",
            reason="Test reason",
            confidence=0.9,
        )
        assert len(config.explanations) == 1
        assert config.explanations[0].source == "rule_engine"

    def test_invalid_batch_size(self):
        import pytest
        from pydantic import ValidationError

        with pytest.raises(ValidationError):
            TrainingConfig(batch_size=0)

    def test_invalid_learning_rate(self):
        import pytest
        from pydantic import ValidationError

        with pytest.raises(ValidationError):
            TrainingConfig(learning_rate=-1)

    def test_default_seed(self):
        config = TrainingConfig()
        assert config.seed == 42

    def test_custom_seed(self):
        config = TrainingConfig(seed=123)
        assert config.seed == 123

    def test_seed_in_dump(self):
        config = TrainingConfig(seed=99)
        data = config.model_dump()
        assert data["seed"] == 99


class TestExplanation:
    def test_confidence_bounds(self):
        import pytest
        from pydantic import ValidationError

        with pytest.raises(ValidationError):
            Explanation(
                decision="test",
                reason="test",
                confidence=1.5,
                source="rule_engine",
            )


class TestGuardrailEngine:
    def test_low_vram_forces_4bit(self, hw_low_gpu, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_low_gpu, small_dataset_profile)
        assert constraints.required_quantization == "4bit"

    def test_high_vram_no_quant(self, hw_high_gpu, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_high_gpu, small_dataset_profile)
        assert constraints.required_quantization == "none"

    def test_no_gpu_4bit(self, hw_no_gpu, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_no_gpu, small_dataset_profile)
        assert constraints.required_quantization == "4bit"
        assert constraints.max_batch_size == 1

    def test_batch_size_always_positive(self, hw_mid_gpu, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_mid_gpu, small_dataset_profile)
        assert constraints.max_batch_size > 0

    def test_every_decision_has_explanation(self, hw_mid_gpu, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_mid_gpu, small_dataset_profile)
        assert len(constraints.explanations) > 0
        for exp in constraints.explanations:
            assert exp.decision
            assert exp.reason
            assert exp.source == "rule_engine"


class TestMultiGPU:
    def test_multi_gpu_distributed(self, hw_multi_gpu, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_multi_gpu, small_dataset_profile)
        assert constraints.num_gpus == 2
        assert constraints.distributed_strategy in ("ddp", "deepspeed", "fsdp")

    def test_8x_a100_fsdp_strategy(self, hw_8x_a100, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_8x_a100, small_dataset_profile)
        assert constraints.num_gpus == 8
        assert constraints.distributed_strategy == "fsdp"
        assert constraints.required_quantization == "none"

    def test_single_gpu_no_distributed(self, hw_mid_gpu, small_dataset_profile):
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_mid_gpu, small_dataset_profile)
        assert constraints.num_gpus == 1
        assert constraints.distributed_strategy == "none"

    def test_hardware_profile_computed_fields(self, hw_8x_a100):
        assert hw_8x_a100.num_gpus == 8
        assert hw_8x_a100.total_vram_gb == 640.0
        assert hw_8x_a100.min_vram_gb == 80.0

    def test_no_gpu_computed_fields(self, hw_no_gpu):
        assert hw_no_gpu.num_gpus == 0
        assert hw_no_gpu.total_vram_gb == 0.0
        assert hw_no_gpu.min_vram_gb == 0.0


class TestDistributedStrategy:
    def test_low_vram_multi_gpu_uses_ddp(self, hw_4x_low_gpu, small_dataset_profile):
        """4x 4GB GPUs with LoRA/quantization should use DDP."""
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_4x_low_gpu, small_dataset_profile)
        assert constraints.distributed_strategy == "ddp"

    def test_high_vram_uses_fsdp(self, hw_8x_a100, small_dataset_profile):
        """8x A100 80GB should use FSDP."""
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_8x_a100, small_dataset_profile)
        assert constraints.distributed_strategy == "fsdp"

    def test_deepspeed_config_fields(self):
        """TrainingConfig should have DeepSpeed fields."""
        config = TrainingConfig(
            deepspeed_stage=3,
            deepspeed_offload_optimizer=True,
            deepspeed_offload_param=True,
        )
        assert config.deepspeed_stage == 3
        assert config.deepspeed_offload_optimizer is True
        assert config.deepspeed_offload_param is True

    def test_single_gpu_no_distributed(self, hw_low_gpu, small_dataset_profile):
        """Single GPU should not use distributed training."""
        engine = GuardrailEngine()
        constraints = engine.compute_constraints(hw_low_gpu, small_dataset_profile)
        assert constraints.distributed_strategy == "none"


class TestExperimentTracking:
    def test_report_to_default_none(self):
        config = TrainingConfig()
        assert config.report_to is None

    def test_report_to_wandb(self):
        config = TrainingConfig(report_to="wandb", wandb_project="my-project")
        assert config.report_to == "wandb"
        assert config.wandb_project == "my-project"
