"""Modal app for serverless kernel evaluation.
This module defines the Modal app that runs kernel evaluation on serverless GPUs.
It reuses the same evaluate.py logic as the SSH backend.
Tiger Style:
- Reuse existing evaluation logic (DRY)
- Explicit dependencies in Modal image
- Clear function signatures
"""
from pathlib import Path
import modal
# ══════════════════════════════════════════════════════════════════════════════
# Modal Image Configuration
# ══════════════════════════════════════════════════════════════════════════════
# This image is cached and reused across function invocations
def build_modal_image() -> modal.Image:
    """Build Modal image with PyTorch, CUTLASS, and evaluation dependencies.
    Uses explicit local code inclusion to avoid pulling in SSH deployment code.
    Returns:
        Modal Image ready for kernel evaluation
    """
    torch_index = "https://download.pytorch.org/whl/cu130"
    torch_version = "torch>=2.6.0"
    image = (
        modal.Image.debian_slim(python_version="3.10")
        # System dependencies
        .apt_install("git", "build-essential")
        # PyTorch (GPU-specific index)
        .pip_install(
            torch_version,
            index_url=torch_index,
            extra_index_url="https://pypi.org/simple",
        )
        # CUDA toolkit dependencies
        .pip_install(
            "triton",
            "ninja",
            index_url="https://pypi.nvidia.com",
            extra_index_url="https://pypi.org/simple",
        )
        # Evaluation dependencies
        .pip_install(
            "numpy",
            "scipy",
            "pytest",
        )
        # Install CUTLASS headers for C++ kernel compilation (v4.3.5)
        .run_commands(
            "git clone --depth 1 --branch v4.3.5 https://github.com/NVIDIA/cutlass.git /usr/local/cutlass",
            # Verify CUTLASS was installed correctly
            "ls -la /usr/local/cutlass/include/cutlass/util/ | head -20",
            "test -f /usr/local/cutlass/include/cutlass/util/packed_stride.hpp && echo 'CUTLASS headers OK' || echo 'CUTLASS headers MISSING'",
        )

        .env({"CUTLASS_PATH": "/usr/local/cutlass/include"})
        # MUST run before add_local_* commands (Modal restriction)
        .run_commands(
            "mkdir -p /root/wafer_core/utils",
            "touch /root/wafer_core/__init__.py",
            "touch /root/wafer_core/utils/__init__.py",
        )
        # We only need the kernel_utils/ subdirectory + exceptions.py + ncu_profile_tools.py
        .add_local_dir(
            local_path=Path(__file__).parent.parent
            / "kernel_utils",  # wafer_core/utils/kernel_utils/
            remote_path="/root/wafer_core/utils/kernel_utils",
        )
        .add_local_file(
            local_path=Path(__file__).parent.parent
            / "exceptions.py",  # wafer_core/utils/exceptions.py
            remote_path="/root/wafer_core/utils/exceptions.py",
        )
        .add_local_file(
            local_path=Path(__file__).parent.parent
            / "ncu_profile_tools.py",  # wafer_core/utils/ncu_profile_tools.py
            remote_path="/root/wafer_core/utils/ncu_profile_tools.py",
        )
    )
    return image
# ══════════════════════════════════════════════════════════════════════════════
# Modal App Definition
# ══════════════════════════════════════════════════════════════════════════════
def create_modal_app(
    app_name: str = "test-kernel-eval",  # Match test script default
) -> modal.App:
    """Create Modal app for kernel evaluation.
    Args:
        app_name: Modal app name
    Returns:
        Modal App instance
    """
    image = build_modal_image()
    return modal.App(name=app_name, image=image)
# Default app (can be overridden)
app = create_modal_app()
# ══════════════════════════════════════════════════════════════════════════════
# Modal Function - Kernel Evaluation
# ══════════════════════════════════════════════════════════════════════════════
@app.function(
    gpu="B200",  # Can be overridden per invocation
    timeout=600,  # 10 minute timeout
    cpu=4,  # CPUs for kernel compilation
    memory=16384,  # 16GB RAM
)
def run_kernel_evaluation(
    kernel_code: str,
    reference_code: str,
    problem_id: str,
    test_suite: str,
    benchmark_suite: str,
    reference_backend: str,
    test_cases: list[dict],
    benchmark_cases: list[dict],
    language: str = "pytorch",
    profile: bool = False,
) -> dict:
    """Run kernel evaluation on Modal GPU.
    This function reuses the same evaluation logic as the SSH backend.
    It creates a temporary workspace, writes the kernel and reference files,
    and runs evaluate.py.
    Args:
        kernel_code: Kernel implementation code
        reference_code: Reference implementation code
        problem_id: Problem identifier
        test_suite: Test suite name (e.g., "gpumode_correctness")
        benchmark_suite: Benchmark suite name (e.g., "gpumode_benchmark")
        reference_backend: Reference backend name (e.g., "reference")
        test_cases: List of test case dicts
        benchmark_cases: List of benchmark case dicts
        language: Language/framework (e.g., "pytorch", "cute")
        profile: Whether to run torch profiler
    Returns:
        Dict with evaluation results including per-test details:
        {
            "compiled": bool,
            "error_message": str | None,
            "correctness_score": float,
            "geomean_speedup": float,
            "all_correct": bool,
            "passed_tests": int,
            "total_tests": int,
            "correctness_tests": list[dict],  # Per-test results with errors
            "performance_tests": list[dict],  # Per-test benchmark results
        }
    """
    import tempfile
    from pathlib import Path
    import torch
    from wafer_core.utils.kernel_utils.evaluate import (
        EvaluationArtifactConfig,
        run_evaluation_full,
    )
    from wafer_core.utils.kernel_utils.task import TestCase
    try:
        # Verify GPU is available
        if not torch.cuda.is_available():
            return {
                "compiled": False,
                "error_message": "CUDA not available in Modal container",
                "correctness_score": 0.0,
                "geomean_speedup": 0.0,
                "all_correct": False,
                "passed_tests": 0,
                "total_tests": len(test_cases),
                "correctness_tests": [],
                "performance_tests": [],
            }
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_path = Path(tmpdir)
            # Write kernel code to file
            kernel_file = tmp_path / "kernel.py"
            kernel_file.write_text(kernel_code)
            # Write reference code to file with the expected name
            # TestCase.generate() imports from 'reference_kernel', not 'reference'
            reference_file = tmp_path / "reference_kernel.py"
            reference_file.write_text(reference_code)
            test_case_objects = [
                TestCase(params=tc, name=f"test_{i + 1}") for i, tc in enumerate(test_cases)
            ]
            # Configure evaluation (no artifacts, no profiling for Modal)
            artifact_config = EvaluationArtifactConfig(
                run_profiling=profile,
                create_artifact=False,  # Don't create artifacts in Modal
                run_dir=None,
            )
            # Change working directory to temp path so TestCase.generate() can import reference_kernel
            # TestCase.generate() does: from reference_kernel import generate_input
            import os
            original_cwd = os.getcwd()
            try:
                os.chdir(tmp_path)
                # Run evaluation using the full results API
                eval_result = run_evaluation_full(
                    implementation_path=str(kernel_file),
                    reference_path=str(reference_file),
                    test_cases=test_case_objects,
                    run_benchmarks=True,  # Always run benchmarks
                    artifact_config=artifact_config,
                )
            finally:
                # Restore original working directory
                os.chdir(original_cwd)
            result = eval_result.to_dict()
            result["gpu_info"] = {
                "available": True,
                "name": torch.cuda.get_device_name(0),
            }
            return result
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        return {
            "compiled": False,
            "error_message": f"{str(e)}\n\nTraceback:\n{error_trace}",
            "correctness_score": 0.0,
            "geomean_speedup": 0.0,
            "all_correct": False,
            "passed_tests": 0,
            "total_tests": len(test_cases),
            "correctness_tests": [],
            "performance_tests": [],
        }
# ══════════════════════════════════════════════════════════════════════════════

# ══════════════════════════════════════════════════════════════════════════════
@app.function()
def health_check() -> dict:
    """Simple health check to verify Modal app is working.
    Returns:
        Dict with status and GPU info
    """
    import torch
    return {
        "status": "ok",
        "cuda_available": torch.cuda.is_available(),
        "gpu_count": torch.cuda.device_count(),
        "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
    }
