from __future__ import annotations

import functools
import linecache
import os
import sys
import time
import warnings
from pathlib import Path
from types import ModuleType
from typing import Any, TYPE_CHECKING

from torch._utils_internal import log_triton_builds


if TYPE_CHECKING:
    from collections.abc import Callable

    from torch._inductor.runtime.triton_heuristics import CachingAutotuner


def _reload_python_module(
    key: str, path: str, set_sys_modules: bool = True
) -> ModuleType:
    with open(path) as f:
        try:
            code = compile(f.read(), path, "exec", dont_inherit=True)
        except Exception as e:
            raise RuntimeError(
                f"Failed to import {path}\n{type(e).__name__}: {e}"
            ) from None
        mod = ModuleType(f"{__name__}.{key}")
        mod.__file__ = path
        mod.key = key  # type: ignore[attr-defined]
        exec(code, mod.__dict__, mod.__dict__)
        if set_sys_modules:
            sys.modules[mod.__name__] = mod
        return mod


@functools.cache
def _set_triton_ptxas_path() -> None:
    if os.environ.get("TRITON_PTXAS_PATH") is not None:
        return
    ptxas = Path(__file__).absolute().parents[1] / "bin" / "ptxas"
    if not ptxas.exists():
        return
    if ptxas.is_file() and os.access(ptxas, os.X_OK):
        os.environ["TRITON_PTXAS_PATH"] = str(ptxas)
    else:
        warnings.warn(f"{ptxas} exists but is not an executable")


def _set_triton_libdevice_path() -> None:
    """
    Use the CUDA toolkit's libdevice instead of Triton's bundled version.
    This ensures Triton's libdevice calls match CUDA eager numerics for bitwise
    precision.  Gated by config.eager_numerics.use_pytorch_libdevice and by
    config.emulate_precision_casts, which also requests eager-like numerics.
    """
    from torch._inductor import config

    if not (
        config.eager_numerics.use_pytorch_libdevice or config.emulate_precision_casts
    ):
        return

    _set_triton_libdevice_path_impl()


def _set_triton_libdevice_path_impl() -> None:
    import torch

    if torch.version.cuda is None:
        return

    try:
        from triton import knobs
    except ImportError:
        return

    env_path = os.environ.get("TRITON_LIBDEVICE_PATH")
    if env_path is not None:
        knobs.nvidia.libdevice_path = env_path
        return

    if knobs.nvidia.libdevice_path is not None:
        return

    try:
        from torch.utils.cpp_extension import CUDA_HOME

        if CUDA_HOME is None:
            warnings.warn(
                "CUDA_HOME not set; using Triton's bundled libdevice which may "
                "cause minor precision differences in pow operations. "
                "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
                "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
                stacklevel=3,
            )
            return
        libdevice = Path(CUDA_HOME) / "nvvm" / "libdevice" / "libdevice.10.bc"
        if libdevice.is_file():
            knobs.nvidia.libdevice_path = str(libdevice)
            # Also set env var so subprocess compile workers inherit it
            os.environ["TRITON_LIBDEVICE_PATH"] = str(libdevice)
        else:
            warnings.warn(
                f"CUDA libdevice not found at {libdevice}; using Triton's bundled "
                "libdevice which may cause minor precision differences in pow operations. "
                "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
                "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
                stacklevel=3,
            )
    except ImportError:
        warnings.warn(
            "torch.utils.cpp_extension not available; using Triton's bundled "
            "libdevice which may cause minor precision differences in pow operations. "
            "To fix: set TRITON_LIBDEVICE_PATH to your CUDA toolkit's libdevice, "
            "e.g., export TRITON_LIBDEVICE_PATH=/usr/local/cuda/nvvm/libdevice/libdevice.10.bc",
            stacklevel=3,
        )


def _worker_compile_pycodecache_kernel(
    kernel_name: str,
    source_code: str,
    main_suffix: str,
    extra_env: dict[str, str],
    precompile_metadata: dict[str, Any] | None = None,
) -> tuple[str, str, int]:
    """
    Subprocess worker for PyCodeCache-based kernel compilation.

    Writes source to PyCodeCache, loads the module, validates the entry point,
    and optionally triggers real GPU compilation (MLIR -> PTX -> CUBIN) via a
    _precompile entry point. Compiled artifacts are persisted to disk cache so
    the parent process can load them without recompilation.

    Used by both CuteDSL and NV Universal GEMM backends.
    """
    # No need to restore: this runs in a subprocess worker that exits after returning.
    os.environ.update(extra_env)

    start_ns = time.time_ns()

    import torch._inductor.codecache as codecache

    key, path = codecache.PyCodeCache.write(source_code)
    mod = codecache.PyCodeCache.load_by_key_path(key, path)

    main_func_name = f"{kernel_name}_{main_suffix}"
    if not hasattr(mod, main_func_name):
        available = [name for name in dir(mod) if callable(getattr(mod, name))]
        raise RuntimeError(
            f"Could not find kernel function '{main_func_name}'. "
            f"Available callables: {available}"
        )

    if precompile_metadata is not None:
        precompile_fn_name = f"{kernel_name}_precompile"
        precompile_fn = getattr(mod, precompile_fn_name, None)
        if precompile_fn is not None:
            precompile_fn(**precompile_metadata)
        else:
            import logging

            logging.getLogger(__name__).warning(
                "Precompile metadata was provided but module has no %s "
                "— the scheduling layer expected this template to support "
                "subprocess precompilation. Kernel will compile lazily on "
                "first call instead.",
                precompile_fn_name,
            )

    elapsed_ns = time.time_ns() - start_ns
    linecache.clearcache()
    return key, path, elapsed_ns // 1000


def _worker_compile_triton(
    load_kernel: Callable[[], CachingAutotuner],
    extra_env: dict[str, str],
    extra_config: dict[str, Any],
) -> tuple[CachingAutotuner, int]:
    _set_triton_ptxas_path()
    os.environ.update(extra_env)
    # Set libdevice path if passed via env from main process
    libdevice_path = extra_env.get("TRITON_LIBDEVICE_PATH")
    if libdevice_path:
        try:
            from triton import knobs

            knobs.nvidia.libdevice_path = libdevice_path
        except ImportError:
            pass
    from torch._inductor import config

    with config.patch(extra_config):
        fail = None
        try:
            start_ns = time.time_ns()
            kernel = load_kernel()
            kernel.precompile(warm_cache_only=True)
            elapsed_ns = time.time_ns() - start_ns
            kernel.prepare_for_pickle()
            # We can release this memory in the compile subprocesses:
            linecache.clearcache()
            return kernel, elapsed_ns // 1000
        except Exception as e:
            fail = str(e)
            raise
        finally:
            log_triton_builds(fail=fail)
