#!/bin/bash
set -eu

# ============================================================
# 生产环境一键回滚脚本
#
# 用法:
#   IMAGE_TAG=<版本号> REGISTRY_TOKEN=<token> sh rollback_production.sh
#
# 功能:
#   1. 拉取指定版本镜像
#   2. 数据库回滚到对应版本（alembic downgrade）
#   3. 重启 api/worker/web 三个服务
#   4. 健康检查确认服务正常
#
# 环境变量:
#   IMAGE_TAG           - 要回滚到的版本标签（必填）
#   REGISTRY_TOKEN      - Registry 访问 token（可选）
#   SKIP_DB_ROLLBACK    - 跳过数据库回滚（1=跳过，默认不跳过）
#   DB_ROLLBACK_REV     - 数据库回滚到的版本（默认自动用镜像里的 head）
# ============================================================

IMAGE_TAG="${IMAGE_TAG:-}"
REGISTRY="${REGISTRY:-git.xiaoxiajianji.com/xiaoxia/xiaoxia-saas}"
REGISTRY_USER="${REGISTRY_USER:-xiaoxia}"
REGISTRY_TOKEN="${REGISTRY_TOKEN:-}"
SKIP_DB_ROLLBACK="${SKIP_DB_ROLLBACK:-0}"
DB_ROLLBACK_REV="${DB_ROLLBACK_REV:-}"

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

if [ -z "$IMAGE_TAG" ]; then
  echo "❌ IMAGE_TAG 是必填参数"
  echo "用法: IMAGE_TAG=v0.1.125 REGISTRY_TOKEN=xxx sh rollback_production.sh"
  exit 1
fi

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

echo "=========================================="
echo "  生产环境回滚  →  $IMAGE_TAG"
echo "=========================================="
echo ""

# 先获取当前版本
CURRENT_VERSION=""
if docker inspect xiaoxia-api-production >/dev/null 2>&1; then
  CURRENT_VERSION=$(docker inspect --format '{{ index .Config.Env 0 }}' xiaoxia-api-production 2>/dev/null | grep APP_VERSION | cut -d= -f2 || echo "unknown")
fi
echo "当前版本: ${CURRENT_VERSION:-unknown}"
echo "回滚目标: $IMAGE_TAG"
echo ""

# 确认
read -p "⚠️  确认要回滚生产环境到 $IMAGE_TAG 吗？(输入 YES 确认): " confirm
if [ "$confirm" != "YES" ]; then
  echo "已取消"
  exit 0
fi
echo ""

# ----- 登录 Registry -----
if [ -n "$REGISTRY_TOKEN" ]; then
  echo "登录 Registry: $REGISTRY"
  REGISTRY_HOST=$(echo "$REGISTRY" | cut -d/ -f1)
  printf '%s' "$REGISTRY_TOKEN" | docker login "$REGISTRY_HOST" -u "$REGISTRY_USER" --password-stdin 2>/dev/null || {
    echo "WARN: docker login failed, will try to pull anyway"
  }
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}"

LOCAL_API="xiaoxia-saas-api:${IMAGE_TAG}"
LOCAL_WORKER="xiaoxia-saas-worker:${IMAGE_TAG}"
LOCAL_WEB="xiaoxia-saas-web:${IMAGE_TAG}"

echo "Pulling API image..."
docker pull "$REGISTRY_API"
echo "Pulling Worker image..."
docker pull "$REGISTRY_WORKER"
echo "Pulling Web image..."
docker pull "$REGISTRY_WEB"

docker tag "$REGISTRY_API" "$LOCAL_API"
docker tag "$REGISTRY_WORKER" "$LOCAL_WORKER"
docker tag "$REGISTRY_WEB" "$LOCAL_WEB"
echo "所有镜像拉取完成"
echo ""

# ----- 检查基础设施容器 -----
echo "检查基础设施容器..."
for c in xiaoxia-postgres-production xiaoxia-redis-production; 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-production 2>/dev/null || true
echo ""

# ----- 数据库回滚 -----
if [ "$SKIP_DB_ROLLBACK" = "1" ]; then
  echo "⏭️  跳过数据库回滚（SKIP_DB_ROLLBACK=1）"
else
  echo "🔄 执行数据库回滚..."
  if [ -n "$DB_ROLLBACK_REV" ]; then
    # 回滚到指定版本
    echo "回滚到版本: $DB_ROLLBACK_REV"
    docker run --rm \
      --env-file "$ENV_FILE" \
      --network xiaoxia-net-production \
      -e APP_ENV=production \
      "$LOCAL_API" sh -c "cd /app && alembic downgrade $DB_ROLLBACK_REV"
  else
    # 用目标镜像的 alembic head 来判断是否需要回滚
    # 先检查当前DB版本和目标版本的关系
    echo "检测数据库当前版本与目标版本..."
    CURRENT_DB_REV=$(docker run --rm \
      --env-file "$ENV_FILE" \
      --network xiaoxia-net-production \
      -e APP_ENV=production \
      "$LOCAL_API" sh -c "cd /app && alembic current" 2>&1 | tail -1 | awk '{print $1}')
    TARGET_DB_HEAD=$(docker run --rm \
      "$LOCAL_API" sh -c "cd /app && alembic head" 2>&1 | tail -1 | awk '{print $1}')

    echo "当前 DB 版本: ${CURRENT_DB_REV:-unknown}"
    echo "目标 DB 版本: ${TARGET_DB_HEAD:-unknown}"

    if [ "$CURRENT_DB_REV" = "$TARGET_DB_HEAD" ]; then
      echo "✅ 数据库版本与目标版本一致，无需回滚"
    else
      echo "⚠️  数据库版本不一致，尝试回滚..."
      echo "注意：自动回滚可能无法正确处理，请确认 DB_ROLLBACK_REV 参数"
      echo "如果需要跳过数据库回滚，请设置 SKIP_DB_ROLLBACK=1"
      exit 1
    fi
  fi
  echo "数据库回滚完成"
fi
echo ""

# ----- 停止旧容器 -----
echo "停止旧容器..."
docker rm -f xiaoxia-api-production 2>/dev/null || true
docker rm -f xiaoxia-worker-production 2>/dev/null || true
docker rm -f xiaoxia-web-production 2>/dev/null || true
echo ""

# ----- 启动新容器 -----
LOG_OPTS="--log-driver json-file --log-opt max-size=50m --log-opt max-file=3"

echo "启动 API 容器..."
docker run -d \
  --name xiaoxia-api-production \
  --env-file "$ENV_FILE" \
  --network xiaoxia-net-production \
  -p 127.0.0.1:8001:8000 \
  -e APP_ENV=production \
  -e APP_VERSION="$IMAGE_TAG" \
  -e GENERATED_FILES_DIR=/app/generated \
  -e GENERATED_FILES_URL_PREFIX=/generated-files \
  -e PUBLIC_API_BASE_URL=https://api.xiaoxiajianji.com \
  -v "$GENERATED_DIR:/app/generated" \
  --restart unless-stopped \
  --cpus 2 \
  --memory 2g \
  --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 \
  "$LOCAL_API"

echo "启动 Worker 容器..."
docker run -d \
  --name xiaoxia-worker-production \
  --env-file "$ENV_FILE" \
  --network xiaoxia-net-production \
  -e APP_ENV=production \
  -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://api.xiaoxiajianji.com \
  -v "$GENERATED_DIR:/app/generated" \
  --restart unless-stopped \
  --cpus 2 \
  --memory 2g \
  --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 \
  "$LOCAL_WORKER"

# Web legacy assets
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"
  echo "Web 容器: legacy assets 已挂载"
else
  echo "Web 容器: 没有 legacy assets"
fi

echo "启动 Web 容器..."
docker run -d \
  --name xiaoxia-web-production \
  --network xiaoxia-net-production \
  -p 127.0.0.1:3002:80 \
  --restart unless-stopped \
  --cpus 0.5 \
  --memory 512m \
  $LEGACY_VOLUME \
  --health-cmd "wget --spider -q http://127.0.0.1:80" \
  --health-interval 30s \
  --health-timeout 5s \
  --health-retries 3 \
  $LOG_OPTS \
  "$LOCAL_WEB"
echo ""

# ----- 健康检查 -----
echo "等待 API 健康..."
i=0
while [ "$i" -lt 40 ]; do
  if curl -sf --max-time 5 http://127.0.0.1:8001/health >/dev/null 2>&1; then
    echo "API is healthy!"
    break
  fi
  i=$((i + 1))
  echo "  Waiting... ($i/40)"
  sleep 3
done

if [ "$i" -ge 40 ]; then
  echo "❌ API 在 120s 内未就绪"
  docker logs --tail 50 xiaoxia-api-production
  exit 1
fi

echo "等待 Web 健康..."
i=0
while [ "$i" -lt 15 ]; do
  if curl -sf --max-time 5 http://127.0.0.1:3002/ >/dev/null 2>&1; then
    echo "Web is healthy!"
    break
  fi
  i=$((i + 1))
  echo "  Waiting... ($i/15)"
  sleep 2
done

if [ "$i" -ge 15 ]; then
  echo "❌ Web 在 30s 内未就绪"
  docker logs --tail 30 xiaoxia-web-production
  exit 1
fi
echo ""

# ----- 清理 -----
echo "清理旧镜像..."
docker image prune -af --filter "until=168h" 2>/dev/null || true
docker builder prune -af --filter "until=168h" 2>/dev/null || true

echo ""
echo "=========================================="
echo "  ✅ 生产环境回滚完成"
echo "=========================================="
echo "API:      http://127.0.0.1:8001"
echo "Web:      http://127.0.0.1:3002"
echo "Version:  $IMAGE_TAG"
docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Image}}" | grep production