#!/bin/bash
set -ex

# Simple build script wrapping uv/pip
# Usage:
#   ./build.sh  # local build (torch-based arch detection, TK only on SM90)
#   ./build.sh --wheel-dir /tmp/wheels  # build + install a reusable wheel
#   ./build.sh --rocm --python-only     # Python + Triton package, no compiled extension
# Environment overrides (if set, they win over auto-detection):
#   TORCH_CUDA_ARCH_LIST
#   CMAKE_ARGS (for FASTVIDEO_KERNEL_BUILD_TK / CMAKE_CUDA_ARCHITECTURES / GPU_BACKEND /
#               FASTVIDEO_KERNEL_BUILD_EXTENSION)

echo "Building fastvideo-kernel..."

# ---------------------------------------------------------------------------
# Neutralise conda-injected compiler toolchains.
#
# Conda compiler packages (gcc_linux-aarch64, gxx_linux-64, etc.) set
# CMAKE_ARGS, CFLAGS, CXXFLAGS, and LDFLAGS on activation.  When multiple
# toolchains are installed the variables can reference a *cross*-compiler
# that doesn't match the host (e.g. aarch64-conda-linux-gnu-c++ on x86_64).
# Even when the correct toolchain is active, the flags it injects
# (-march=nocona, -mtune=haswell, …) can conflict with nvcc's host-compiler
# expectations.  Clear them so CMake discovers the system compiler instead.
# ---------------------------------------------------------------------------
if [[ -n "${CONDA_PREFIX:-}" ]]; then
    _need_clean=0
    # Detect conda cross-compiler that doesn't match the host.
    _host_arch="$(uname -m)"
    if [[ "${CXX:-}" == *"conda"* ]] || [[ "${CC:-}" == *"conda"* ]]; then
        _need_clean=1
    fi
    if [[ "${CMAKE_ARGS:-}" == *"conda"* ]]; then
        _need_clean=1
    fi
    if (( _need_clean )); then
        echo "NOTE: Clearing conda-injected compiler settings (CC/CXX/CMAKE_ARGS/CFLAGS/...)"
        echo "      to use the system compiler for CUDA extension builds."
        unset CC CXX CMAKE_ARGS CFLAGS CXXFLAGS LDFLAGS
    fi
    unset _need_clean _host_arch
fi

GPU_BACKEND=CUDA
WHEEL_DIR=""
PYTHON_ONLY=0
while [[ $# -gt 0 ]]; do
    case "$1" in
        --rocm)
            GPU_BACKEND=ROCM
            shift
            ;;
        --python-only)
            PYTHON_ONLY=1
            shift
            ;;
        --wheel-dir)
            if [[ $# -lt 2 ]]; then
                echo "ERROR: --wheel-dir requires a directory path" >&2
                exit 2
            fi
            WHEEL_DIR="$2"
            shift 2
            ;;
        *)
            echo "ERROR: unknown argument: $1" >&2
            exit 2
            ;;
    esac
done

has_cmake_arg() {
    local key="$1"
    [[ "${CMAKE_ARGS:-}" =~ (^|[[:space:]])-D${key}(=|$) ]]
}

extension_mode_from_cmake_args() {
    [[ "${CMAKE_ARGS:-}" =~ (^|[[:space:]])-DFASTVIDEO_KERNEL_BUILD_EXTENSION=([^[:space:]]*) ]] \
        && printf '%s' "${BASH_REMATCH[2]^^}"
}
if (( PYTHON_ONLY )); then
    _explicit_mode="$(extension_mode_from_cmake_args || true)"
    if [[ -n "${_explicit_mode}" && "${_explicit_mode}" != "OFF" ]]; then
        echo "ERROR: --python-only conflicts with CMAKE_ARGS -DFASTVIDEO_KERNEL_BUILD_EXTENSION=${_explicit_mode}" >&2
        exit 2
    fi
    if ! has_cmake_arg "FASTVIDEO_KERNEL_BUILD_EXTENSION"; then
        CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_EXTENSION=OFF"
    fi
elif [[ "$(extension_mode_from_cmake_args)" == "OFF" ]]; then
    PYTHON_ONLY=1
fi

# The vendored headers only feed the compiled extension, so a Python-only build
# neither initializes nor checks them.
if ! (( PYTHON_ONLY )); then
    # Ensure only the kernel's required headers are initialized. A repository-wide
    # update also clones the unrelated VBench evaluation submodule. Skip outside a
    # git checkout (e.g. Docker contexts that exclude .git), where the submodule
    # contents must already be present.
    if git rev-parse --git-dir >/dev/null 2>&1; then
        git submodule update --init --recursive include/cutlass include/tk
    fi
    # Fail fast with a clear message if the headers are still missing (e.g. a
    # Docker context that excluded .git AND the submodule contents) instead of
    # dying later in a wall of nvcc include errors. CUTLASS is consumed by the
    # always-built turbodiffusion sources, so it is a hard error; ThunderKittens
    # only feeds the TK-gated Hopper kernels, so a missing tree just warns (the
    # TK gate resolves later, and non-SM90/ROCm builds never touch it).
    if [ ! -d include/cutlass/include ]; then
        echo "ERROR: include/cutlass/include is missing. Outside a git checkout the" >&2
        echo "       CUTLASS sources must already be present (run" >&2
        echo "       'git submodule update --init --recursive include/cutlass include/tk'" >&2
        echo "       in the source checkout, or include them in the build context)." >&2
        exit 1
    fi
    if [ ! -d include/tk/include ]; then
        echo "WARNING: include/tk/include is missing; ThunderKittens (Hopper sm_90a)" >&2
        echo "         kernels cannot be built. Fine for non-SM90/ROCm targets." >&2
    fi
fi

# Install build dependencies
uv pip install scikit-build-core cmake ninja

detect_with_torch() {
    # Prefer the active venv's python directly over `uv run --active --no-project`,
    # which on some uv versions provisions its own interpreter and misses packages
    # installed into VIRTUAL_ENV.
    local py
    if [[ -n "${VIRTUAL_ENV:-}" && -x "${VIRTUAL_ENV}/bin/python" ]]; then
        py="${VIRTUAL_ENV}/bin/python"
    else
        py="$(command -v python3 || command -v python)"
    fi
    "${py}" -c "import torch
if not torch.cuda.is_available():
    raise RuntimeError('torch.cuda.is_available() is false')
mj, mn = torch.cuda.get_device_capability(0)
print(f'{mj}.{mn}')"
}

if [ "${GPU_BACKEND}" = "CUDA" ] && ! (( PYTHON_ONLY )); then
    # Compute capability drives the arch/TK defaults below. Prefer an explicit
    # TORCH_CUDA_ARCH_LIST (works on GPU-less build machines such as CI/Docker);
    # only probe a live GPU via torch when no arch was provided.
    if [ -n "${TORCH_CUDA_ARCH_LIST:-}" ]; then
        echo "Using TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} (skipping torch GPU probe)"
        first_arch="${TORCH_CUDA_ARCH_LIST%%[;, ]*}"   # first entry, e.g. 9.0a
        first_arch="${first_arch%[af]}"                # strip trailing a/f suffix
        cc_major="${first_arch%%.*}"
        cc_minor="${first_arch##*.}"
    else
        detected_cc="$(detect_with_torch)" || {
            echo "ERROR: torch-based CUDA arch detection failed and TORCH_CUDA_ARCH_LIST is unset." >&2
            echo "       Set TORCH_CUDA_ARCH_LIST (e.g. 9.0a) for GPU-less builds, or build where CUDA is available." >&2
            exit 1
        }
        cc_major="${detected_cc%%.*}"
        cc_minor="${detected_cc##*.}"
        echo "Detected compute capability via torch: ${detected_cc}"
    fi
    cmake_arch="${cc_major}${cc_minor}"

    # Respect explicit overrides.
    if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then
        case "${cc_major}.${cc_minor}" in
            # Architecture-conditional kernels need the `a` target: Hopper TK (9.0),
            # data-center Blackwell VSA tcgen05 (10.0/10.3), and the attn_qat_infer
            # FP4 kernels that CMake's AUTO gate builds for 12.0a/12.1a.
            9.0|10.0|10.3|12.0|12.1)
                export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}a"
                ;;
            *)
                export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}"
                ;;
        esac
    fi

    if [[ "${TORCH_CUDA_ARCH_LIST}" =~ (^|[\;,[:space:]])(12\.1a|121a|sm_121a)([\;,[:space:]]|$) ]]; then
        cuda_compiler="${CUDACXX:-}"
        if [ -z "${cuda_compiler}" ] && [ -n "${CUDA_HOME:-}" ] && [ -x "${CUDA_HOME}/bin/nvcc" ]; then
            cuda_compiler="${CUDA_HOME}/bin/nvcc"
        fi
        if [ -z "${cuda_compiler}" ]; then
            cuda_compiler="$(command -v nvcc || true)"
        fi
        if [ -z "${cuda_compiler}" ]; then
            echo "ERROR: sm_121a requires CUDA Toolkit 13.0+, but nvcc was not found." >&2
            echo "       Set CUDACXX or CUDA_HOME to a CUDA 13 toolkit." >&2
            exit 1
        fi
        cuda_major="$("${cuda_compiler}" --version 2>/dev/null | sed -n 's/.*release \([0-9][0-9]*\)\..*/\1/p' | tail -n 1)"
        if [[ ! "${cuda_major}" =~ ^[0-9]+$ ]] || (( cuda_major < 13 )); then
            echo "ERROR: sm_121a requires CUDA Toolkit 13.0+; ${cuda_compiler} reports CUDA ${cuda_major:-unknown}." >&2
            exit 1
        fi
    fi

    # ThunderKittens build targeting:
    # - SM90: compile Hopper/TK kernels with 90a.
    # - Others (e.g., SM100): compile non-TK path with detected arch.
    if ! has_cmake_arg "CMAKE_CUDA_ARCHITECTURES"; then
        if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then
            CMAKE_ARGS="${CMAKE_ARGS:-} -DCMAKE_CUDA_ARCHITECTURES=90a"
        else
            CMAKE_ARGS="${CMAKE_ARGS:-} -DCMAKE_CUDA_ARCHITECTURES=${cmake_arch}"
        fi
    fi

    if ! has_cmake_arg "FASTVIDEO_KERNEL_BUILD_TK"; then
        # Let CMake detect both the target GPU arch and host support. Explicit ON
        # remains a request to fail if TK cannot be built on the host.
        CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=AUTO"
    fi
fi

if ! has_cmake_arg "GPU_BACKEND"; then
    CMAKE_ARGS="${CMAKE_ARGS:-} -DGPU_BACKEND=${GPU_BACKEND}"
fi
export CMAKE_ARGS

echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-<unset>}"
echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
echo "GPU_BACKEND: ${GPU_BACKEND:-<unset>}"
# Build and install
# Use -v for verbose output
if [ -n "${WHEEL_DIR}" ]; then
    mkdir -p "${WHEEL_DIR}"
    uv build --wheel -v --no-build-isolation --out-dir "${WHEEL_DIR}" .
    wheel_path="$(find "${WHEEL_DIR}" -maxdepth 1 -type f \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) | sort | tail -n 1)"
    if [ -z "${wheel_path}" ]; then
        echo "ERROR: no fastvideo-kernel wheel was built in ${WHEEL_DIR}" >&2
        exit 1
    fi
    uv pip install "${wheel_path}" --reinstall-package fastvideo-kernel --no-deps
else
    uv pip install . -v --no-build-isolation
fi