"""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 spawns a subprocess to call Modal, avoiding event loop
    conflicts between the caller's event loop and Modal's asyncio.

    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 json
    import os
    import tempfile
    import time
    from pathlib import Path

    start_time = time.time()

    # Write request to temp file
    request_dict = request_to_dict(request)

    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".json", delete=False
    ) as request_file:
        json.dump(request_dict, request_file)
        request_path = request_file.name

    try:
        # Create a Python script that calls Modal using Function.lookup
        # This calls the deployed function without needing to rebuild the image
        script = f'''
import json
import modal

# Load request
with open("{request_path}") as f:
    request = json.load(f)

# Look up the deployed function
compile_fn = modal.Function.from_name("cuda-compile", "compile_cuda")

# Call the function remotely
result = compile_fn.remote(request)

# Output result as JSON
print(json.dumps(result))
'''

        # Run in subprocess to avoid event loop conflicts
        env = os.environ.copy()
        if modal_token_id:
            env["MODAL_TOKEN_ID"] = modal_token_id
        if modal_token_secret:
            env["MODAL_TOKEN_SECRET"] = modal_token_secret

        # Use the same Python interpreter that's running this code
        import sys
        python_executable = sys.executable

        # Use asyncio.create_subprocess_exec for async subprocess execution
        proc = await asyncio.create_subprocess_exec(
            python_executable, "-c", script,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )
        stdout_bytes, stderr_bytes = await proc.communicate()

        if proc.returncode != 0:
            stderr = stderr_bytes.decode() if stderr_bytes else "Unknown error"
            # Check for common Modal auth errors
            if "MODAL_TOKEN" in stderr or "AuthError" in stderr or "not authenticated" in stderr.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: {stderr}",
                compilation_time_ms=int((time.time() - start_time) * 1000),
            )

        # Parse result
        stdout = stdout_bytes.decode() if stdout_bytes else "{}"
        try:
            response_dict = json.loads(stdout)
            return response_from_dict(response_dict)
        except json.JSONDecodeError as e:
            return CompileResponse.error(
                f"Failed to parse Modal response: {e}\nOutput: {stdout[:500]}",
                compilation_time_ms=int((time.time() - start_time) * 1000),
            )

    finally:
        # Clean up temp file
        Path(request_path).unlink(missing_ok=True)


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()

    # Check if nvcc is available
    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
                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}

            # Build base nvcc command
            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),
        )
