"""
Tests for Smart Executor and Serverless Execution

Tests the fixes for:
1. VM-per-job scaling problem (serverless mode for short jobs)
2. Serialization limitations (source mode alternative to pickle)
"""


import pytest

from verlex_server.gateway.backend.serverless_executor import (
    ExecutionTier,
    LocalDockerExecutor,
    ServerlessConfig,
    select_execution_tier,
    should_use_serverless,
)
from verlex_server.gateway.backend.smart_executor import (
    ExecutionBackend,
    ExecutionPlan,
    ResourceRequirements,
    SerializationMethod,
    SmartExecutor,
)
from verlex_server.gateway.core.source_serialization import (
    HybridSerializer,
    SourceBundle,
    can_extract_source,
    extract_source,
    smart_serialize,
)

# =============================================================================
# Tests for Execution Tier Selection
# =============================================================================

class TestExecutionTierSelection:
    """Test automatic selection between serverless and VM execution."""

    def test_short_cpu_job_uses_serverless(self):
        """Short CPU-only jobs should use serverless."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=60,  # 1 minute
            gpu_required=False,
            gpu_type=None,
            memory_gb=4.0,
        )
        assert tier == ExecutionTier.SERVERLESS
        assert effective_gpu is None

    def test_long_job_uses_vm(self):
        """Jobs over 10 minutes should use VM."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=1200,  # 20 minutes
            gpu_required=False,
            gpu_type=None,
            memory_gb=4.0,
        )
        assert tier == ExecutionTier.VM
        assert effective_gpu is None

    def test_a100_gpu_prefers_serverless(self):
        """A100 GPU jobs prefer serverless (Cloud Run supports A100)."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=60,  # Short job
            gpu_required=True,
            gpu_type="A100",
            memory_gb=16.0,
        )
        assert tier == ExecutionTier.SERVERLESS
        assert effective_gpu == "A100"

    def test_h100_gpu_prefers_serverless(self):
        """H100 GPU jobs prefer serverless (Cloud Run supports H100)."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=60,
            gpu_required=True,
            gpu_type="H100",
            memory_gb=32.0,
        )
        assert tier == ExecutionTier.SERVERLESS
        assert effective_gpu == "H100"

    def test_t4_gpu_substituted_for_serverless(self):
        """T4 GPU is substituted for serverless (provider-specific)."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=300,  # 5 minutes
            gpu_required=True,
            gpu_type="T4",
            memory_gb=16.0,
        )
        assert tier == ExecutionTier.SERVERLESS
        # T4 → L4 (GCP Cloud Run) or T4 → V100 (Azure ACI)
        assert effective_gpu in ("L4", "V100")

    def test_a10g_gpu_substituted_for_serverless(self):
        """A10G GPU is substituted for serverless (provider-specific)."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=300,
            gpu_required=True,
            gpu_type="A10G",
            memory_gb=16.0,
        )
        assert tier == ExecutionTier.SERVERLESS
        # A10G → L4 (GCP Cloud Run) or A10G → V100 (Azure ACI)
        assert effective_gpu in ("L4", "V100")

    def test_l4_gpu_uses_serverless(self):
        """L4 GPU jobs use serverless (native on GCP or substituted on Azure)."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=300,
            gpu_required=True,
            gpu_type="L4",
            memory_gb=16.0,
        )
        assert tier == ExecutionTier.SERVERLESS
        # L4 native on GCP Cloud Run, or L4 → V100 on Azure ACI
        assert effective_gpu in ("L4", "V100")

    def test_high_memory_uses_vm(self):
        """Jobs exceeding ALL serverless provider memory limits should use VM."""
        # The threshold is now dynamic — based on the max memory across all
        # serverless providers (e.g., Fargate supports up to 120GB).
        # Use a value that exceeds every provider's limit.
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=60,
            gpu_required=False,
            gpu_type=None,
            memory_gb=256.0,
        )
        assert tier == ExecutionTier.VM
        assert effective_gpu is None

    def test_should_use_serverless_helper(self):
        """Test the convenience helper function."""
        assert should_use_serverless(60, None, 4.0) is True
        assert should_use_serverless(1200, None, 4.0) is False
        assert should_use_serverless(60, "A100", 16.0) is True  # Cloud Run supports A100
        assert should_use_serverless(300, "L4", 16.0) is True


class TestServerlessConfig:
    """Test serverless configuration validation."""

    def test_valid_gcp_config(self):
        """Valid Cloud Run config should have no issues."""
        config = ServerlessConfig(
            cpu=4.0,
            memory_mb=8192,
            timeout_seconds=600,
        )
        from verlex_server.gateway.backend.serverless_executor import ServerlessProvider
        issues = config.validate_for_provider(ServerlessProvider.GCP_CLOUD_RUN)
        assert len(issues) == 0

    def test_cpu_exceeds_gcp_limit(self):
        """CPU over 8 should fail for Cloud Run."""
        config = ServerlessConfig(cpu=16.0)
        from verlex_server.gateway.backend.serverless_executor import ServerlessProvider
        issues = config.validate_for_provider(ServerlessProvider.GCP_CLOUD_RUN)
        assert any("8 vCPU" in issue for issue in issues)

    def test_memory_exceeds_gcp_limit(self):
        """Memory over 32GB should fail for Cloud Run."""
        config = ServerlessConfig(memory_mb=64 * 1024)  # 64GB
        from verlex_server.gateway.backend.serverless_executor import ServerlessProvider
        issues = config.validate_for_provider(ServerlessProvider.GCP_CLOUD_RUN)
        assert any("32GB" in issue for issue in issues)


# =============================================================================
# Tests for Source Code Serialization
# =============================================================================

class TestSourceExtraction:
    """Test source code extraction (Modal-style serialization)."""

    def test_extract_simple_function(self):
        """Should extract source from a simple function."""
        def simple_function(x, y):
            return x + y

        bundle = extract_source(simple_function, 1, 2)

        assert bundle.entry_point == "simple_function"
        assert "def simple_function" in bundle.main_source
        assert bundle.args == [1, 2]
        assert bundle.checksum  # Should have checksum

    def test_extract_function_with_imports(self):
        """Should detect imports in function source."""
        def function_with_imports():
            import json
            import os
            return json.dumps({"cwd": os.getcwd()})

        bundle = extract_source(function_with_imports)

        # json and os are stdlib, shouldn't be in pip_packages
        assert "json" not in bundle.pip_packages
        assert "os" not in bundle.pip_packages

    def test_extract_function_with_third_party_imports(self):
        """Should identify pip packages from imports."""
        # Skip if numpy not installed - the function won't serialize properly without it
        try:
            import numpy
        except ImportError:
            pytest.skip("numpy not installed")

        def ml_function():
            import numpy as np
            return np.zeros(10)

        can_extract, reason = can_extract_source(ml_function)
        assert can_extract is True

    def test_generate_runner_script(self):
        """Should generate executable runner script."""
        def add(a, b):
            return a + b

        bundle = extract_source(add, 5, 3)
        script = bundle.generate_runner_script()

        assert "def add" in script
        assert "__GATEWAY_RESULT_START__" in script
        assert "__GATEWAY_RESULT_END__" in script
        assert "_result = add(*_args, **_kwargs)" in script

    def test_cannot_extract_builtin(self):
        """Built-in functions should not be extractable."""
        can_extract, reason = can_extract_source(len)
        assert can_extract is False
        assert "Source code not available" in reason or "Cannot get source" in reason

    def test_cannot_extract_lambda(self):
        """Lambdas have limited source but can be extracted."""
        def func(x):
            return x * 2
        can_extract, reason = can_extract_source(func)
        # Lambdas can often be extracted (single line)
        # The result depends on the context


class TestHybridSerializer:
    """Test hybrid serialization (auto-selects source vs pickle)."""

    def test_normal_function_uses_source(self):
        """Normal functions should prefer source serialization."""
        def normal_function(x):
            return x * 2

        serializer = HybridSerializer(prefer_source=True)
        data, method = serializer.serialize(normal_function, args=(5,))

        assert method == "source"
        assert isinstance(data, SourceBundle)

    def test_lambda_uses_pickle(self):
        """Lambdas should use pickle serialization."""
        # Use an actual lambda - lambdas cannot be source-extracted
        func = lambda x: x * 2  # noqa: E731

        serializer = HybridSerializer(prefer_source=True)
        data, method = serializer.serialize(func, args=(5,))

        assert method == "pickle"

    def test_force_source_method(self):
        """Should respect forced source method."""
        def func(x):
            return x

        serializer = HybridSerializer(prefer_source=False)  # Prefer pickle
        data, method = serializer.serialize(func, args=(1,), force_method="source")

        assert method == "source"

    def test_force_pickle_method(self):
        """Should respect forced pickle method."""
        def func(x):
            return x

        serializer = HybridSerializer(prefer_source=True)
        data, method = serializer.serialize(func, args=(1,), force_method="pickle")

        assert method == "pickle"


class TestSmartSerialize:
    """Test the smart_serialize convenience function."""

    def test_smart_serialize_auto(self):
        """Auto mode should choose intelligently."""
        def simple_func():
            return 42

        data, method = smart_serialize(simple_func)
        assert method in ["source", "pickle"]

    def test_smart_serialize_with_args(self):
        """Should pass args to serialized bundle."""
        def add(a, b):
            return a + b

        data, method = smart_serialize(add, 1, 2, method="source")

        assert method == "source"
        assert data.args == [1, 2]


# =============================================================================
# Tests for Smart Executor
# =============================================================================

class TestSmartExecutor:
    """Test the smart executor that combines all optimizations."""

    def test_plan_short_cpu_job(self):
        """Short CPU job should plan for serverless."""
        executor = SmartExecutor()
        resources = ResourceRequirements(
            cpu=2.0,
            memory_gb=4.0,
            estimated_duration_seconds=60,
        )

        def simple_job():
            return 42

        plan = executor.plan(simple_job, resources)

        assert plan.backend == ExecutionBackend.SERVERLESS
        assert plan.estimated_cold_start_seconds < 10  # Serverless is fast

    def test_plan_long_gpu_job(self):
        """GPU job exceeding serverless timeout should plan for VM."""
        executor = SmartExecutor()
        resources = ResourceRequirements(
            cpu=8.0,
            memory_gb=32.0,
            gpu_type="A100",
            gpu_count=1,
            estimated_duration_seconds=7200,  # 2 hours — exceeds Cloud Run 3600s limit
        )

        def training_job():
            return {"loss": 0.01}

        plan = executor.plan(training_job, resources)

        assert plan.backend == ExecutionBackend.VM
        assert plan.estimated_cold_start_seconds >= 30

    def test_plan_with_forced_backend(self):
        """Should respect forced backend."""
        executor = SmartExecutor(force_backend=ExecutionBackend.LOCAL)
        resources = ResourceRequirements()

        def func():
            return 1

        plan = executor.plan(func, resources)

        assert plan.backend == ExecutionBackend.LOCAL

    def test_plan_shows_cost_estimate(self):
        """Plan should include cost estimate."""
        executor = SmartExecutor()
        resources = ResourceRequirements(
            cpu=2.0,
            memory_gb=4.0,
            estimated_duration_seconds=300,
        )

        def func():
            return 1

        plan = executor.plan(func, resources)

        assert plan.estimated_cost_usd >= 0
        assert "USD" in str(plan) or "$" in str(plan)

    def test_resource_requirements_to_serverless_config(self):
        """ResourceRequirements should convert to ServerlessConfig."""
        resources = ResourceRequirements(
            cpu=4.0,
            memory_gb=8.0,
            timeout_seconds=600,
            pip_packages=["numpy", "pandas"],
        )

        config = resources.to_serverless_config()

        assert config.cpu == 4.0
        assert config.memory_mb == 8192
        assert config.timeout_seconds == 600
        assert "numpy" in config.pip_packages


# =============================================================================
# Tests for Local Docker Executor
# =============================================================================

class TestLocalDockerExecutor:
    """Test local Docker execution."""

    @pytest.mark.asyncio
    async def test_docker_availability_check(self):
        """Should check if Docker is available."""
        executor = LocalDockerExecutor()
        available = await executor.is_available()
        # This is environment-dependent, just check it returns bool
        assert isinstance(available, bool)

    def test_pricing_is_free(self):
        """Local execution should be free."""
        executor = LocalDockerExecutor()
        config = ServerlessConfig()
        pricing = executor.get_pricing(config)

        assert pricing["estimated_total"] == 0.0


# =============================================================================
# Integration Tests
# =============================================================================

class TestIntegration:
    """Integration tests for the full execution flow."""

    @pytest.mark.asyncio
    async def test_local_execution_simple_function(self):
        """Should execute a simple function locally."""
        executor = SmartExecutor(force_backend=ExecutionBackend.LOCAL)
        resources = ResourceRequirements(estimated_duration_seconds=5)

        def add_numbers():
            return 1 + 1

        # Only run if Docker is available
        docker = LocalDockerExecutor()
        if not await docker.is_available():
            # Fall back to direct execution
            result = await executor.execute(
                add_numbers,
                resources,
                output_callback=lambda x: None,
            )
            # May succeed with direct execution or fail gracefully
            return

        result = await executor.execute(
            add_numbers,
            resources,
            output_callback=lambda x: None,
        )

        # Result depends on Docker availability
        assert result.start_time is not None
        assert result.end_time is not None

    def test_execution_plan_str_representation(self):
        """ExecutionPlan should have readable string representation."""
        plan = ExecutionPlan(
            backend=ExecutionBackend.SERVERLESS,
            serialization=SerializationMethod.SOURCE,
            provider="gcp",
            service="cloud_run",
            region="us-central1",
            estimated_cost_usd=0.001,
            estimated_cold_start_seconds=2.0,
            reason="Short CPU job",
        )

        plan_str = str(plan)

        assert "serverless" in plan_str.lower()
        assert "cloud_run" in plan_str
        assert "gcp" in plan_str
        assert "$0.001" in plan_str or "0.001" in plan_str


# =============================================================================
# Edge Cases and Error Handling
# =============================================================================

class TestEdgeCases:
    """Test edge cases and error handling."""

    def test_boundary_duration_10_minutes(self):
        """10 minutes is the threshold - should be serverless."""
        tier, _ = select_execution_tier(
            estimated_duration_seconds=600,  # Exactly 10 min
            gpu_required=False,
            gpu_type=None,
            memory_gb=4.0,
        )
        # 600s is the threshold, should still be serverless
        assert tier == ExecutionTier.SERVERLESS

    def test_just_over_10_minutes(self):
        """Just over 10 minutes should use VM."""
        tier, _ = select_execution_tier(
            estimated_duration_seconds=601,  # 10 min + 1 sec
            gpu_required=False,
            gpu_type=None,
            memory_gb=4.0,
        )
        assert tier == ExecutionTier.VM

    def test_unknown_gpu_type(self):
        """Unknown GPU type should use VM to be safe."""
        tier, effective_gpu = select_execution_tier(
            estimated_duration_seconds=60,
            gpu_required=True,
            gpu_type="UNKNOWN_GPU",
            memory_gb=4.0,
        )
        assert tier == ExecutionTier.VM
        assert effective_gpu == "UNKNOWN_GPU"  # No substitution available

    def test_empty_function(self):
        """Should handle empty function."""
        def empty():
            pass

        bundle = extract_source(empty)
        assert bundle.entry_point == "empty"
        assert "def empty" in bundle.main_source


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
