#!/usr/bin/env bash
set -euo pipefail

echo "🚀 Builder (Apex + Q8 + FlashAttention) — runtime, GPU visível, cache persistente"

# Persistência e caches
if [ -d /data ]; then
  export HF_HOME="${HF_HOME:-/data/.cache/huggingface}"
  export TORCH_HOME="${TORCH_HOME:-/data/.cache/torch}"
else
  export HF_HOME="${HF_HOME:-/app/.cache/huggingface}"
  export TORCH_HOME="${TORCH_HOME:-/app/.cache/torch}"
fi
export HF_HUB_CACHE="${HF_HUB_CACHE:-$HF_HOME/hub}"
mkdir -p "$HF_HOME" "$HF_HUB_CACHE" "$TORCH_HOME"
mkdir -p /app/.cache && ln -sf "$HF_HOME" /app/.cache/huggingface

# Repo de wheels
export SELF_HF_REPO_ID="${SELF_HF_REPO_ID:-carlex3321/aduc-sdr3}"

# Transfer accel
export HF_HUB_ENABLE_HF_TRANSFER="${HF_HUB_ENABLE_HF_TRANSFER:-1}"
export HF_HUB_DOWNLOAD_TIMEOUT="${HF_HUB_DOWNLOAD_TIMEOUT:-60}"

# Work dirs
mkdir -p /app/wheels /app/cuda_cache /app/wheels/src
chmod -R 777 /app/wheels || true
export CUDA_CACHE_PATH="/app/cuda_cache"
[ -f "/NGC-DL-CONTAINER-LICENSE" ] && cp -f /NGC-DL-CONTAINER-LICENSE /app/wheels/NGC-DL-CONTAINER-LICENSE || true

# Build deps
python -m pip install -U pip build setuptools wheel hatchling hatch-vcs scikit-build-core cmake ninja packaging "huggingface_hub[hf_transfer]" || true

# Tags do ambiente
PY_TAG="$(python -c 'import sys; print(f"cp{sys.version_info[0]}{sys.version_info[1]}")' 2>/dev/null || echo cp310)"
TORCH_VER="$(python - <<'PY'
try:
    import torch, re
    v = torch.__version__
    print(re.sub(r'\+.*$', '', v))
except Exception:
    print("unknown")
PY
)"
CU_TAG="$(python - <<'PY'
try:
    import torch
    cu = getattr(torch.version, "cuda", None)
    print("cu"+cu.replace(".","")) if cu else print("")
except Exception:
    print("")
PY
)"
echo "[env] PY_TAG=${PY_TAG} TORCH_VER=${TORCH_VER} CU_TAG=${CU_TAG}"

# Checkers
check_apex() {
python - <<'PY'
try:
    from apex.normalization import FusedLayerNorm, FusedRMSNorm
    import importlib; importlib.import_module("fused_layer_norm_cuda")
    ok = True
except Exception:
    ok = False
raise SystemExit(0 if ok else 1)
PY
}
check_q8() {
python - <<'PY'
import importlib.util
spec = importlib.util.find_spec("ltx_q8_kernels") or importlib.util.find_spec("q8_kernels")
raise SystemExit(0 if spec else 1)
PY
}
check_flash() {
python - <<'PY'
ok = False
try:
    import importlib
    for name in ("flash_attn_2_cuda","flash_attn.ops.layer_norm","flash_attn.layers.layer_norm","flash_attn"):
        try:
            importlib.import_module(name); ok=True; break
        except Exception:
            pass
except Exception:
    ok = False
raise SystemExit(0 if ok else 1)
PY
}

# Baixar wheel do Hub
install_from_hf () {
  local PKG="$1"
  python - "$PKG" "$PY_TAG" "$CU_TAG" <<'PY' || exit 0
import os, sys
from huggingface_hub import HfApi, hf_hub_download, HfFolder
pkg, py_tag, cu_tag = sys.argv[1], sys.argv[2], sys.argv[3]
repo = os.environ.get("SELF_HF_REPO_ID","carlex3321/aduc-sdr3")
api = HfApi(token=os.getenv("HF_TOKEN") or HfFolder.get_token())
try:
    files = api.list_repo_files(repo_id=repo, repo_type="model")
except Exception:
    raise SystemExit(0)
cands = [f for f in files if f.endswith(".whl") and f.rsplit("/",1)[-1].startswith(pkg+"-") and py_tag in f]
pref = [f for f in cands if cu_tag and cu_tag in f] or cands
if not pref: raise SystemExit(0)
target = sorted(pref, reverse=True)[0]
print(target)
path = hf_hub_download(repo_id=repo, filename=target, repo_type="model", local_dir="/app/wheels")
print(path)
PY
}

# Builders
build_apex () {
  local SRC="/app/wheels/src/apex"
  if [ -d "$SRC/.git" ]; then
    git -C "$SRC" fetch --all -p || true
    git -C "$SRC" reset --hard HEAD || true
    git -C "$SRC" clean -fdx || true
  else
    rm -rf "$SRC"
    git clone --depth 1 https://github.com/NVIDIA/apex "$SRC"
  fi
  export APEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_ALL_CONTRIB_EXT=0
  python -m pip wheel --no-build-isolation --no-deps "$SRC" -w /app/wheels || true
  local W="$(ls -t /app/wheels/apex-*.whl 2>/dev/null | head -n1 || true)"
  if [ -n "${W}" ]; then
    python -m pip install -U --no-deps "${W}" || true
  else
    python -m pip install --no-build-isolation "$SRC" || true
  fi
}
Q8_REPO="${Q8_REPO:-https://github.com/Lightricks/LTX-Video-Q8-Kernels.git}"
Q8_COMMIT="${Q8_COMMIT:-f3066edea210082799ca5a2bbf9ef0321c5dd8fc}"
build_q8 () {
  local SRC="/app/wheels/src/q8_kernels"
  rm -rf "$SRC"
  git clone --filter=blob:none "$Q8_REPO" "$SRC"
  git -C "$SRC" checkout "$Q8_COMMIT"
  git -C "$SRC" submodule update --init --recursive
  python -m pip wheel --no-build-isolation "$SRC" -w /app/wheels || true
  local W="$(ls -t /app/wheels/q8_kernels-*.whl 2>/dev/null | head -n1 || true)"
  if [ -n "${W}" ]; then
    python -m pip install -U --no-deps "${W}" || true
  else
    python -m pip install --no-build-isolation "$SRC" || true
  fi
}
FLASH_ATTENTION_TAG="${FLASH_ATTENTION_TAG:-v2.6.3}"
build_flash () {
  set -e
  local SRC="/app/wheels/src/flash-attn"
  rm -rf "$SRC"
  git clone --depth 1 --branch "$FLASH_ATTENTION_TAG" https://github.com/Dao-AILab/flash-attention.git "$SRC"
  export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.9}"
  export MAX_JOBS="${MAX_JOBS:-$(nproc)}"
  export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
  python -m pip wheel --no-build-isolation --no-deps "$SRC" -w /app/wheels || true
  local W="$(ls -t /app/wheels/flash_attn-*.whl 2>/dev/null | head -n1 || true)"
  if [ -n "${W}" ]; then
    python -m pip install -U --no-deps "${W}" || true
  else
    python -m pip install --no-build-isolation "$SRC" || true
  fi
}

# Orquestrador
ensure_pkg () {
  local PKG="$1"; local CHECK_FN="$2"; local BUILD_FN="$3"
  if ${CHECK_FN}; then
    echo "[flow] ${PKG}: já instalado"; return 0
  fi
  echo "[flow] ${PKG}: tentando wheel do Hub (${SELF_HF_REPO_ID})"
  HF_OUT="$(install_from_hf "$PKG" || true)"
  if [ -n "${HF_OUT:-}" ]; then
    WHEEL_PATH="$(printf "%s\n" "${HF_OUT}" | tail -n1)"
    python -m pip install -U --no-build-isolation "${WHEEL_PATH}" || true
    if ${CHECK_FN}; then
      echo "[flow] ${PKG}: sucesso via Hub (${WHEEL_PATH})"; return 0
    fi
  fi
  echo "[flow] ${PKG}: compilando (fallback)"
  ${BUILD_FN} || true
  ${CHECK_FN} || echo "[flow] ${PKG}: falhou após build; seguindo"
}

# Execução
ensure_pkg "apex"       check_apex   build_apex   || true
#ensure_pkg "q8_kernels" check_q8     build_q8     || true
ensure_pkg "flash_attn" check_flash  build_flash  || true

# Upload das wheels
python - <<'PY'
import os
from huggingface_hub import HfApi, HfFolder
repo=os.environ.get("SELF_HF_REPO_ID","carlex3321/aduc-sdr3")
token=os.getenv("HF_TOKEN") or HfFolder.get_token()
if not token:
    raise SystemExit(0)
api=HfApi(token=token)
api.upload_folder(
    folder_path="/app/wheels",
    repo_id=repo,
    repo_type="model",
    allow_patterns=["*.whl","NGC-DL-CONTAINER-LICENSE"],
    ignore_patterns=["**/src/**","**/*.log","**/logs/**",".git/**"],
)
print("Upload de wheels concluído.")
PY

chmod -R 777 /app/wheels || true
echo "✅ Builder finalizado."