#!/bin/sh
# ===========================================
# Staging 部署脚本（SSH 模式，并行优化版）
# ===========================================
set -eu

retry_cmd() {
  local max_attempts=$1
  local backoff=$2
  shift 2
  local attempt=1
  while [ $attempt -le $max_attempts ]; do
    if "$@"; then
      return 0
    fi
    echo "  attempt $attempt/$max_attempts failed, retrying in ${backoff}s..."
    sleep $backoff
    backoff=$((backoff * 2))
    attempt=$((attempt + 1))
  done
  echo "  ERROR: failed after $max_attempts retries"
  return 1
}

retry_docker_login() {
  echo "Logging in to registry (up to 3 retries)"
  export REGISTRY_TOKEN REGISTRY_HOST REGISTRY_USER
  if retry_cmd 3 5 sh -c 'printf "%s" "$REGISTRY_TOKEN" | docker login "$REGISTRY_HOST" -u "$REGISTRY_USER" --password-stdin'; then
    return 0
  fi
  echo "WARN: docker login failed after retries, will try pull anyway"
  return 0
}

retry_docker_pull() {
  local image=$1
  echo "Pulling $image (up to 3 retries)"
  retry_cmd 3 10 docker pull "$image"
}

IMAGE_TAG="${IMAGE_TAG:-}"
REGISTRY="${REGISTRY:-xiaoxia-registry.cn-hangzhou.cr.aliyuncs.com/xiaoxiakeji}"
REGISTRY_USER="${ACR_USERNAME:-${REGISTRY_USER:-nick0415343655}}"
REGISTRY_TOKEN="${ACR_PASSWORD:-${REGISTRY_TOKEN:-}}"

ENV_FILE="${ENV_FILE:-/var/lib/xiaoxia-saas-staging/.env}"
GENERATED_DIR="${GENERATED_DIR:-/var/lib/xiaoxia-saas-staging/generated}"
LEGACY_ASSETS_DIR="${LEGACY_ASSETS_DIR:-/var/lib/xiaoxia-saas-staging/legacy-assets}"

SKIP_MIGRATION="${SKIP_MIGRATION:-false}"
SKIP_ROLLBACK="${SKIP_ROLLBACK:-false}"

if [ -z "$IMAGE_TAG" ]; then
  echo "ERROR: IMAGE_TAG is required"
  exit 1
fi

test -f "$ENV_FILE"
mkdir -p "$GENERATED_DIR"
mkdir -p "$LEGACY_ASSETS_DIR"

echo "==========================================="
echo "  Staging 部署 - $IMAGE_TAG (并行优化版)"
echo "==========================================="

echo "Recording current image versions for rollback..."
PREV_API_IMAGE=""
PREV_WORKER_IMAGE=""
PREV_WEB_IMAGE=""
for c in xiaoxia-api-staging xiaoxia-worker-staging xiaoxia-web-staging; do
  if docker inspect "$c" >/dev/null 2>&1; then
    img=$(docker inspect -f '{{.Config.Image}}' "$c")
    case "$c" in
      xiaoxia-api-staging) PREV_API_IMAGE="$img" ;;
      xiaoxia-worker-staging) PREV_WORKER_IMAGE="$img" ;;
      xiaoxia-web-staging) PREV_WEB_IMAGE="$img" ;;
    esac
    echo "  $c -> $img"
  else
    echo "  $c -> (not running)"
  fi
done

rollback() {
  echo ""
  echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
  echo "  部署失败，正在自动回滚到上一版本..."
  echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
  echo ""

  if [ "$SKIP_ROLLBACK" = "true" ]; then
    echo "SKIP_ROLLBACK=true，跳过自动回滚"
    exit 1
  fi

  echo "Stopping new containers..."
  docker rm -f xiaoxia-api-staging 2>/dev/null || true
  docker rm -f xiaoxia-worker-staging 2>/dev/null || true
  docker rm -f xiaoxia-web-staging 2>/dev/null || true

  LOG_OPTS="--log-driver json-file --log-opt max-size=50m --log-opt max-file=3"

  if [ -n "$PREV_API_IMAGE" ]; then
    echo "Rolling back API to: $PREV_API_IMAGE"
    docker run -d \
      --name xiaoxia-api-staging \
      --env-file "$ENV_FILE" \
      --network xiaoxia-net-staging \
      -p 127.0.0.1:8000:8000 \
      -e APP_ENV=staging \
      -e APP_VERSION="$(echo $PREV_API_IMAGE | grep -oE '[^:]+$')" \
      -e GENERATED_FILES_DIR=/app/generated \
      -e GENERATED_FILES_URL_PREFIX=/generated-files \
      -e PUBLIC_API_BASE_URL=https://staging-api.xiaoxiajianji.com \
      -v "$GENERATED_DIR:/app/generated" \
      --restart unless-stopped \
      --health-cmd "python -c \"import urllib.request; urllib.request.urlopen('http://localhost:8000/health', timeout=5)\"" \
      --health-interval 30s \
      --health-timeout 10s \
      --health-retries 3 \
      --health-start-period 40s \
      $LOG_OPTS \
      "$PREV_API_IMAGE" &
  fi

  if [ -n "$PREV_WORKER_IMAGE" ]; then
    echo "Rolling back Worker to: $PREV_WORKER_IMAGE"
    docker run -d \
      --name xiaoxia-worker-staging \
      --env-file "$ENV_FILE" \
      --network xiaoxia-net-staging \
      -e APP_ENV=staging \
      -e APP_VERSION="$(echo $PREV_WORKER_IMAGE | grep -oE '[^:]+$')" \
      -e WORKER_CONCURRENCY=1 \
      -e WORKER_MAX_TASKS_PER_CHILD=100 \
      -e GENERATED_FILES_DIR=/app/generated \
      -e GENERATED_FILES_URL_PREFIX=/generated-files \
      -e PUBLIC_API_BASE_URL=https://staging-api.xiaoxiajianji.com \
      -v "$GENERATED_DIR:/app/generated" \
      --restart unless-stopped \
      --health-cmd "sh -c \"grep -q celery /proc/1/cmdline || exit 1\"" \
      --health-interval 30s \
      --health-timeout 10s \
      --health-retries 3 \
      --health-start-period 30s \
      $LOG_OPTS \
      "$PREV_WORKER_IMAGE" &
  fi

  if [ -n "$PREV_WEB_IMAGE" ]; then
    echo "Rolling back Web to: $PREV_WEB_IMAGE"
    LEGACY_VOLUME=""
    if [ -d "$LEGACY_ASSETS_DIR" ] && [ "$(ls -A "$LEGACY_ASSETS_DIR" 2>/dev/null)" ]; then
      LEGACY_VOLUME="-v ${LEGACY_ASSETS_DIR}:/usr/share/nginx/html/assets-legacy/assets:ro"
    fi
    docker run -d \
      --name xiaoxia-web-staging \
      --network xiaoxia-net-staging \
      -p 127.0.0.1:3001:80 \
      --restart unless-stopped \
      $LEGACY_VOLUME \
      --health-cmd "wget --spider -q http://127.0.0.1:80" \
      --health-interval 30s \
      --health-timeout 5s \
      --health-retries 3 \
      $LOG_OPTS \
      "$PREV_WEB_IMAGE" &
  fi

  wait

  if [ -n "$PREV_API_IMAGE" ]; then
    echo "Waiting for rolled-back API to become healthy..."
    i=0
    while [ "$i" -lt 40 ]; do
      if curl -sf --max-time 5 http://127.0.0.1:8000/health >/dev/null 2>&1; then
        echo "Rolled-back API is healthy!"
        break
      fi
      i=$((i + 1))
      echo "  Waiting... ($i/40)"
      sleep 3
    done
    if [ "$i" -ge 40 ]; then
      echo "WARN: Rolled-back API did not become healthy within 120s"
      docker logs --tail 30 xiaoxia-api-staging
    fi
  fi

  echo ""
  echo "==========================================="
  echo "  回滚完成"
  echo "==========================================="
  echo "Previous API:    ${PREV_API_IMAGE:-none}"
  echo "Previous Worker: ${PREV_WORKER_IMAGE:-none}"
  echo "Previous Web:    ${PREV_WEB_IMAGE:-none}"
  echo ""
  echo "部署失败，已自动回滚到上一版本"
  docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Image}}" | grep staging
  exit 1
}

if [ -n "$REGISTRY_TOKEN" ]; then
  echo "=========================================="
  echo "  Login to Registry (with retries)"
  echo "=========================================="
  REGISTRY_HOST=$(echo "$REGISTRY" | cut -d/ -f1)
  retry_docker_login
fi

# ---- 并行 Pull 三个镜像 ----
REGISTRY_API="${REGISTRY}/xiaoxia-saas-api:${IMAGE_TAG}"
REGISTRY_WORKER="${REGISTRY}/xiaoxia-saas-worker:${IMAGE_TAG}"
REGISTRY_WEB="${REGISTRY}/xiaoxia-saas-web:${IMAGE_TAG}"

echo "=========================================="
echo "  Pull images (parallel, up to 3 retries each)"
echo "=========================================="
PULL_LOG_DIR="/tmp/staging-pull-$$"
mkdir -p "$PULL_LOG_DIR"

retry_docker_pull "$REGISTRY_API" > "$PULL_LOG_DIR/api.log" 2>&1 &
PID_API=$!
retry_docker_pull "$REGISTRY_WORKER" > "$PULL_LOG_DIR/worker.log" 2>&1 &
PID_WORKER=$!
retry_docker_pull "$REGISTRY_WEB" > "$PULL_LOG_DIR/web.log" 2>&1 &
PID_WEB=$!

wait $PID_API $PID_WORKER $PID_WEB

echo ""
echo "Pull 结果:"
PULL_FAILED=0
for svc in api worker web; do
  if tail -1 "$PULL_LOG_DIR/$svc.log" 2>/dev/null | grep -qE "Status:|Downloaded|already exists|is up to date"; then
    echo "  OK $svc"
  elif grep -qE "Digest:|Status: Downloaded" "$PULL_LOG_DIR/$svc.log" 2>/dev/null; then
    echo "  OK $svc"
  else
    # 检查docker pull返回值不直接，用镜像是否存在来判断
    img_var="REGISTRY_$(echo $svc | tr '[:lower:]' '[:upper:]')"
    img_val=$(eval echo "\$$img_var")
    if docker image inspect "$img_val" >/dev/null 2>&1; then
      echo "  OK $svc"
    else
      echo "  FAIL $svc"
      tail -5 "$PULL_LOG_DIR/$svc.log" 2>/dev/null || true
      PULL_FAILED=$((PULL_FAILED + 1))
    fi
  fi
done

rm -rf "$PULL_LOG_DIR"

if [ "$PULL_FAILED" -gt 0 ]; then
  echo ""
  echo "ERROR: $PULL_FAILED 个镜像 pull 失败"
  exit 1
fi

echo "All images pulled."

echo "Backing up legacy assets from current web container..."
if docker inspect xiaoxia-web-staging >/dev/null 2>&1; then
  _tmpdir="/tmp/legacy-assets-$$"
  rm -rf "$_tmpdir"
  mkdir -p "$_tmpdir"
  docker cp xiaoxia-web-staging:/usr/share/nginx/html/assets/. "$_tmpdir/" 2>/dev/null || true
  if [ -d "$_tmpdir" ] && [ "$(ls -A "$_tmpdir" 2>/dev/null)" ]; then
    cp -an "$_tmpdir"/. "$LEGACY_ASSETS_DIR"/ 2>/dev/null || true
    echo "Legacy assets backed up: $(ls "$_tmpdir" | wc -l) files"
  fi
  rm -rf "$_tmpdir"
else
  echo "No existing web container, skipping legacy assets backup"
fi

if [ -d "$LEGACY_ASSETS_DIR" ]; then
  find "$LEGACY_ASSETS_DIR" -type f -mtime +7 -delete 2>/dev/null || true
  echo "Legacy assets cleanup done (retain 7 days)"
fi

echo "Checking infrastructure containers..."
for c in xiaoxia-postgres-staging xiaoxia-redis-staging; do
  if ! docker inspect "$c" >/dev/null 2>&1; then
    echo "ERROR: Required container not found: $c"
    exit 1
  fi
  state=$(docker inspect -f '{{.State.Status}}' "$c")
  if [ "$state" != "running" ]; then
    echo "ERROR: Container not running: $c ($state)"
    exit 1
  fi
done

docker network create xiaoxia-net-staging 2>/dev/null || true

if [ "$SKIP_MIGRATION" != "true" ]; then
  echo "Running database migrations..."
  docker run --rm \
    --env-file "$ENV_FILE" \
    --network xiaoxia-net-staging \
    -e APP_ENV=staging \
    "$REGISTRY_API" sh -c "cd /app && alembic upgrade head" || {
    echo "ERROR: Database migration failed"
    exit 1
  }
  echo "Migrations completed."
else
  echo "Skipping migrations (SKIP_MIGRATION=true)"
fi

echo "Stopping old containers..."
docker rm -f xiaoxia-api-staging 2>/dev/null || true
docker rm -f xiaoxia-worker-staging 2>/dev/null || true
docker rm -f xiaoxia-web-staging 2>/dev/null || true

LOG_OPTS="--log-driver json-file --log-opt max-size=50m --log-opt max-file=3"

# ---- 并行启动三个容器 ----
echo "Starting all containers (parallel)..."

LEGACY_VOLUME=""
if [ -d "$LEGACY_ASSETS_DIR" ] && [ "$(ls -A "$LEGACY_ASSETS_DIR" 2>/dev/null)" ]; then
  LEGACY_VOLUME="-v ${LEGACY_ASSETS_DIR}:/usr/share/nginx/html/assets-legacy/assets:ro"
fi

docker run -d \
  --name xiaoxia-api-staging \
  --env-file "$ENV_FILE" \
  --network xiaoxia-net-staging \
  -p 127.0.0.1:8000:8000 \
  -e APP_ENV=staging \
  -e APP_VERSION="$IMAGE_TAG" \
  -e GENERATED_FILES_DIR=/app/generated \
  -e GENERATED_FILES_URL_PREFIX=/generated-files \
  -e PUBLIC_API_BASE_URL=https://staging-api.xiaoxiajianji.com \
  -v "$GENERATED_DIR:/app/generated" \
  --restart unless-stopped \
  --health-cmd "python -c \"import urllib.request; urllib.request.urlopen('http://localhost:8000/health', timeout=5)\"" \
  --health-interval 30s \
  --health-timeout 10s \
  --health-retries 3 \
  --health-start-period 40s \
  $LOG_OPTS \
  "$REGISTRY_API" &
PID_API_START=$!

docker run -d \
  --name xiaoxia-worker-staging \
  --env-file "$ENV_FILE" \
  --network xiaoxia-net-staging \
  -e APP_ENV=staging \
  -e APP_VERSION="$IMAGE_TAG" \
  -e WORKER_CONCURRENCY=1 \
  -e WORKER_MAX_TASKS_PER_CHILD=100 \
  -e GENERATED_FILES_DIR=/app/generated \
  -e GENERATED_FILES_URL_PREFIX=/generated-files \
  -e PUBLIC_API_BASE_URL=https://staging-api.xiaoxiajianji.com \
  -v "$GENERATED_DIR:/app/generated" \
  --restart unless-stopped \
  --health-cmd "sh -c \"grep -q celery /proc/1/cmdline || exit 1\"" \
  --health-interval 30s \
  --health-timeout 10s \
  --health-retries 3 \
  --health-start-period 30s \
  $LOG_OPTS \
  "$REGISTRY_WORKER" &
PID_WORKER_START=$!

docker run -d \
  --name xiaoxia-web-staging \
  --network xiaoxia-net-staging \
  -p 127.0.0.1:3001:80 \
  --restart unless-stopped \
  $LEGACY_VOLUME \
  --health-cmd "wget --spider -q http://127.0.0.1:80" \
  --health-interval 30s \
  --health-timeout 5s \
  --health-retries 3 \
  $LOG_OPTS \
  "$REGISTRY_WEB" &
PID_WEB_START=$!

wait $PID_API_START $PID_WORKER_START $PID_WEB_START

START_FAILED=0
for c in xiaoxia-api-staging xiaoxia-worker-staging xiaoxia-web-staging; do
  if ! docker inspect "$c" >/dev/null 2>&1; then
    echo "  FAIL $c: not created"
    START_FAILED=$((START_FAILED + 1))
  else
    state=$(docker inspect -f '{{.State.Status}}' "$c")
    if [ "$state" = "running" ] || [ "$state" = "starting" ]; then
      echo "  OK $c: $state"
    else
      echo "  FAIL $c: $state"
      docker logs --tail 20 "$c" 2>/dev/null || true
      START_FAILED=$((START_FAILED + 1))
    fi
  fi
done

if [ "$START_FAILED" -gt 0 ]; then
  echo "ERROR: $START_FAILED 个容器启动失败"
  rollback
fi

# ---- 并行等待 API 和 Web 健康 ----
echo ""
echo "Waiting for API + Web health (parallel)..."

HEALTH_LOG_DIR="/tmp/staging-health-$$"
mkdir -p "$HEALTH_LOG_DIR"

(
  i=0
  while [ "$i" -lt 40 ]; do
    if curl -sf --max-time 5 http://127.0.0.1:8000/health >/dev/null 2>&1; then
      echo "API healthy after $((i * 3))s"
      exit 0
    fi
    i=$((i + 1))
    sleep 3
  done
  echo "API FAILED after 120s"
  exit 1
) > "$HEALTH_LOG_DIR/api.log" 2>&1 &
PID_API_HEALTH=$!

(
  i=0
  while [ "$i" -lt 15 ]; do
    if curl -sf --max-time 5 http://127.0.0.1:3001/ >/dev/null 2>&1; then
      echo "Web healthy after $((i * 2))s"
      exit 0
    fi
    i=$((i + 1))
    sleep 2
  done
  echo "Web FAILED after 30s"
  exit 1
) > "$HEALTH_LOG_DIR/web.log" 2>&1 &
PID_WEB_HEALTH=$!

set +e
wait $PID_API_HEALTH
API_EXIT=$?
wait $PID_WEB_HEALTH
WEB_EXIT=$?
set -e

echo ""
echo "健康检查结果:"
API_OK=0
WEB_OK=0
if [ "$API_EXIT" -eq 0 ]; then
  echo "  OK API: $(cat "$HEALTH_LOG_DIR/api.log")"
  API_OK=1
else
  echo "  FAIL API: 120s未就绪"
  docker logs --tail 50 xiaoxia-api-staging
fi

if [ "$WEB_EXIT" -eq 0 ]; then
  echo "  OK Web: $(cat "$HEALTH_LOG_DIR/web.log")"
  WEB_OK=1
else
  echo "  FAIL Web: 30s未就绪"
  docker logs --tail 30 xiaoxia-web-staging
fi

rm -rf "$HEALTH_LOG_DIR"

if [ "$API_OK" -eq 0 ] || [ "$WEB_OK" -eq 0 ]; then
  echo ""
  echo "ERROR: 健康检查失败"
  rollback
fi

echo "Cleaning up old images..."
docker image prune -af --filter "until=168h" 2>/dev/null || true
docker builder prune -af --filter "until=168h" 2>/dev/null || true

echo ""
echo "=== Staging deployment complete (并行优化版) ==="
echo "API:      http://127.0.0.1:8000"
echo "Web:      http://127.0.0.1:3001"
echo "Version:  $IMAGE_TAG"
docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Image}}" | grep staging
