"""
Tests for source_code passthrough from client to LLM analyzer.

Verifies that:
1. The client extracts source code and includes it in the request payload.
2. The LLM analyzer uses pre-extracted source instead of inspect.getsource().
3. The static AST fallback also works with pre-extracted source.
4. The full chain (auto_detect_resources → llm_detect_resources → LLMAnalyzer)
   passes source_code through correctly.
"""

import base64
import inspect
import textwrap
from unittest.mock import MagicMock, patch

import cloudpickle
import pytest

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


# ============================================================================
# Test Functions
# ============================================================================

def sample_torch_function():
    """Sample function that would need GPU."""
    import torch
    model = torch.nn.Linear(100, 10)
    x = torch.randn(32, 100)
    return model(x)


def simple_cpu_function():
    x = 0
    for i in range(100):
        x += i
    return x


# ============================================================================
# Test: Client sends source_code
# ============================================================================

class TestClientSourceExtraction:
    """Test that the client extracts and sends source_code."""

    def test_client_includes_source_code_in_request(self):
        """Verify the client builds request_data with source_code field."""
        func = simple_cpu_function
        source = textwrap.dedent(inspect.getsource(func))

        # Simulate what the client does
        payload = cloudpickle.dumps((func, (), {}))
        payload_b64 = base64.b64encode(payload).decode("utf-8")

        # This mirrors the client's request_data construction
        import sys
        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"

        extracted_source = None
        try:
            extracted_source = textwrap.dedent(inspect.getsource(func))
        except (OSError, TypeError):
            pass

        request_data = {
            "function": {
                "payload": payload_b64,
                "type": "cloudpickle",
                "python_version": python_version,
                "source_code": extracted_source,
            },
        }

        assert request_data["function"]["source_code"] is not None
        assert "simple_cpu_function" in request_data["function"]["source_code"]
        assert "for i in range" in request_data["function"]["source_code"]


# ============================================================================
# Test: LLM Analyzer uses pre-extracted source
# ============================================================================

class TestLLMAnalyzerSourcePassthrough:
    """Test that LLMAnalyzer.analyze() uses source_code param."""

    def test_analyze_uses_provided_source_code(self):
        """When source_code is provided, inspect.getsource should NOT be called."""
        from verlex_server.gateway.core.llm_analyzer import LLMAnalyzer

        analyzer = LLMAnalyzer()
        source = textwrap.dedent(inspect.getsource(simple_cpu_function))

        # Mock the client property to avoid needing a real API key
        with patch.object(LLMAnalyzer, 'client', new_callable=lambda: property(lambda self: MagicMock())):
            # Mock _call_llm to avoid actual API call
            with patch.object(analyzer, '_call_llm') as mock_call_llm:
                mock_call_llm.return_value = {
                    "cpu": 2,
                    "memory_gb": 4,
                    "gpu_type": None,
                    "gpu_count": 0,
                    "workload_type": "general",
                    "reasoning": ["Simple CPU function"],
                    "confidence": 0.9,
                    "execution_tier": "serverless",
                    "execution_tier_reason": "Simple workload",
                    "recommended_provider": "gcp",
                    "recommended_instance": "n1-standard-2",
                    "use_spot": True,
                    "estimated_cost_per_hour": 0.01,
                    "estimated_duration_seconds": 60,
                    "estimated_total_cost": 0.01,
                }

                # Should NOT raise even though inspect.getsource would fail
                # on a cloudpickle'd function (we pass source directly)
                with patch('verlex_server.gateway.core.llm_analyzer.inspect.getsource',
                           side_effect=OSError("source not available")):
                    rec = analyzer.analyze(
                        simple_cpu_function,
                        source_code=source,
                    )

                # Verify _call_llm received the source we provided
                call_args = mock_call_llm.call_args
                assert source in call_args[0][0]  # first positional arg is source_code

    def test_analyze_falls_back_to_inspect_when_no_source(self):
        """When source_code is None, should fall back to inspect.getsource()."""
        from verlex_server.gateway.core.llm_analyzer import LLMAnalyzer, _cache

        # Clear the global cache so a previous test's result doesn't interfere
        _cache.clear()

        analyzer = LLMAnalyzer()
        # Disable caching for this test
        analyzer.config.cache_enabled = False

        with patch.object(LLMAnalyzer, 'client', new_callable=lambda: property(lambda self: MagicMock())):
            with patch.object(analyzer, '_call_llm') as mock_call_llm:
                mock_call_llm.return_value = {
                    "cpu": 2, "memory_gb": 4, "gpu_type": None, "gpu_count": 0,
                    "workload_type": "general", "reasoning": ["test"],
                    "confidence": 0.9, "execution_tier": "serverless",
                    "execution_tier_reason": "test", "use_spot": True,
                    "estimated_cost_per_hour": 0.01,
                    "estimated_duration_seconds": 60,
                    "estimated_total_cost": 0.01,
                }

                # Without source_code param, should use inspect.getsource
                rec = analyzer.analyze(simple_cpu_function)
                assert mock_call_llm.called

    def test_analyze_raises_when_no_source_and_inspect_fails(self):
        """When source_code is None and inspect fails, should raise ValueError."""
        from verlex_server.gateway.core.llm_analyzer import LLMAnalyzer

        analyzer = LLMAnalyzer()

        with patch('verlex_server.gateway.core.llm_analyzer.inspect.getsource',
                   side_effect=OSError("source not available")):
            with pytest.raises(ValueError, match="Cannot extract source"):
                analyzer.analyze(simple_cpu_function, source_code=None)


# ============================================================================
# Test: Static analysis (_analyze_code) uses pre-extracted source
# ============================================================================

class TestStaticAnalysisSourcePassthrough:
    """Test that _analyze_code uses source_code param."""

    def test_analyze_code_with_provided_source(self):
        """_analyze_code should parse provided source instead of calling inspect."""
        source = textwrap.dedent(inspect.getsource(sample_torch_function))

        with patch('verlex_server.gateway.core.profiling.inspect.getsource',
                   side_effect=OSError("source not available")):
            # Without source_code, should get empty analyzer (inspect fails)
            analyzer_empty = _analyze_code(sample_torch_function)
            assert len(analyzer_empty.imports) == 0

            # With source_code, should detect imports
            analyzer_with_source = _analyze_code(sample_torch_function, source_code=source)
            assert "torch" in analyzer_with_source.imports

    def test_analyze_code_falls_back_to_inspect(self):
        """Without source_code, _analyze_code should still try inspect."""
        # For locally-defined functions, inspect.getsource() works
        analyzer = _analyze_code(simple_cpu_function)
        # Should have parsed something (range call, etc.)
        assert isinstance(analyzer.function_calls, list)


# ============================================================================
# Test: auto_detect_resources passes source_code through
# ============================================================================

class TestAutoDetectSourcePassthrough:
    """Test the full chain: auto_detect_resources → llm/static with source_code."""

    def test_auto_detect_with_source_code_static_fallback(self):
        """auto_detect_resources should pass source_code to static analysis."""
        source = textwrap.dedent(inspect.getsource(sample_torch_function))

        # Disable LLM so it falls through to static analysis
        rec = auto_detect_resources(
            sample_torch_function,
            use_llm=False,
            source_code=source,
        )

        assert isinstance(rec, ResourceRecommendation)
        # With source code, static analysis should detect torch → GPU workload
        assert rec.gpu_type is not None, "Should detect torch and recommend GPU"
        assert rec.workload_type in (WorkloadType.ML_INFERENCE, WorkloadType.ML_TRAINING)

    def test_auto_detect_without_source_fails_gracefully_for_unpicklable(self):
        """Without source, auto_detect still works (inspect fallback or defaults)."""
        rec = auto_detect_resources(
            simple_cpu_function,
            use_llm=False,
            source_code=None,
        )

        assert isinstance(rec, ResourceRecommendation)
        # Should still produce a valid recommendation via inspect fallback
        assert rec.cpu >= 1


# ============================================================================
# Test: Simulated cloudpickle round-trip (the real-world scenario)
# ============================================================================

class TestCloudpickleRoundTrip:
    """Simulate the real production scenario: function pickled on client,
    unpickled on server, and source_code passed alongside."""

    def test_cloudpickle_function_with_source_code(self):
        """Cloudpickle'd function + source_code should enable analysis."""
        # Client side: serialize function AND extract source
        func = sample_torch_function
        pickled = cloudpickle.dumps((func, (), {}))
        source = textwrap.dedent(inspect.getsource(func))

        # Server side: deserialize and analyze
        data = cloudpickle.loads(pickled)
        restored_func, args, kwargs = data

        # With source_code, static analysis should work even on restored function
        rec = auto_detect_resources(
            restored_func,
            use_llm=False,
            source_code=source,
        )

        assert isinstance(rec, ResourceRecommendation)
        assert rec.gpu_type is not None, "Should detect torch imports from provided source"

    def test_cloudpickle_function_without_source_still_works(self):
        """Even without source_code, the function should still get defaults."""
        func = simple_cpu_function
        pickled = cloudpickle.dumps((func, (), {}))
        data = cloudpickle.loads(pickled)
        restored_func, args, kwargs = data

        # Without source_code — in production, inspect.getsource might still
        # work for local functions but would fail for remote ones
        rec = auto_detect_resources(
            restored_func,
            use_llm=False,
        )

        assert isinstance(rec, ResourceRecommendation)
        assert rec.cpu >= 1
