"""Integration tests for Modal-based CUDA compilation.
These tests require Modal to be configured with valid credentials.
They are skipped if Modal is not available or not authenticated.
Run with:
    MODAL_TOKEN_ID=... MODAL_TOKEN_SECRET=... pytest test_modal_integration.py -v
"""
import os
from pathlib import Path
import pytest

pytestmark = pytest.mark.skipif(
    not os.environ.get("MODAL_TOKEN_ID") or not os.environ.get("MODAL_TOKEN_SECRET"),
    reason="Modal credentials not configured (set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET)",
)
# Test data directory
TEST_DATA_DIR = Path(__file__).parent / "test_data"
# =============================================================================
# Test CUDA Code Samples
# =============================================================================
SIMPLE_KERNEL = """\
__global__ void test_kernel() {
    // Empty kernel for testing
}
"""
VECTOR_ADD_KERNEL = """\
__global__ void vector_add(float* a, float* b, float* c, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}
"""
SHARED_MEMORY_KERNEL = """\
#define TILE_SIZE 16
__global__ void tiled_matmul(float* A, float* B, float* C, int N) {
    __shared__ float As[TILE_SIZE][TILE_SIZE];
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    int row = by * TILE_SIZE + ty;
    int col = bx * TILE_SIZE + tx;
    float sum = 0.0f;
    for (int t = 0; t < (N + TILE_SIZE - 1) / TILE_SIZE; t++) {
        if (row < N && t * TILE_SIZE + tx < N)
            As[ty][tx] = A[row * N + t * TILE_SIZE + tx];
        else
            As[ty][tx] = 0.0f;
        if (t * TILE_SIZE + ty < N && col < N)
            Bs[ty][tx] = B[(t * TILE_SIZE + ty) * N + col];
        else
            Bs[ty][tx] = 0.0f;
        __syncthreads();
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += As[ty][k] * Bs[k][tx];
        }
        __syncthreads();
    }
    if (row < N && col < N) {
        C[row * N + col] = sum;
    }
}
"""
SYNTAX_ERROR_KERNEL = """\
__global__ void broken_kernel() {
    // Missing semicolon below
    int x = 1
}
"""
# =============================================================================
# Tests
# =============================================================================
class TestModalCompileSimple:
    """Test simple kernel compilation via Modal."""
    @pytest.mark.slow
    def test_simple_kernel_ptx(self) -> None:
        """Test compiling a simple kernel to PTX."""
        from wafer_core.tools.compile import (
            CompileRequest,
            OutputFormat,
        )
        from wafer_core.tools.compile.modal_compile import compile_cuda
        request = CompileRequest(
            files={"kernel.cu": SIMPLE_KERNEL},
            arch="sm_90a",
            output=(OutputFormat.PTX,),
        )
        # Call Modal function
        result = compile_cuda.remote({
            "files": request.files,
            "arch": request.arch,
            "flags": list(request.flags),
            "output": [fmt.value for fmt in request.output],
        })
        assert result["success"] is True
        assert result["ptx"] is not None
        assert ".version" in result["ptx"]
        assert ".target" in result["ptx"]
        assert "test_kernel" in result["ptx"]
    @pytest.mark.slow
    def test_vector_add_kernel_sass(self) -> None:
        """Test compiling vector_add kernel to SASS."""
        from wafer_core.tools.compile import (
            CompileRequest,
            OutputFormat,
        )
        from wafer_core.tools.compile.modal_compile import compile_cuda
        request = CompileRequest(
            files={"kernel.cu": VECTOR_ADD_KERNEL},
            arch="sm_90a",
            output=(OutputFormat.SASS,),
        )
        result = compile_cuda.remote({
            "files": request.files,
            "arch": request.arch,
            "flags": list(request.flags),
            "output": [fmt.value for fmt in request.output],
        })
        assert result["success"] is True
        assert result["sass"] is not None
        # SASS should contain assembly instructions
        assert "vector_add" in result["sass"]
    @pytest.mark.slow
    def test_both_ptx_and_sass(self) -> None:
        """Test generating both PTX and SASS."""
        from wafer_core.tools.compile import (
            CompileRequest,
            OutputFormat,
        )
        from wafer_core.tools.compile.modal_compile import compile_cuda
        request = CompileRequest(
            files={"kernel.cu": VECTOR_ADD_KERNEL},
            arch="sm_90a",
            output=(OutputFormat.PTX, OutputFormat.SASS),
        )
        result = compile_cuda.remote({
            "files": request.files,
            "arch": request.arch,
            "flags": list(request.flags),
            "output": [fmt.value for fmt in request.output],
        })
        assert result["success"] is True
        assert result["ptx"] is not None
        assert result["sass"] is not None
class TestModalCompileMultiFile:
    """Test multi-file compilation via Modal."""
    @pytest.mark.slow
    def test_kernel_with_header(self) -> None:
        """Test compiling a kernel that includes a header."""
        from wafer_core.tools.compile.modal_compile import compile_cuda
        main_cu = '#include "utils.cuh"\n' + """
__global__ void apply_square(float* data, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        data[idx] = square(data[idx]);
    }
}
"""
        header = """
#pragma once
__device__ float square(float x) {
    return x * x;
}
"""
        result = compile_cuda.remote({
            "files": {
                "main.cu": main_cu,
                "utils.cuh": header,
            },
            "arch": "sm_90a",
            "flags": [],
            "output": ["ptx"],
        })
        assert result["success"] is True
        assert result["ptx"] is not None
        assert "apply_square" in result["ptx"]
class TestModalCompileArchitectures:
    """Test compilation for different GPU architectures."""
    @pytest.mark.slow
    @pytest.mark.parametrize("arch", ["sm_80", "sm_89", "sm_90a"])
    def test_different_architectures(self, arch: str) -> None:
        """Test compiling for different architectures."""
        from wafer_core.tools.compile.modal_compile import compile_cuda
        result = compile_cuda.remote({
            "files": {"kernel.cu": SIMPLE_KERNEL},
            "arch": arch,
            "flags": [],
            "output": ["ptx"],
        })
        assert result["success"] is True
        assert result["ptx"] is not None
        assert f".target {arch}" in result["ptx"]
class TestModalCompileErrors:
    """Test error handling in Modal compilation."""
    @pytest.mark.slow
    def test_syntax_error(self) -> None:
        """Test that syntax errors are reported."""
        from wafer_core.tools.compile.modal_compile import compile_cuda
        result = compile_cuda.remote({
            "files": {"kernel.cu": SYNTAX_ERROR_KERNEL},
            "arch": "sm_90a",
            "flags": [],
            "output": ["ptx"],
        })
        assert result["success"] is False
        assert result["stderr"] != ""
        # Error should mention the missing semicolon or syntax error
        assert "error" in result["stderr"].lower()
    @pytest.mark.slow
    def test_missing_include(self) -> None:
        """Test that missing includes are reported."""
        from wafer_core.tools.compile.modal_compile import compile_cuda
        kernel = '#include "nonexistent.h"\n__global__ void test() {}'
        result = compile_cuda.remote({
            "files": {"kernel.cu": kernel},
            "arch": "sm_90a",
            "flags": [],
            "output": ["ptx"],
        })
        assert result["success"] is False
        assert "nonexistent.h" in result["stderr"] or "cannot open" in result["stderr"].lower()
class TestModalCompileFlags:
    """Test custom compiler flags."""
    @pytest.mark.slow
    def test_optimization_flags(self) -> None:
        """Test compilation with optimization flags."""
        from wafer_core.tools.compile.modal_compile import compile_cuda
        result = compile_cuda.remote({
            "files": {"kernel.cu": VECTOR_ADD_KERNEL},
            "arch": "sm_90a",
            "flags": ["-O3"],
            "output": ["ptx"],
        })
        assert result["success"] is True
        assert result["ptx"] is not None
    @pytest.mark.slow
    def test_lineinfo_flag(self) -> None:
        """Test compilation with lineinfo flag."""
        from wafer_core.tools.compile.modal_compile import compile_cuda
        result = compile_cuda.remote({
            "files": {"kernel.cu": VECTOR_ADD_KERNEL},
            "arch": "sm_90a",
            "flags": ["-lineinfo"],
            "output": ["ptx"],
        })
        assert result["success"] is True
        assert result["ptx"] is not None
class TestModalHealthCheck:
    """Test Modal health check function."""
    @pytest.mark.slow
    def test_health_check(self) -> None:
        """Test the health check endpoint."""
        from wafer_core.tools.compile.modal_compile import health_check
        result = health_check.remote()
        assert result["status"] == "ok"
        assert result["nvcc_available"] is True
        assert "nvcc_version" in result
