"""Tests for the cloud CUDA compiler.

These tests verify request validation, serialization, and compilation logic.
Integration tests that actually invoke nvcc are in test_modal_integration.py.
"""

import pytest

from wafer_core.tools.compile.types import (
    CompileRequest,
    CompileResponse,
    OutputFormat,
    VALID_ARCHITECTURES,
)


class TestCompileRequest:
    """Test request validation and serialization."""

    def test_single_file_request(self) -> None:
        """Test creating a request with a single file."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
        )
        assert request.files == {"kernel.cu": "__global__ void test() {}"}
        assert request.arch == "sm_90a"  # default
        assert request.flags == ()
        assert request.output == (OutputFormat.PTX, OutputFormat.SASS)
        assert request.main_cu_file == "kernel.cu"

    def test_multi_file_request(self) -> None:
        """Test creating a request with multiple files."""
        request = CompileRequest(
            files={
                "main.cu": '#include "utils.cuh"\n__global__ void test() {}',
                "utils.cuh": "__device__ float square(float x) { return x * x; }",
            },
        )
        assert len(request.files) == 2
        assert "main.cu" in request.files
        assert "utils.cuh" in request.files

    def test_custom_arch(self) -> None:
        """Test specifying a custom architecture."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            arch="sm_80",
        )
        assert request.arch == "sm_80"

    def test_custom_flags(self) -> None:
        """Test specifying custom compiler flags."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            flags=("-O3", "--maxrregcount=64", "-lineinfo"),
        )
        assert request.flags == ("-O3", "--maxrregcount=64", "-lineinfo")

    def test_ptx_only_output(self) -> None:
        """Test requesting only PTX output."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            output=(OutputFormat.PTX,),
        )
        assert request.output == (OutputFormat.PTX,)

    def test_sass_only_output(self) -> None:
        """Test requesting only SASS output."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            output=(OutputFormat.SASS,),
        )
        assert request.output == (OutputFormat.SASS,)

    def test_invalid_arch_rejected(self) -> None:
        """Test that invalid architecture is rejected."""
        with pytest.raises(ValueError, match="Invalid architecture"):
            CompileRequest(
                files={"kernel.cu": "__global__ void test() {}"},
                arch="sm_999",
            )

    def test_empty_files_rejected(self) -> None:
        """Test that empty files dict is rejected."""
        with pytest.raises(ValueError, match="At least one file is required"):
            CompileRequest(files={})

    def test_no_cu_file_rejected(self) -> None:
        """Test that missing .cu file is rejected."""
        with pytest.raises(ValueError, match="At least one .cu file is required"):
            CompileRequest(files={"utils.cuh": "// header only"})

    def test_empty_output_rejected(self) -> None:
        """Test that empty output list is rejected."""
        with pytest.raises(ValueError, match="At least one output format is required"):
            CompileRequest(
                files={"kernel.cu": "__global__ void test() {}"},
                output=(),
            )

    def test_request_is_frozen(self) -> None:
        """Test that request is immutable."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
        )
        with pytest.raises(AttributeError):
            request.arch = "sm_80"  # type: ignore[misc]


class TestCompileResponse:
    """Test response creation and helper methods."""

    def test_success_response(self) -> None:
        """Test creating a successful response."""
        response = CompileResponse(
            success=True,
            ptx=".version 8.0\n.target sm_90a",
            sass="MOV R0, R1;",
            compilation_time_ms=150,
        )
        assert response.success is True
        assert response.ptx is not None
        assert response.sass is not None
        assert response.stderr == ""
        assert response.compilation_time_ms == 150

    def test_error_response_helper(self) -> None:
        """Test the error response helper method."""
        response = CompileResponse.error(
            "kernel.cu(10): error: expected a ';'",
            compilation_time_ms=50,
        )
        assert response.success is False
        assert response.ptx is None
        assert response.sass is None
        assert "expected a ';'" in response.stderr
        assert response.compilation_time_ms == 50

    def test_response_is_frozen(self) -> None:
        """Test that response is immutable."""
        response = CompileResponse(success=True)
        with pytest.raises(AttributeError):
            response.success = False  # type: ignore[misc]


class TestValidArchitectures:
    """Test the set of valid architectures."""

    def test_common_architectures_valid(self) -> None:
        """Test that common architectures are in the valid set."""
        common_archs = ["sm_80", "sm_86", "sm_89", "sm_90", "sm_90a"]
        for arch in common_archs:
            assert arch in VALID_ARCHITECTURES, f"{arch} should be valid"

    def test_hopper_architectures(self) -> None:
        """Test Hopper architectures."""
        assert "sm_90" in VALID_ARCHITECTURES
        assert "sm_90a" in VALID_ARCHITECTURES

    def test_blackwell_architectures(self) -> None:
        """Test Blackwell architectures."""
        assert "sm_100" in VALID_ARCHITECTURES
        assert "sm_100a" in VALID_ARCHITECTURES


class TestOutputFormat:
    """Test output format enum."""

    def test_ptx_value(self) -> None:
        """Test PTX enum value."""
        assert OutputFormat.PTX.value == "ptx"
        assert OutputFormat.PTX == "ptx"  # str enum comparison

    def test_sass_value(self) -> None:
        """Test SASS enum value."""
        assert OutputFormat.SASS.value == "sass"
        assert OutputFormat.SASS == "sass"  # str enum comparison


# =============================================================================
# Test CUDA Code Examples
# =============================================================================

# These are example CUDA kernels used in tests.
# They are validated for correct syntax in integration tests.

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];
    }
}
"""

MATRIX_MULTIPLY_KERNEL = """\
__global__ void matmul(float* A, float* B, float* C, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}
"""

REDUCTION_KERNEL = """\
__global__ void reduce_sum(float* input, float* output, int n) {
    extern __shared__ float sdata[];

    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    sdata[tid] = (i < n) ? input[i] : 0.0f;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        output[blockIdx.x] = sdata[0];
    }
}
"""

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;
    }
}
"""


class TestSimpleKernels:
    """Test that simple kernel code can be used in requests."""

    def test_vector_add_kernel(self) -> None:
        """Test vector add kernel can be packaged in a request."""
        request = CompileRequest(
            files={"vector_add.cu": VECTOR_ADD_KERNEL},
        )
        assert "vector_add" in request.files["vector_add.cu"]

    def test_matrix_multiply_kernel(self) -> None:
        """Test matrix multiply kernel can be packaged in a request."""
        request = CompileRequest(
            files={"matmul.cu": MATRIX_MULTIPLY_KERNEL},
        )
        assert "matmul" in request.files["matmul.cu"]

    def test_reduction_kernel(self) -> None:
        """Test reduction kernel can be packaged in a request."""
        request = CompileRequest(
            files={"reduce.cu": REDUCTION_KERNEL},
        )
        assert "__shared__" in request.files["reduce.cu"]

    def test_shared_memory_kernel(self) -> None:
        """Test shared memory kernel can be packaged in a request."""
        request = CompileRequest(
            files={"tiled_matmul.cu": SHARED_MEMORY_KERNEL},
        )
        assert "TILE_SIZE" in request.files["tiled_matmul.cu"]


# Multi-file test examples

HEADER_FILE = """\
#pragma once

__device__ float square(float x) {
    return x * x;
}

__device__ float cube(float x) {
    return x * x * x;
}
"""

MAIN_FILE_WITH_HEADER = """\
#include "utils.cuh"

__global__ void apply_square(float* data, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        data[idx] = square(data[idx]);
    }
}
"""

NESTED_HEADER_A = """\
#pragma once

#define MY_CONSTANT 42
"""

NESTED_HEADER_B = """\
#pragma once

#include "constants.cuh"

__device__ int get_constant() {
    return MY_CONSTANT;
}
"""

MAIN_WITH_NESTED = """\
#include "helper.cuh"

__global__ void use_constant(int* out) {
    *out = get_constant();
}
"""


class TestMultiFile:
    """Test multi-file compilation requests."""

    def test_kernel_with_header(self) -> None:
        """Test kernel that includes a header."""
        request = CompileRequest(
            files={
                "main.cu": MAIN_FILE_WITH_HEADER,
                "utils.cuh": HEADER_FILE,
            },
        )
        assert len(request.files) == 2
        assert request.main_cu_file == "main.cu"

    def test_kernel_with_multiple_headers(self) -> None:
        """Test kernel with multiple headers."""
        extra_header = "__device__ float add(float a, float b) { return a + b; }"
        request = CompileRequest(
            files={
                "main.cu": '#include "utils.cuh"\n#include "math.cuh"\n'
                + MAIN_FILE_WITH_HEADER.split("\n", 1)[1],
                "utils.cuh": HEADER_FILE,
                "math.cuh": extra_header,
            },
        )
        assert len(request.files) == 3

    def test_nested_includes(self) -> None:
        """Test nested header includes."""
        request = CompileRequest(
            files={
                "main.cu": MAIN_WITH_NESTED,
                "helper.cuh": NESTED_HEADER_B,
                "constants.cuh": NESTED_HEADER_A,
            },
        )
        assert len(request.files) == 3
        assert "MY_CONSTANT" in request.files["constants.cuh"]

    def test_relative_include_paths(self) -> None:
        """Test that relative include paths are preserved in files dict."""
        request = CompileRequest(
            files={
                "src/main.cu": '#include "../include/utils.cuh"\n__global__ void test() {}',
                "include/utils.cuh": "// utils",
            },
        )
        assert "src/main.cu" in request.files
        assert "include/utils.cuh" in request.files


# PyTorch extension examples

PYTORCH_ACCESSOR_KERNEL = """\
#include <torch/types.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void fused_add_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ out,
    int64_t size
) {
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}
"""

PYTORCH_PACKED_ACCESSOR_KERNEL = """\
#include <torch/types.h>
#include <ATen/ATen.h>
#include <cuda.h>
#include <cuda_runtime.h>

template <typename scalar_t>
__global__ void accessor_kernel(
    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> input,
    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> output
) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    int col = blockIdx.y * blockDim.y + threadIdx.y;

    if (row < input.size(0) && col < input.size(1)) {
        output[row][col] = input[row][col] * 2;
    }
}
"""


class TestPyTorchHeaders:
    """Test code using PyTorch CUDA types."""

    def test_pytorch_accessor(self) -> None:
        """Test PyTorch accessor kernel can be packaged."""
        request = CompileRequest(
            files={"pytorch_kernel.cu": PYTORCH_ACCESSOR_KERNEL},
        )
        assert "torch/types.h" in request.files["pytorch_kernel.cu"]

    def test_pytorch_packed_accessor(self) -> None:
        """Test PyTorch packed accessor kernel can be packaged."""
        request = CompileRequest(
            files={"accessor_kernel.cu": PYTORCH_PACKED_ACCESSOR_KERNEL},
        )
        assert "PackedTensorAccessor" in request.files["accessor_kernel.cu"]


# CUTLASS examples

CUTLASS_TYPES_KERNEL = """\
#include <cutlass/cutlass.h>
#include <cutlass/numeric_types.h>

__global__ void use_cutlass_types() {
    cutlass::half_t a = cutlass::half_t(1.0f);
    cutlass::half_t b = cutlass::half_t(2.0f);
    cutlass::half_t c = a + b;
}
"""

CUTLASS_GEMM_INCLUDE = """\
#include <cutlass/cutlass.h>
#include <cutlass/gemm/device/gemm.h>

// Just testing include paths work
using ColumnMajor = cutlass::layout::ColumnMajor;
"""


class TestCUTLASSHeaders:
    """Test code using CUTLASS headers."""

    def test_cutlass_types(self) -> None:
        """Test CUTLASS types kernel can be packaged."""
        request = CompileRequest(
            files={"cutlass_kernel.cu": CUTLASS_TYPES_KERNEL},
        )
        assert "cutlass/cutlass.h" in request.files["cutlass_kernel.cu"]

    def test_cutlass_gemm_include(self) -> None:
        """Test CUTLASS GEMM include can be packaged."""
        request = CompileRequest(
            files={"gemm_kernel.cu": CUTLASS_GEMM_INCLUDE},
        )
        assert "cutlass/gemm" in request.files["gemm_kernel.cu"]


class TestOutputFormats:
    """Test output format expectations for responses."""

    def test_ptx_output_structure(self) -> None:
        """Test PTX output has expected structure."""
        # PTX typically starts with version/target info
        mock_ptx = """\
.version 8.0
.target sm_90a
.address_size 64

.visible .entry vector_add(
    .param .u64 a,
    .param .u64 b,
    .param .u64 c,
    .param .u32 n
)
{
    .reg .pred %p<2>;
    .reg .f32 %f<4>;
    // ...
}
"""
        response = CompileResponse(
            success=True,
            ptx=mock_ptx,
            compilation_time_ms=100,
        )
        assert response.ptx is not None
        assert ".version" in response.ptx
        assert ".target" in response.ptx
        assert ".entry" in response.ptx

    def test_sass_output_structure(self) -> None:
        """Test SASS output has expected structure."""
        # SASS is assembly with register operations
        mock_sass = """\
        code for sm_90a
                Function : vector_add
        .headerflags    @"EF_CUDA_SM90 EF_CUDA_PTX_SM(8.0)"
        /*0000*/   MOV R1, c[0x0][0x28] ;
        /*0010*/   S2R R0, SR_CTAID.X ;
        /*0020*/   IMAD.U32 R0, R0, c[0x0][0x0], R0 ;
"""
        response = CompileResponse(
            success=True,
            sass=mock_sass,
            compilation_time_ms=100,
        )
        assert response.sass is not None
        assert "MOV" in response.sass or "code for" in response.sass

    def test_both_outputs(self) -> None:
        """Test response can have both PTX and SASS."""
        response = CompileResponse(
            success=True,
            ptx=".version 8.0",
            sass="MOV R0, R1",
            compilation_time_ms=150,
        )
        assert response.ptx is not None
        assert response.sass is not None


class TestArchitectures:
    """Test different GPU architecture specifications."""

    def test_sm_90a_hopper(self) -> None:
        """Test Hopper architecture request."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            arch="sm_90a",
        )
        assert request.arch == "sm_90a"

    def test_sm_89_ada(self) -> None:
        """Test Ada Lovelace architecture request."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            arch="sm_89",
        )
        assert request.arch == "sm_89"

    def test_sm_80_ampere(self) -> None:
        """Test Ampere architecture request."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            arch="sm_80",
        )
        assert request.arch == "sm_80"

    def test_sm_100_blackwell(self) -> None:
        """Test Blackwell architecture request."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            arch="sm_100",
        )
        assert request.arch == "sm_100"


class TestCompilerFlags:
    """Test custom nvcc compiler flags."""

    def test_optimization_O3(self) -> None:
        """Test O3 optimization flag."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            flags=("-O3",),
        )
        assert "-O3" in request.flags

    def test_lineinfo_flag(self) -> None:
        """Test lineinfo debug flag."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            flags=("-lineinfo",),
        )
        assert "-lineinfo" in request.flags

    def test_maxrregcount(self) -> None:
        """Test maxrregcount register limit flag."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            flags=("--maxrregcount=64",),
        )
        assert "--maxrregcount=64" in request.flags

    def test_multiple_flags(self) -> None:
        """Test multiple compiler flags."""
        request = CompileRequest(
            files={"kernel.cu": "__global__ void test() {}"},
            flags=("-O3", "-lineinfo", "--maxrregcount=64"),
        )
        assert len(request.flags) == 3


class TestErrors:
    """Test compilation error reporting."""

    def test_syntax_error_reported(self) -> None:
        """Test syntax error is included in response."""
        response = CompileResponse.error(
            "kernel.cu(10): error: expected a ';'"
        )
        assert response.success is False
        assert "expected a ';'" in response.stderr

    def test_missing_include_error(self) -> None:
        """Test missing include error."""
        response = CompileResponse.error(
            "kernel.cu(1): fatal error: cannot open include file: missing.h"
        )
        assert response.success is False
        assert "cannot open include file" in response.stderr

    def test_undefined_symbol_error(self) -> None:
        """Test undefined symbol error."""
        response = CompileResponse.error(
            "kernel.cu(5): error: identifier 'undefined_func' is undefined"
        )
        assert response.success is False
        assert "undefined" in response.stderr

    def test_error_line_number_included(self) -> None:
        """Test that error messages include line numbers."""
        response = CompileResponse.error(
            "kernel.cu(42): error: some error message"
        )
        assert "(42)" in response.stderr
