"""Modal function for cloud CUDA compilation.
This module provides a CPU-only Modal function that compiles CUDA code
and returns PTX/SASS assembly. No GPU is required - nvcc can generate
PTX/SASS for any architecture without a physical GPU.
Usage:
    # From Python
    from wafer_core.tools.compile.modal_compile import compile_cuda
    result = compile_cuda.remote(request_dict)
    # Deploy Modal app
    modal deploy modal_compile.py
"""
from enum import Enum
import modal
# Inline OutputFormat to avoid import chain issues with Modal deployment
class OutputFormat(str, Enum):
    """Supported output formats."""
    PTX = "ptx"
    SASS = "sass"
# ══════════════════════════════════════════════════════════════════════════════
# Modal Image Configuration
# ══════════════════════════════════════════════════════════════════════════════
# No GPU driver needed - nvcc can generate PTX/SASS without a GPU
compile_image = (
    # Start with NVIDIA's CUDA development image (includes nvcc)
    modal.Image.from_registry(
        "nvidia/cuda:13.0.1-devel-ubuntu22.04",
        add_python="3.12",
    )
    # System dependencies
    .apt_install("git", "build-essential", "wget")
    # Install PyTorch headers (CPU-only, we just need the headers)
    .pip_install(
        "torch>=2.5.0",
        index_url="https://download.pytorch.org/whl/cpu",
        extra_index_url="https://pypi.org/simple",
    )
    # Install CUTLASS headers (v4.3.5)
    .run_commands(
        "git clone --depth 1 --branch v4.3.5 https://github.com/NVIDIA/cutlass.git /usr/local/cutlass",
    )

    .env(
        {
            "CUDA_HOME": "/usr/local/cuda",
            "PATH": "/usr/local/cuda/bin:$PATH",
            "CUTLASS_PATH": "/usr/local/cutlass/include",
        }
    )
    # Verify nvcc is working
    .run_commands(
        "nvcc --version",
        "ls -la /usr/local/cuda/bin/nvcc",
    )
)
app = modal.App(name="cuda-compile", image=compile_image)
# ══════════════════════════════════════════════════════════════════════════════
# Modal Function - CUDA Compilation
# ══════════════════════════════════════════════════════════════════════════════
@app.function(
    # CPU only - no GPU needed for compilation!
    cpu=4,
    memory=8192,  # 8GB RAM
    timeout=120,  # 2 minute timeout
    # Keep one container warm to avoid cold starts (~5-10s savings)
    min_containers=1,
)
@modal.concurrent(max_inputs=4)  # Allow concurrent compilations for better throughput
def compile_cuda(request: dict) -> dict:
    """Compile CUDA code and return PTX/SASS.
    This function runs on a CPU-only container with nvcc installed.
    No GPU is required because nvcc can cross-compile for any target architecture.
    Args:
        request: Dict with the following fields:
            - files: dict[str, str] - Mapping of filename to content
            - arch: str - Target architecture (e.g., "sm_90a")
            - flags: list[str] - Additional nvcc flags
            - output: list[str] - Output formats ("ptx", "sass")
    Returns:
        Dict with the following fields:
            - success: bool - Whether compilation succeeded
            - ptx: str | None - Generated PTX code
            - sass: str | None - Generated SASS code
            - stderr: str - Compiler warnings/errors
            - compilation_time_ms: int - Time taken in ms
    """
    import os
    import subprocess
    import tempfile
    import time
    from concurrent.futures import ThreadPoolExecutor, as_completed
    from pathlib import Path
    start_time = time.time()
    files: dict[str, str] = request.get("files", {})
    arch: str = request.get("arch", "sm_90a")
    flags: list[str] = request.get("flags", [])
    output_formats: list[str] = request.get("output", ["ptx", "sass"])
    # Validate request
    if not files:
        return {
            "success": False,
            "ptx": None,
            "sass": None,
            "stderr": "No files provided",
            "compilation_time_ms": int((time.time() - start_time) * 1000),
        }

    cu_files = [f for f in files.keys() if f.endswith(".cu")]
    if not cu_files:
        return {
            "success": False,
            "ptx": None,
            "sass": None,
            "stderr": "No .cu file found",
            "compilation_time_ms": int((time.time() - start_time) * 1000),
        }
    main_cu_file = cu_files[0]
    nvcc_env = {
        **os.environ,
        "CUDA_HOME": "/usr/local/cuda",
        "PATH": f"/usr/local/cuda/bin:{os.environ.get('PATH', '')}",
    }
    def compile_ptx(tmpdir: str, base_cmd: list[str], main_cu_path: Path) -> tuple[str | None, str | None]:
        """Compile to PTX. Returns (ptx_content, error_message)."""
        ptx_output = Path(tmpdir) / "output.ptx"
        ptx_cmd = base_cmd + [
            "--ptx",
            "-o",
            str(ptx_output),
            str(main_cu_path),
        ]
        ptx_result = subprocess.run(
            ptx_cmd,
            capture_output=True,
            text=True,
            timeout=60,
            cwd=tmpdir,
            env=nvcc_env,
        )
        if ptx_result.returncode != 0:
            return None, ptx_result.stderr or ptx_result.stdout
        if ptx_output.exists():
            return ptx_output.read_text(), None
        return None, "PTX output file not created"
    def compile_sass(tmpdir: str, base_cmd: list[str], main_cu_path: Path) -> tuple[str | None, str | None]:
        """Compile to SASS (via cubin). Returns (sass_content, error_message)."""
        cubin_output = Path(tmpdir) / "output.cubin"
        cubin_cmd = base_cmd + [
            "--cubin",
            "-o",
            str(cubin_output),
            str(main_cu_path),
        ]
        cubin_result = subprocess.run(
            cubin_cmd,
            capture_output=True,
            text=True,
            timeout=60,
            cwd=tmpdir,
            env=nvcc_env,
        )
        if cubin_result.returncode != 0:
            return None, cubin_result.stderr or cubin_result.stdout
        if not cubin_output.exists():
            return None, "cubin output file not created"
        # Disassemble cubin to SASS
        sass_result = subprocess.run(
            ["cuobjdump", "--dump-sass", str(cubin_output)],
            capture_output=True,
            text=True,
            timeout=30,
            cwd=tmpdir,
        )
        if sass_result.returncode == 0:
            return sass_result.stdout, None
        return None, f"SASS disassembly failed: {sass_result.stderr}"
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_path = Path(tmpdir)
            # Write all files to temp directory, preserving subdirectory structure
            for filename, content in files.items():
                file_path = (tmp_path / filename).resolve()
                if not file_path.is_relative_to(tmp_path):
                    return {
                        "success": False,
                        "ptx": None,
                        "sass": None,
                        "stderr": f"Invalid filename: {filename}",
                        "compilation_time_ms": int((time.time() - start_time) * 1000),
                    }
                file_path.parent.mkdir(parents=True, exist_ok=True)
                file_path.write_text(content)

            main_cu_path = tmp_path / main_cu_file
            include_dir = main_cu_path.parent
            base_cmd = [
                "nvcc",
                "-arch",
                arch,
                f"-I{include_dir}",
                "-I/usr/local/lib/python3.12/site-packages/torch/include",
                "-I/usr/local/lib/python3.12/site-packages/torch/include/torch/csrc/api/include",
                "-I/usr/local/cutlass/include",
            ]
            base_cmd.extend(flags)

            want_ptx = OutputFormat.PTX.value in output_formats
            want_sass = OutputFormat.SASS.value in output_formats
            results: dict[str, str | None] = {"ptx": None, "sass": None}
            errors: list[str] = []
            # Run compilations in parallel if both are requested
            if want_ptx and want_sass:
                with ThreadPoolExecutor(max_workers=2) as executor:
                    futures = {
                        executor.submit(compile_ptx, tmpdir, base_cmd, main_cu_path): "ptx",
                        executor.submit(compile_sass, tmpdir, base_cmd, main_cu_path): "sass",
                    }
                    for future in as_completed(futures):
                        output_type = futures[future]
                        try:
                            content, error = future.result()
                            if content:
                                results[output_type] = content
                            if error:
                                errors.append(f"{output_type.upper()}: {error}")
                        except Exception as e:
                            errors.append(f"{output_type.upper()} compilation error: {e}")
            elif want_ptx:
                content, error = compile_ptx(tmpdir, base_cmd, main_cu_path)
                if content:
                    results["ptx"] = content
                if error:
                    errors.append(error)
            elif want_sass:
                content, error = compile_sass(tmpdir, base_cmd, main_cu_path)
                if content:
                    results["sass"] = content
                if error:
                    errors.append(error)
            if not results["ptx"] and not results["sass"]:
                return {
                    "success": False,
                    "ptx": None,
                    "sass": None,
                    "stderr": "\n".join(errors) if errors else "No output generated",
                    "compilation_time_ms": int((time.time() - start_time) * 1000),
                }
            # Partial success if we got at least one output
            stderr = ""
            if errors and (results["ptx"] or results["sass"]):
                stderr = "\n".join(errors)
            return {
                "success": True,
                "ptx": results["ptx"],
                "sass": results["sass"],
                "stderr": stderr,
                "compilation_time_ms": int((time.time() - start_time) * 1000),
            }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "ptx": None,
            "sass": None,
            "stderr": "Compilation timed out",
            "compilation_time_ms": int((time.time() - start_time) * 1000),
        }
    except Exception as e:
        import traceback
        return {
            "success": False,
            "ptx": None,
            "sass": None,
            "stderr": f"Internal error: {e}\n{traceback.format_exc()}",
            "compilation_time_ms": int((time.time() - start_time) * 1000),
        }
# ══════════════════════════════════════════════════════════════════════════════
# Health Check
# ══════════════════════════════════════════════════════════════════════════════
@app.function(cpu=1, memory=1024, timeout=30)
def health_check() -> dict:
    """Verify the compilation environment is working.
    Returns:
        Dict with status and version info
    """
    import subprocess
    try:
        nvcc_result = subprocess.run(
            ["nvcc", "--version"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        version_line = nvcc_result.stdout.strip().split("\n")[-1]
        return {
            "status": "ok",
            "nvcc_version": version_line,
            "nvcc_available": nvcc_result.returncode == 0,
        }
    except Exception as e:
        return {
            "status": "error",
            "error": str(e),
            "nvcc_available": False,
        }
