"""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
# ══════════════════════════════════════════════════════════════════════════════

# Build a CPU-only image with nvcc for CUDA compilation
# No GPU driver needed - nvcc can generate PTX/SASS without a GPU
compile_image = (
    # Start with NVIDIA's CUDA development image (includes nvcc)
    # Use debian-slim + manual CUDA install for smaller image
    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",
    )
    # Set environment variables for CUDA and libraries
    .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",
    )
)

# Create Modal app
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
)
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 pathlib import Path

    start_time = time.time()

    # Extract request fields
    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),
        }

    # Find the main .cu file
    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]

    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
                file_path.parent.mkdir(parents=True, exist_ok=True)
                file_path.write_text(content)

            # Determine the directory of the main .cu file for include paths
            main_cu_path = tmp_path / main_cu_file
            include_dir = main_cu_path.parent

            results: dict[str, str | None] = {"ptx": None, "sass": None}

            # Build base nvcc command with common flags
            base_cmd = [
                "nvcc",
                "-arch",
                arch,
                # Include the temp directory for user headers
                f"-I{include_dir}",
                # Include PyTorch headers
                "-I/usr/local/lib/python3.12/site-packages/torch/include",
                "-I/usr/local/lib/python3.12/site-packages/torch/include/torch/csrc/api/include",
                # Include CUTLASS headers
                "-I/usr/local/cutlass/include",
                # Standard CUDA headers are already in the default path
            ]

            # Add user-specified flags
            base_cmd.extend(flags)

            # Generate PTX if requested
            if OutputFormat.PTX.value in output_formats:
                ptx_output = tmp_path / "output.ptx"
                ptx_cmd = base_cmd + [
                    "--ptx",  # Generate 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={
                        **os.environ,
                        "CUDA_HOME": "/usr/local/cuda",
                        "PATH": f"/usr/local/cuda/bin:{os.environ.get('PATH', '')}",
                    },
                )

                if ptx_result.returncode != 0:
                    return {
                        "success": False,
                        "ptx": None,
                        "sass": None,
                        "stderr": ptx_result.stderr or ptx_result.stdout,
                        "compilation_time_ms": int((time.time() - start_time) * 1000),
                    }

                if ptx_output.exists():
                    results["ptx"] = ptx_output.read_text()

            # Generate SASS if requested
            if OutputFormat.SASS.value in output_formats:
                # First compile to cubin, then disassemble to SASS
                cubin_output = tmp_path / "output.cubin"
                cubin_cmd = base_cmd + [
                    "--cubin",  # Generate cubin (binary)
                    "-o",
                    str(cubin_output),
                    str(main_cu_path),
                ]

                cubin_result = subprocess.run(
                    cubin_cmd,
                    capture_output=True,
                    text=True,
                    timeout=60,
                    cwd=tmpdir,
                    env={
                        **os.environ,
                        "CUDA_HOME": "/usr/local/cuda",
                        "PATH": f"/usr/local/cuda/bin:{os.environ.get('PATH', '')}",
                    },
                )

                if cubin_result.returncode != 0:
                    # If we already have PTX, that's a partial success
                    if results["ptx"]:
                        return {
                            "success": True,
                            "ptx": results["ptx"],
                            "sass": None,
                            "stderr": f"SASS generation failed: {cubin_result.stderr}",
                            "compilation_time_ms": int(
                                (time.time() - start_time) * 1000
                            ),
                        }
                    return {
                        "success": False,
                        "ptx": None,
                        "sass": None,
                        "stderr": cubin_result.stderr or cubin_result.stdout,
                        "compilation_time_ms": int((time.time() - start_time) * 1000),
                    }

                # Disassemble cubin to SASS using cuobjdump
                if cubin_output.exists():
                    sass_cmd = [
                        "cuobjdump",
                        "--dump-sass",
                        str(cubin_output),
                    ]

                    sass_result = subprocess.run(
                        sass_cmd,
                        capture_output=True,
                        text=True,
                        timeout=30,
                        cwd=tmpdir,
                    )

                    if sass_result.returncode == 0:
                        results["sass"] = sass_result.stdout
                    else:
                        # SASS generation failed but we might have PTX
                        if results["ptx"]:
                            return {
                                "success": True,
                                "ptx": results["ptx"],
                                "sass": None,
                                "stderr": f"SASS disassembly failed: {sass_result.stderr}",
                                "compilation_time_ms": int(
                                    (time.time() - start_time) * 1000
                                ),
                            }

            # Check if we got any output
            if not results["ptx"] and not results["sass"]:
                return {
                    "success": False,
                    "ptx": None,
                    "sass": None,
                    "stderr": "No output generated",
                    "compilation_time_ms": int((time.time() - start_time) * 1000),
                }

            return {
                "success": True,
                "ptx": results["ptx"],
                "sass": results["sass"],
                "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:
        # Get nvcc version
        nvcc_result = subprocess.run(
            ["nvcc", "--version"],
            capture_output=True,
            text=True,
            timeout=10,
        )

        # Parse version from output
        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,
        }
