#!/bin/bash
# flash_attn 휠을 Gitea 컨테이너 레지스트리에 push하는 스크립트
#
# 사용법:
#   # 현재 빌드된 comfyui 이미지에서 휠 추출하는 경우:
#   ./push-to-gitea.sh --from-image localhost/comfyui:latest
#
#   # 직접 flash-attn Containerfile로 빌드하는 경우:
#   ./push-to-gitea.sh --build

set -e

REGISTRY="10.89.0.8:3000"
REPO="gm/flash-attn-builder"
TAG="cu128-torch211"
FULL_IMAGE="${REGISTRY}/${REPO}:${TAG}"
GITEA_USER="gm"
GITEA_TOKEN="1f8061173dca537f12d3ac54b38fa8e5bf68493f"

case "$1" in
  --from-image)
    # 기존 이미지에서 builder 레이어 참조해 휠 추출
    SOURCE_IMAGE="${2:-localhost/comfyui:latest}"
    echo "==> $SOURCE_IMAGE 에서 flash_attn 휠 추출..."

    # 빌드된 이미지에서 builder 중간 레이어 검색
    # (multi-stage 빌드 결과 이미지엔 wheels가 없으므로 builder 이미지 활용)
    echo "==> flash-attn-builder 이미지 직접 빌드 (캐시 활용)..."
    cd "$(dirname "$0")"
    podman build \
      --tag "localhost/flash-attn-builder-tmp:latest" \
      --target builder \
      --file Containerfile \
      .

    echo "==> 휠 추출 → 최소 이미지 생성..."
    WHEEL_FILE=$(podman run --rm localhost/flash-attn-builder-tmp:latest ls /wheels/ | grep flash_attn)
    echo "    휠: $WHEEL_FILE"

    podman build \
      --tag "localhost/flash-attn-builder:${TAG}" \
      --file Containerfile \
      .
    ;;

  --build|*)
    echo "==> flash-attn-builder 이미지 빌드..."
    cd "$(dirname "$0")"
    # --cpu-quota/period: 2코어 하드 캡 (유휴시에도 초과 불가)
    # --cpu-shares 512: 경쟁 시 추가로 낮은 우선순위
    # MAX_JOBS=2: 컴파일러 병렬도를 CPU 캡에 맞춤
    podman build \
      --cpu-period=100000 \
      --cpu-quota=200000 \
      --cpu-shares=512 \
      --tag "localhost/flash-attn-builder:${TAG}" \
      --file Containerfile \
      .
    ;;
esac

echo ""
echo "==> Gitea 레지스트리 로그인..."
echo "${GITEA_TOKEN}" | podman login "${REGISTRY}" \
  --username "${GITEA_USER}" \
  --password-stdin

echo "==> 이미지 태그 및 push..."
podman tag "localhost/flash-attn-builder:${TAG}" "${FULL_IMAGE}"
podman push "${FULL_IMAGE}"

echo ""
echo "완료: ${FULL_IMAGE}"
echo "ComfyUI Containerfile에서 다음과 같이 사용:"
echo "  FROM ${FULL_IMAGE} AS flash-attn-wheels"
echo "  COPY --from=flash-attn-wheels /wheels/ /wheels/"