import os
import re
import subprocess
from pathlib import Path

import setuptools
import torch
from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension

ROOT = Path(__file__).resolve().parent

# cutlass headers for the blockwise GEMM kernels. Pinned to the upstream commit the
# build was validated against and fetched into a cache dir, rather than carried as a
# git submodule (keeps the public-repo sync clean and the clone CI-cacheable).
CUTLASS_REPO = "https://github.com/NVIDIA/cutlass.git"
CUTLASS_REF = "afa1772203677c5118fcd82537a9c8fefbcc7008"  # v3.8.0


def _nvidia_include_dirs() -> list[str]:
    """Include dirs from pip-installed nvidia packages (e.g. cusparse headers)."""
    try:
        import nvidia  # noqa: PLC0415

        return [str(p) for pkg in Path(nvidia.__path__[0]).iterdir() if (p := pkg / "include").is_dir()]
    except ImportError:
        return []


def _arch_tokens() -> list[str]:
    """Normalized entries from TORCH_CUDA_ARCH_LIST (e.g. ['8.9', '9.0'])."""
    raw = os.environ.get("TORCH_CUDA_ARCH_LIST", "")
    return [re.sub(r"\+PTX$", "", tok).strip() for tok in raw.replace(",", " ").split() if tok.strip()]


# Arch codes the blockwise FP8 GEMM kernels support, as nvcc `sm_<code>` targets.
# SM89 ("geforce") is the generic fp8 kernel: it runs on Ada and is also the kernel
# dispatched on Blackwell (sm_100a datacenter / sm_120 consumer), so it is compiled for
# those too. The SM90 ("deep_gemm") kernel is Hopper-only and needs sm_90a (wgmma/TMA);
# it is declared-always / defined-conditionally, so it stubs out on every non-Hopper
# pass and can share a multi-arch fat binary. Ampere has no fp8 path. Entries are
# filtered to what the local nvcc can actually target (see _nvcc_arch_nums), so the
# Blackwell codes are inert until built with CUDA 12.8+.
# NOTE: the Blackwell (100a/120) path is implemented but not yet validated on real
# Blackwell hardware -- needs a B200 + CUDA 12.8 build/run pass.
_BLOCKWISE_ARCHES = ["89", "90a", "100a", "120"]


def _nvcc_arch_nums() -> set[str]:
    """Architecture numbers this nvcc can target, e.g. {'80', '86', '89', '90'}."""
    try:
        out = subprocess.check_output([f"{CUDA_HOME}/bin/nvcc", "--list-gpu-arch"], text=True)
    except (OSError, subprocess.CalledProcessError):
        return set()
    return {m.group(1) for tok in out.split() if (m := re.match(r"compute_(\d+a?)$", tok.strip()))}


def _blockwise_gencode() -> tuple[list[str], bool]:
    """Return (``-gencode`` flags for blockwise_cpp, build_sm90).
    Honors ``TORCH_CUDA_ARCH_LIST`` when set (mapping 8.9 -> sm_89, 9.0/9.0a -> sm_90a,
    ignoring arches the kernels do not support); unset builds for every supported arch
    this nvcc can target. The flags apply uniformly to all sources -- safe because the
    SM90 source stubs itself out on non-sm_90a passes.
    """
    supported = _nvcc_arch_nums()
    # nvcc may report an arch either plain ("90") or suffixed ("90a"); accept either.
    base = [a for a in _BLOCKWISE_ARCHES if a in supported or a.rstrip("a") in supported]
    env = _arch_tokens()
    if env:
        sel = []
        for tok in env:
            if tok.startswith("8.9"):
                sel.append("89")
            elif tok.startswith("9.0"):
                sel.append("90a")
            elif tok.startswith("10.0"):
                sel.append("100a")
            elif tok.startswith("12.0"):
                sel.append("120")
            # Ampere (8.0/8.6) and other arches have no fp8 blockwise kernel.
        archs = [a for a in dict.fromkeys(sel) if a in base] or base
    else:
        archs = base
    flags = [f"-gencode=arch=compute_{a},code=sm_{a}" for a in archs]
    return flags, ("90a" in archs)


def _cutlass_include() -> str:
    """Return the cutlass include dir, fetching the pinned commit on first use.
    Honors ``CUTLASS_DIR`` (a prebuilt cutlass checkout, e.g. a system copy) and
    otherwise caches a shallow clone of ``CUTLASS_REF`` under
    ``LTX_KERNELS_CACHE_DIR`` (default ``~/.cache/ltx-kernels``), so it is reused
    across builds and can be restored from a CI cache. Keeps ``uv sync`` /
    ``pip install -e`` self-contained without a git submodule.
    """
    if env := os.environ.get("CUTLASS_DIR"):
        return str(Path(env) / "include")
    cache_root = Path(os.environ.get("LTX_KERNELS_CACHE_DIR", Path.home() / ".cache" / "ltx-kernels"))
    dest = cache_root / f"cutlass-{CUTLASS_REF}"
    if not (dest / "include" / "cutlass" / "cutlass.h").is_file():
        dest.mkdir(parents=True, exist_ok=True)
        # Blobless partial clone of the exact pinned commit (GitHub allows fetching an
        # arbitrary SHA), sparse-checked-out to include/ only: cutlass is header-only and
        # the rest of the repo (tools/test/examples/python, ~85% by size) is unused.
        subprocess.run(["git", "init", "-q", str(dest)], check=True)
        subprocess.run(["git", "-C", str(dest), "remote", "add", "origin", CUTLASS_REPO], check=True)
        # Cone-mode sparse checkout of just include/. Use init + set (not "set --cone",
        # whose inline flag postdates git 2.35 and is silently parsed as a pattern on older git).
        subprocess.run(["git", "-C", str(dest), "sparse-checkout", "init", "--cone"], check=True)
        subprocess.run(["git", "-C", str(dest), "sparse-checkout", "set", "include"], check=True)
        subprocess.run(
            ["git", "-C", str(dest), "fetch", "-q", "--depth", "1", "--filter=blob:none", "origin", CUTLASS_REF],
            check=True,
        )
        subprocess.run(["git", "-C", str(dest), "checkout", "-q", "FETCH_HEAD"], check=True)
    return str(dest / "include")


if __name__ == "__main__":
    if CUDA_HOME is None:
        raise RuntimeError(
            "CUDA toolkit not found (CUDA_HOME is None). ltx-kernels compiles CUDA extensions "
            "and must be built on a host with the CUDA toolkit installed (nvcc on PATH or "
            "CUDA_HOME set)."
        )

    ext_modules = []

    # all2all_cpp -- unchanged.
    all2all_args = ["-O3", "-Wall", "-Wextra", "-Werror", "-Wno-unused-parameter", "-Wno-attributes"]
    ext_modules.append(
        CUDAExtension(
            name="all2all_cpp",
            include_dirs=[str(ROOT / "csrc/all2all"), str(ROOT / "csrc/include"), *_nvidia_include_dirs()],
            sources=[
                "csrc/all2all/all2all.cpp",
                "csrc/all2all/cuda/all2all_heads.cu",
                "csrc/all2all/cuda/allgather.cu",
            ],
            extra_compile_args={"cxx": all2all_args, "nvcc": ["-O3"]},
        )
    )

    # ops_cpp -- arch-independent element ops (rms_norm_rope, rms_norm_split_rope,
    # fp6 pack/unpack). Arch is driven by TORCH_CUDA_ARCH_LIST / torch defaults.
    ext_modules.append(
        CUDAExtension(
            name="ops_cpp",
            sources=[
                "csrc/ops/ops_api.cpp",
                "csrc/ops/fp6_bitpack.cpp",
                "csrc/ops/fp6_pack.cu",
                "csrc/ops/rms_norm_rope.cpp",
                "csrc/ops/rms_norm_rope_cuda.cu",
                "csrc/ops/rms_norm_split_rope.cpp",
                "csrc/ops/rms_norm_split_rope_cuda.cu",
            ],
            include_dirs=[str(ROOT / "csrc/ops/include"), *_nvidia_include_dirs()],
            extra_compile_args={
                "cxx": ["-O3", "-std=c++17"],
                "nvcc": [
                    "-O3",
                    "-std=c++17",
                    "-U__CUDA_NO_HALF_OPERATORS__",
                    "-U__CUDA_NO_HALF_CONVERSIONS__",
                    "-U__CUDA_NO_HALF2_OPERATORS__",
                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                    "--expt-relaxed-constexpr",
                    "--expt-extended-lambda",
                    "--use_fast_math",
                ],
            },
        )
    )

    # blockwise_cpp -- FP8 GEMM. The SM89 (GeForce) kernel is always built; the SM90
    # (deep_gemm) kernel + -D__SM90__ are added whenever sm_90a is among the targets.
    # Arches are an explicit -gencode list (see _blockwise_gencode): TORCH_CUDA_ARCH_LIST
    # when set, else every supported arch this nvcc can target ("build for everything").
    # The list is uniform across sources -- the SM90 source declares-always /
    # defines-conditionally, so it compiles (as a stub) for non-sm_90a arches too. Note
    # blockwise is unsupported on Ampere and fails at *runtime* there, by design.
    cutlass_include = _cutlass_include()
    gencode, build_sm90 = _blockwise_gencode()
    blockwise_sources = [
        "csrc/blockwise/api.cpp",
        "csrc/blockwise/kernels/geforce/gemm.cu",
    ]
    abi = f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}"
    blockwise_cxx = ["-O3", "-std=c++17", "-fPIC", "-Wno-psabi", "-Wno-deprecated-declarations", abi]
    blockwise_nvcc = [
        "-O3",
        "-std=c++17",
        "--ptxas-options=-O2",
        "--expt-relaxed-constexpr",
        "--expt-extended-lambda",
        "-U__CUDA_NO_HALF_OPERATORS__",
        "-U__CUDA_NO_HALF_CONVERSIONS__",
        "-U__CUDA_NO_HALF2_OPERATORS__",
        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
        *gencode,
    ]
    if build_sm90:
        blockwise_sources.append("csrc/blockwise/kernels/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d_bias.cu")
        blockwise_cxx.append("-D__SM90__")
        blockwise_nvcc.append("-D__SM90__")

    ext_modules.append(
        CUDAExtension(
            name="blockwise_cpp",
            sources=blockwise_sources,
            include_dirs=[
                f"{CUDA_HOME}/include",
                f"{CUDA_HOME}/include/cccl",
                str(ROOT / "csrc/blockwise"),
                str(ROOT / "csrc/blockwise/kernels/deep_gemm/include"),
                cutlass_include,
                *_nvidia_include_dirs(),
            ],
            libraries=["cuda", "cudart", "nvrtc"],
            library_dirs=[f"{CUDA_HOME}/lib64", f"{CUDA_HOME}/lib64/stubs"],
            extra_compile_args={"cxx": blockwise_cxx, "nvcc": blockwise_nvcc},
        )
    )

    setuptools.setup(
        ext_modules=ext_modules,
        cmdclass={"build_ext": BuildExtension},
    )