"""
Tests for auto-resource detection.

GateWay should automatically detect optimal resources based on code analysis,
ensuring users don't overpay for resources they don't need.
"""

from verlex_server.gateway.core.profiling import (
    ResourceRecommendation,
    WorkloadType,
    auto_detect_resources,
)

# ============================================================================
# Test Functions (these simulate real user code)
# ============================================================================

def simple_cpu_function():
    """Should detect: CPU-only, minimal resources."""
    x = 0
    for i in range(100):
        x += i
    return x


def pytorch_inference():
    """Should detect: GPU (T4), inference workload."""
    import torch
    model = torch.nn.Linear(100, 10)
    x = torch.randn(32, 100)
    return model(x)


def pytorch_training():
    """Should detect: GPU (T4/A10), training workload with backward()."""
    import torch
    model = torch.nn.Linear(100, 10)
    optimizer = torch.optim.Adam(model.parameters())
    x = torch.randn(32, 100)
    y = model(x)
    loss = y.sum()
    loss.backward()
    optimizer.step()
    return loss.item()


def heavy_training():
    """Should detect: Larger GPU for heavy training."""
    import torch
    model = torch.nn.Linear(1000, 100)
    optimizer = torch.optim.Adam(model.parameters())

    for _epoch in range(100):
        x = torch.randn(64, 1000)
        y = model(x)
        loss = y.sum()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    return loss.item()


def data_processing():
    """Should detect: CPU-heavy, more memory, pandas workload."""
    import numpy as np
    import pandas as pd

    df = pd.DataFrame(np.random.randn(100000, 50))
    result = df.groupby(df[0] > 0).mean()
    return result


def transformers_model():
    """Should detect: Large GPU for HuggingFace transformers."""
    from transformers import AutoModelForCausalLM
    model = AutoModelForCausalLM.from_pretrained("gpt2")
    return model


# ============================================================================
# Tests
# ============================================================================

class TestAutoResourceDetection:
    """Test automatic resource detection."""

    def test_simple_cpu_function(self):
        """Simple functions should get minimal resources."""
        rec = auto_detect_resources(simple_cpu_function)

        assert isinstance(rec, ResourceRecommendation)
        assert rec.gpu_type is None, "Simple CPU function shouldn't need GPU"
        assert rec.cpu <= 4, "Should not over-provision CPU"
        assert rec.memory_gb <= 8, "Should not over-provision memory"
        assert rec.workload_type == WorkloadType.GENERAL

    def test_pytorch_inference_detection(self):
        """PyTorch inference should get small GPU."""
        rec = auto_detect_resources(pytorch_inference)

        assert rec.gpu_type is not None, "PyTorch should get GPU"
        assert rec.gpu_type in ("T4", "L4"), f"Inference should use cost-effective GPU, got {rec.gpu_type}"
        assert rec.gpu_count == 1
        assert "torch" in rec.frameworks or "pytorch" in str(rec.frameworks).lower()

    def test_pytorch_training_detection(self):
        """PyTorch training should detect training workload."""
        rec = auto_detect_resources(pytorch_training)

        assert rec.gpu_type is not None, "Training should get GPU"
        assert rec.workload_type == WorkloadType.ML_TRAINING
        assert rec.memory_gb >= 8, "Training needs adequate memory"
        assert "backward" in str(rec.reasoning).lower() or "training" in str(rec.reasoning).lower()

    def test_heavy_training_gets_more_resources(self):
        """Heavy training should get beefier GPU."""
        rec = auto_detect_resources(heavy_training)

        assert rec.gpu_type is not None
        # More training indicators should result in better GPU
        assert rec.ml_training_indicators if hasattr(rec, 'ml_training_indicators') else True
        assert len(rec.reasoning) > 0

    def test_recommendation_has_required_fields(self):
        """Recommendations should have all required fields."""
        rec = auto_detect_resources(simple_cpu_function)

        # Check required fields
        assert hasattr(rec, 'cpu')
        assert hasattr(rec, 'memory_gb')
        assert hasattr(rec, 'gpu_type')
        assert hasattr(rec, 'gpu_count')
        assert hasattr(rec, 'reasoning')
        assert hasattr(rec, 'workload_type')
        assert hasattr(rec, 'frameworks')
        assert hasattr(rec, 'confidence')

        # Check types
        assert isinstance(rec.cpu, int)
        assert isinstance(rec.memory_gb, (int, float))
        assert isinstance(rec.reasoning, list)
        assert isinstance(rec.confidence, float)
        assert 0 <= rec.confidence <= 1

    def test_to_resource_dict(self):
        """to_resource_dict should return valid @resource config."""
        rec = auto_detect_resources(pytorch_training)
        config = rec.to_resource_dict()

        assert "cpu" in config
        assert "memory" in config
        assert isinstance(config["memory"], str)
        assert "GB" in config["memory"]

        if rec.gpu_type:
            assert "gpu" in config

    def test_recommendations_include_reasoning(self):
        """Every recommendation should explain why."""
        rec = auto_detect_resources(pytorch_training)

        assert len(rec.reasoning) > 0, "Should provide reasoning"
        # Each reason should be a non-empty string
        for reason in rec.reasoning:
            assert isinstance(reason, str)
            assert len(reason) > 0

    def test_savings_calculation(self):
        """Should calculate potential savings vs over-provisioning."""
        rec = auto_detect_resources(simple_cpu_function)

        # Simple CPU function should have savings vs GPU over-provisioning
        # This tests the "right-sizing" feature
        assert hasattr(rec, 'savings_vs_default')


class TestCostEfficiency:
    """Test that auto-detection prefers cost-efficient options."""

    def test_small_models_get_small_gpus(self):
        """Small models shouldn't get A100s."""
        rec = auto_detect_resources(pytorch_inference)

        if rec.gpu_type:
            expensive_gpus = {"A100", "A100-80", "H100"}
            assert rec.gpu_type not in expensive_gpus, \
                f"Small model shouldn't get {rec.gpu_type}"

    def test_cpu_workload_no_gpu(self):
        """CPU-only workloads shouldn't get GPU."""
        rec = auto_detect_resources(simple_cpu_function)

        assert rec.gpu_type is None, \
            f"CPU workload shouldn't get GPU, got {rec.gpu_type}"
        assert rec.gpu_count == 0


class TestWorkloadTypeDetection:
    """Test workload type classification."""

    def test_ml_training_detected(self):
        """ML training should be detected."""
        rec = auto_detect_resources(pytorch_training)
        assert rec.workload_type == WorkloadType.ML_TRAINING

    def test_general_workload_detected(self):
        """Simple functions should be GENERAL."""
        rec = auto_detect_resources(simple_cpu_function)
        assert rec.workload_type == WorkloadType.GENERAL


# Run with: pytest tests/test_auto_resources.py -v
