"""CUDA compiler wrapper.
This module provides the core compilation logic that can be used locally
or through Modal. It handles request/response conversion and provides
a clean interface for the API and CLI.
"""
import subprocess
from dataclasses import asdict
from wafer_core.tools.compile.types import (
    CompileRequest,
    CompileResponse,
    OutputFormat,
)
def request_to_dict(request: CompileRequest) -> dict:
    """Convert a CompileRequest to a dict for Modal invocation.
    Args:
        request: The compile request
    Returns:
        Dict suitable for passing to Modal function
    """
    return {
        "files": dict(request.files),
        "arch": request.arch,
        "flags": list(request.flags),
        "output": [fmt.value for fmt in request.output],
    }
def response_from_dict(data: dict) -> CompileResponse:
    """Convert a Modal response dict to a CompileResponse.
    Args:
        data: Dict returned from Modal function
    Returns:
        CompileResponse object
    """
    return CompileResponse(
        success=data.get("success", False),
        ptx=data.get("ptx"),
        sass=data.get("sass"),
        stderr=data.get("stderr", ""),
        compilation_time_ms=data.get("compilation_time_ms", 0),
    )
async def compile_cuda_remote(
    request: CompileRequest,
    *,
    modal_token_id: str | None = None,
    modal_token_secret: str | None = None,
) -> CompileResponse:
    """Compile CUDA code using Modal (remote execution).
    This function calls the deployed Modal function directly using asyncio.to_thread
    to avoid blocking the event loop.
    Args:
        request: The compile request
        modal_token_id: Optional Modal token ID (uses env var if not provided)
        modal_token_secret: Optional Modal token secret
    Returns:
        CompileResponse with PTX/SASS or error
    """
    import asyncio
    import os
    import time
    from contextlib import contextmanager
    @contextmanager
    def temporary_env_vars(env_updates: dict[str, str]):
        """Context manager to temporarily set environment variables.
        Saves original values, sets new values, yields, then restores originals.
        This ensures we don't leak credentials between concurrent requests.
        """
        original_values: dict[str, str | None] = {}
        for key, value in env_updates.items():
            original_values[key] = os.environ.get(key)
            os.environ[key] = value
        try:
            yield
        finally:
            for key, original in original_values.items():
                if original is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = original
    start_time = time.time()
    request_dict = request_to_dict(request)

    env_updates: dict[str, str] = {}
    if modal_token_id:
        env_updates["MODAL_TOKEN_ID"] = modal_token_id
    if modal_token_secret:
        env_updates["MODAL_TOKEN_SECRET"] = modal_token_secret
    def call_modal() -> dict:
        """Call Modal function synchronously (runs in thread pool)."""
        import modal
        # Look up the deployed function
        compile_fn = modal.Function.from_name("cuda-compile", "compile_cuda")
        # Call the function remotely
        return compile_fn.remote(request_dict)
    try:
        # Run Modal call in thread pool with temporary credentials
        # The context manager ensures env vars are restored after the call
        def call_modal_with_env() -> dict:
            with temporary_env_vars(env_updates):
                return call_modal()
        result = await asyncio.to_thread(call_modal_with_env)
        return response_from_dict(result)
    except Exception as e:
        error_str = str(e)

        if "MODAL_TOKEN" in error_str or "AuthError" in error_str or "not authenticated" in error_str.lower():
            return CompileResponse.error(
                "Modal not configured. Set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET environment variables, "
                "or run 'modal token new' to authenticate.",
                compilation_time_ms=int((time.time() - start_time) * 1000),
            )
        return CompileResponse.error(
            f"Compilation failed: {error_str}",
            compilation_time_ms=int((time.time() - start_time) * 1000),
        )
def compile_cuda_local(request: CompileRequest) -> CompileResponse:
    """Compile CUDA code locally using nvcc.
    This function requires nvcc to be installed locally.
    Primarily useful for testing without Modal.
    Args:
        request: The compile request
    Returns:
        CompileResponse with PTX/SASS or error
    """
    import os
    import tempfile
    import time
    from pathlib import Path
    start_time = time.time()

    try:
        subprocess.run(
            ["nvcc", "--version"],
            capture_output=True,
            check=True,
            timeout=10,
        )
    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
        return CompileResponse.error(
            "nvcc not found. Install CUDA toolkit or use Modal for remote compilation.",
            compilation_time_ms=int((time.time() - start_time) * 1000),
        )
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_path = Path(tmpdir)
            # Write all files to temp directory
            for filename, content in request.files.items():
                file_path = (tmp_path / filename).resolve()
                if not file_path.is_relative_to(tmp_path):
                    return CompileResponse.error(
                        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 / request.main_cu_file
            include_dir = main_cu_path.parent
            results: dict[str, str | None] = {"ptx": None, "sass": None}

            base_cmd = [
                "nvcc",
                "-arch",
                request.arch,
                f"-I{include_dir}",
            ]
            base_cmd.extend(request.flags)
            # Generate PTX if requested
            if OutputFormat.PTX in request.output:
                ptx_output = tmp_path / "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,
                )
                if ptx_result.returncode != 0:
                    return CompileResponse.error(
                        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 in request.output:
                cubin_output = tmp_path / "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,
                )
                if cubin_result.returncode != 0:
                    if results["ptx"]:
                        return CompileResponse(
                            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 CompileResponse.error(
                        cubin_result.stderr or cubin_result.stdout,
                        compilation_time_ms=int((time.time() - start_time) * 1000),
                    )
                if cubin_output.exists():
                    sass_result = subprocess.run(
                        ["cuobjdump", "--dump-sass", str(cubin_output)],
                        capture_output=True,
                        text=True,
                        timeout=30,
                        cwd=tmpdir,
                    )
                    if sass_result.returncode == 0:
                        results["sass"] = sass_result.stdout
            if not results["ptx"] and not results["sass"]:
                return CompileResponse.error(
                    "No output generated",
                    compilation_time_ms=int((time.time() - start_time) * 1000),
                )
            return CompileResponse(
                success=True,
                ptx=results["ptx"],
                sass=results["sass"],
                stderr="",
                compilation_time_ms=int((time.time() - start_time) * 1000),
            )
    except subprocess.TimeoutExpired:
        return CompileResponse.error(
            "Compilation timed out",
            compilation_time_ms=int((time.time() - start_time) * 1000),
        )
    except Exception as e:
        import traceback
        return CompileResponse.error(
            f"Internal error: {e}\n{traceback.format_exc()}",
            compilation_time_ms=int((time.time() - start_time) * 1000),
        )
