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

REPO_ID="${1:-}"
if [[ -z "$REPO_ID" ]]; then
  echo "Usage: $0 <USER_OR_ORG>/<REPO_NAME>"
  exit 2
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

SSH_KEY_PATH="${SSH_KEY_PATH:-$HOME/.ssh/hf_codex_upload}"
SSH_CMD="ssh -i ${SSH_KEY_PATH} -o IdentitiesOnly=yes -o StrictHostKeyChecking=no"
if [[ -n "${TMP_DIR:-}" ]]; then
  WORK_DIR="$TMP_DIR"
  mkdir -p "$WORK_DIR"
else
  WORK_DIR="$(mktemp -d /tmp/hf_model_repo_XXXXXX)"
fi
REPO_DIR="${WORK_DIR}/repo"

mkdir -p "$(dirname "$SSH_KEY_PATH")"
chmod 700 "$(dirname "$SSH_KEY_PATH")"

if [[ ! -f "$SSH_KEY_PATH" ]]; then
  ssh-keygen -t ed25519 -f "$SSH_KEY_PATH" -N "" -C "hf-upload" >/dev/null
  echo "[INFO] Generated SSH key: $SSH_KEY_PATH"
fi

echo ""
echo "1) Add this public key to https://huggingface.co/settings/keys :"
cat "${SSH_KEY_PATH}.pub"
echo ""
echo "2) After adding the key, run this script again to push."
echo ""

AUTH_MSG="$(ssh -o BatchMode=yes -i "$SSH_KEY_PATH" -o IdentitiesOnly=yes -T git@hf.co 2>&1 || true)"
if echo "$AUTH_MSG" | grep -q "Hi anonymous"; then
  echo "[STOP] SSH key not linked yet. Add the key in Hugging Face settings, then rerun:"
  echo "  $0 $REPO_ID"
  exit 1
fi

rm -rf "$REPO_DIR"
GIT_SSH_COMMAND="$SSH_CMD" git clone "git@hf.co:${REPO_ID}" "$REPO_DIR" >/dev/null
rsync -a --delete \
  --exclude '.git/' \
  --exclude 'output/**' \
  --exclude '**/__pycache__/**' \
  --exclude 'assets/best_model.pt' \
  --exclude 'assets/best_model.pt.part-*' \
  "$SCRIPT_DIR/" "$REPO_DIR/"

pushd "$REPO_DIR" >/dev/null
git config user.name "${GIT_AUTHOR_NAME:-hf-upload}"
git config user.email "${GIT_AUTHOR_EMAIL:-hf-upload@local}"
git add -A
if ! git diff --cached --quiet; then
  git commit -m "Upload inference bundle" >/dev/null
fi

GIT_SSH_COMMAND="$SSH_CMD" \
  GIT_TERMINAL_PROMPT=0 \
  git push -u origin main

popd >/dev/null
echo "[DONE] https://huggingface.co/${REPO_ID}"