#!/bin/bash
# Prepare a FlagTree/FlagOS image to build and run DeepInfer's Triton backend.
#
# The kernels themselves need nothing but Triton, and FlagTree *is* a Triton fork -- that part
# works out of the box. What does not is everything around it: `deepinfer._C` is a torch C++
# extension, and the vendored DeepGEMM/DeepEP/DeepSelect are CUDA extensions. None of them
# involve Triton, and none of them ship in a vendor image, so they have to be built here.
# This script assembles a build environment where that succeeds.
#
# Run it *inside* the target container, from the repo root:
#
#     bash scripts/setup_flagtree_image.sh env     # system packages, venv, headers, deps
#     bash scripts/setup_flagtree_image.sh build   # _C and the three vendor wheels
#     bash scripts/setup_flagtree_image.sh all
#
# Before `env`, the caller stages reference packages on the host, because `docker cp` only
# exists outside the container:
#
#     mkdir -p /data/ftbuild/ref-triton /data/ftbuild/ref-tilelang /data/ftbuild/ref-site
#     docker cp <triton-image>:/usr/local/lib/python3.12/dist-packages/.  /data/ftbuild/ref-triton/
#     docker cp <tilelang-image>:/usr/local/lib/python3.12/dist-packages/. /data/ftbuild/ref-tilelang/
#     docker cp <engine-image>:/usr/local/lib/python3.12/dist-packages/.  /data/ftbuild/ref-site/
#
# `ref-triton` needs packages matching this image's torch build; `ref-site` supplies the
# Python packages from requirements.txt that the engine imports at load time.
#
# Every step is idempotent, so re-running after a failure resumes rather than starting over.

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SCRATCH="${SCRATCH:-/data/ftbuild}"
VENV_DIR="${VENV_DIR:-$SCRATCH/venv}"
WHEEL_DIR="${WHEEL_DIR:-$SCRATCH/wheels}"

REF_TRITON_DIR="${REF_TRITON_DIR:-$SCRATCH/ref-triton}"
REF_TILELANG_DIR="${REF_TILELANG_DIR:-$SCRATCH/ref-tilelang}"
REF_SITE_PACKAGES="${REF_SITE_PACKAGES:-$SCRATCH/ref-site}"

# DeepEP calls ncclCommQueryProperties and reads props.ginMinStride, both of which arrived in
# NCCL 2.32.3. Vendor images ship 2.28.9, so this one has to be replaced.
NCCL_VERSION="${NCCL_VERSION:-2.32.3}"
NCCL_URL="${NCCL_URL:-https://mirrors.aliyun.com/pypi/packages/5b/29/6b277e63c92d91f9cb4d1a3a554e148983de39d54baa652bb52c798af78e/nvidia_nccl_cu13-${NCCL_VERSION}-py3-none-manylinux_2_27_x86_64.whl}"

MAX_JOBS="${MAX_JOBS:-$(nproc)}"
PY_TAG="python$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')"
SITE="$VENV_DIR/lib/$PY_TAG/site-packages"

say() { printf '\n\033[1m== %s\033[0m\n' "$*"; }

# Copy one or more site-packages entries, expanded from shell globs, replacing what is there.
# Globbing the dist-info directory rather than hardcoding its name keeps this working when the
# reference image bumps a version.
stage() {
    local label="$1" src="$2" dst="$3"
    shift 3
    local count=0 path
    for pattern in "$@"; do
        for path in $src/$pattern; do
            [ -e "$path" ] || continue
            rm -rf "${dst:?}/$(basename "$path")"
            cp -a "$path" "$dst/"
            count=$((count + 1))
        done
    done
    if [ "$count" -eq 0 ]; then
        echo "  !! $label: nothing matched under $src (looked for: $*)" >&2
        return 1
    fi
    echo "  $label: $count entr$([ "$count" -eq 1 ] && echo y || echo ies)"
}

step_env() {
    say "preflight"
    command -v nvcc >/dev/null || { echo "nvcc missing: this image has no CUDA toolkit" >&2; exit 1; }
    python3 -c 'import torch; assert torch.version.cuda' || { echo "torch is not a CUDA build" >&2; exit 1; }
    echo "  torch    $(python3 -c 'import torch; print(torch.__version__)')"
    echo "  nvcc     $(nvcc --version | tail -1 | sed 's/.*, //')"
    echo "  gcc      $(gcc -dumpversion)"
    echo "  repo     $REPO_ROOT"
    echo "  venv     $VENV_DIR"

    say "system packages"
    # libdw/libelf: DeepGEMM's bundled deep_jit resolves stack frames through libdwfl.
    # libzstd: DeepEP links it. pciutils: DeepSelect's setup.py shells out to lspci to pick
    # the build target. git: vendor submodule tooling.
    if command -v apt-get >/dev/null; then
        export DEBIAN_FRONTEND=noninteractive
        apt-get update -qq || echo "  (apt update had warnings; continuing)"
        apt-get install -y -qq libdw-dev libelf-dev libzstd-dev pciutils git >/dev/null
        echo "  installed"
    else
        echo "  no apt-get here; make sure libdw, libelf, libzstd and lspci are present"
    fi

    say "venv"
    # setup.py refuses to build unless CONDA_PREFIX or VIRTUAL_ENV is set, and only uses the
    # value as a guard -- it never reads it. A venv is the lightest way to satisfy it while
    # inheriting the image's torch.
    if [ ! -x "$VENV_DIR/bin/python" ]; then
        python3 -m venv --system-site-packages "$VENV_DIR"
    fi
    # shellcheck disable=SC1091
    source "$VENV_DIR/bin/activate"
    python -c 'import torch; print("  torch visible in venv:", torch.__version__)'

    say "CUDA headers"
    # The image ships the CUDA runtime, not the toolkit headers; torch's own headers include
    # cusparse.h, so building any extension fails without it. The nvidia-*-cu13 wheels carry
    # the headers, so borrow them rather than installing a full toolkit.
    if [ ! -f /usr/local/cuda/include/cusparse.h ]; then
        for hdr in /usr/local/lib/"$PY_TAG"/dist-packages/nvidia/cu13/include \
                   /usr/local/lib/"$PY_TAG"/dist-packages/nvidia/cuda_runtime/include; do
            [ -d "$hdr" ] && cp -n "$hdr"/*.h /usr/local/cuda/include/ 2>/dev/null
        done
    fi
    if [ -f /usr/local/cuda/include/cusparse.h ]; then
        echo "  cusparse.h present"
    else
        echo "  !! cusparse.h still missing; _C will not compile" >&2
        return 1
    fi

    say "NCCL $NCCL_VERSION"
    if python -c "import nvidia.nccl" 2>/dev/null &&
        [ "$(python -c 'import importlib.metadata as m; print(m.version("nvidia-nccl-cu13"))' 2>/dev/null)" = "$NCCL_VERSION" ]; then
        echo "  already $NCCL_VERSION"
    else
        mkdir -p "$SCRATCH/pipdl"
        local whl="$SCRATCH/pipdl/nvidia_nccl_cu13-${NCCL_VERSION}.whl"
        # pip's own download stalls on this host for a 300MB wheel; curl with resume finishes.
        [ -s "$whl" ] || curl -sSL -C - --retry 5 --retry-delay 3 --max-time 900 -o "$whl" "$NCCL_URL"
        # pip validates the filename, so it has to look like one.
        mv "$whl" "$SCRATCH/pipdl/nvidia_nccl_cu13-${NCCL_VERSION}-py3-none-manylinux_2_27_x86_64.whl"
        pip install -q --no-deps --force-reinstall \
            "$SCRATCH/pipdl/nvidia_nccl_cu13-${NCCL_VERSION}-py3-none-manylinux_2_27_x86_64.whl"
        echo "  installed"
    fi

    say "Triton (= FlagTree)"
    # FlagTree is a fork of Triton: the import name stays `triton`, and `flagtree` only holds
    # FlagPrism, which triton imports at load time. Both halves plus their metadata are needed,
    # because the driver is registered through an entry point that only dist-info declares.
    stage triton "$REF_TRITON_DIR" "$SITE" \
        triton flagtree 'tokenspeed*triton*.dist-info' 'flagtree*.dist-info'
    python -c 'import triton, flagtree; print("  triton", triton.__version__, "= flagtree", flagtree.__version__)'

    say "TileLang runtime"
    # Only the fallback path needs these -- the kernels Triton does not implement yet are still
    # served by TileLang. tile_kernels is the quantiser used by the reference implementations.
    # tvm_ffi must come along: the image's own 0.1.9 lacks tvm_ffi.dataclasses.Enum, which
    # tilelang 0.1.15 imports during startup.
    if [ -d "$REF_TILELANG_DIR/tilelang" ]; then
        stage tilelang "$REF_TILELANG_DIR" "$SITE" \
            tilelang tile_kernels tvm_ffi \
            'tilelang-*.dist-info' 'tile_kernels-*.dist-info' 'apache_tvm_ffi-*.dist-info'
        python -c 'import tilelang, tile_kernels; print("  tilelang", tilelang.__version__)'
    else
        echo "  skipped: no $REF_TILELANG_DIR (end-to-end will fail on the MTP fallback path)"
    fi

    say "Python dependencies"
    if [ -d "$REF_SITE_PACKAGES" ]; then
        # pillow stays as the image ships it: its wheels carry vendored libs (libtiff and
        # friends) under pillow.libs, and a copy from another image points at different ones.
        python "$REPO_ROOT/scripts/copy_site_packages.py" "$REF_SITE_PACKAGES" "$SITE" --keep-system pillow
    else
        echo "  no $REF_SITE_PACKAGES; falling back to pip (may stall on this host)"
        pip install -q -r "$REPO_ROOT/requirements.txt" || echo "  !! pip failed; stage ref-site instead"
    fi

    say "env ready"
    echo "  source $VENV_DIR/bin/activate"
}

step_build() {
    # shellcheck disable=SC1091
    source "$VENV_DIR/bin/activate"
    cd "$REPO_ROOT"
    mkdir -p "$SCRATCH/tmp" "$WHEEL_DIR"
    export TMPDIR="$SCRATCH/tmp"
    export DEEPINFER_ACCELERATOR=nvidia
    export MAX_JOBS

    say "_C extension"
    bash scripts/build.sh build

    say "vendor wheels"
    for pkg in DeepGEMM DeepEP DeepSelect; do
        [ -f "vendor/$pkg/setup.py" ] || {
            echo "  !! vendor/$pkg is empty -- run 'git submodule update --init' on a git checkout" >&2
            return 1
        }
        echo "  building $pkg"
        rm -rf "vendor/$pkg/build" "vendor/$pkg"/*.egg-info
        local envs=()
        # Matches scripts/build_vendor.sh: DeepGEMM only rebuilds when forced, DeepSelect needs
        # its platform target spelled out.
        [ "$pkg" = DeepGEMM ] && envs+=(DG_FORCE_BUILD=1)
        [ "$pkg" = DeepSelect ] && envs+=(DEEP_SELECT_BUILD_TARGET_PLATFORM=CUDA)
        (
            cd "vendor/$pkg"
            env "${envs[@]}" python setup.py bdist_wheel --dist-dir "$WHEEL_DIR" >/dev/null
        )
    done
    ls -lh "$WHEEL_DIR"

    say "install wheels"
    # --no-deps and a forced reinstall: the image has its own deep_ep, and pip would otherwise
    # keep it because the version number happens to look newer.
    pip install -q --no-deps --force-reinstall \
        "$WHEEL_DIR"/deep_gemm-*.whl "$WHEEL_DIR"/deep_ep-*.whl "$WHEEL_DIR"/deep_select-*.whl
    pip list 2>/dev/null | grep -iE "^(deep_gemm|deep_ep|deep_select|tilelang|triton) " || true
}

step_verify() {
    # shellcheck disable=SC1091
    source "$VENV_DIR/bin/activate"
    cd "$REPO_ROOT"
    export PYTHONPATH="$REPO_ROOT:$REPO_ROOT/examples:$REPO_ROOT/scripts"
    export DEEPINFER_OPS_BACKEND=triton DEEPINFER_OPS_STRICT=1

    say "imports"
    python - <<'PY'
import torch, triton, flagtree
print("  torch   ", torch.__version__, "cap", torch.cuda.get_device_capability(0))
print("  triton  ", triton.__version__, "(= FlagTree", flagtree.__version__ + ")")
import deep_gemm, deep_ep, deep_select
print("  vendor   deep_gemm", deep_gemm.__version__, "| deep_ep", deep_ep.__version__)
from deepinfer._C import engine
print("  _C      ", len([a for a in dir(engine) if not a.startswith("_")]), "symbols")
PY

    say "ops backend"
    python -c '
from deepinfer.model_executor.layers.tilelang_ops import ops
print("  backend", ops._BACKEND, "strict", ops._STRICT, "kernels", len(ops._SELECTED))'
}

case "${1:-all}" in
    env) step_env ;;
    build) step_build ;;
    verify) step_verify ;;
    all)
        step_env
        step_build
        step_verify
        ;;
    *)
        echo "usage: $0 {env|build|verify|all}" >&2
        exit 2
        ;;
esac