#!/usr/bin/env bash
# ============================================================
# llmsproxy 原子化滚动部署脚本
#
# 策略：
#   1. /tmp (tmpfs) 构建，/usr/local/bin (ext4) 部署
#   2. 构建后先 copy 到目标文件系统（.llmsproxy.tmp），
#      再 rename(2) 做原子替换（同一文件系统内 mv 是原子的）
#   3. runtime.json 不替换（保留运行时密钥、审计状态）
#   4. systemctl restart → 服务中断仅在 Go 进程重启瞬间 (<500ms)
#   5. 健康检查确认新进程正常
#
# 注意事项：
#   - 本服务依赖自身（推理依赖），restart 前后会短暂中断，
#     但重试机制（systemd RestartSec=5 + 客户端 retry）可自愈
#   - 适配器文件替换后需重启才能生效（VM 启动时批量加载）
# ============================================================
set -euo pipefail

# ---------- 配置 ----------
readonly TARGET_BIN="/usr/local/bin/llmsproxy"
readonly TARGET_ADAPTERS="/etc/llmsproxy/adapters"
readonly REPO_DIR="/home/program/llmsproxy"
readonly HEALTH_URL="http://127.0.0.1:8081/v1/models"
readonly BIN_NAME="llmsproxy"

# 构建目录在 /tmp（tmpfs 更快）
readonly TMP_DIR="/tmp/llmsproxy-deploy"
# 同文件系统暂存文件（用于原子 rename）
readonly STAGING_BIN="${TARGET_BIN}.staging"
# 回滚备份：健康检查失败时从这里恢复旧二进制
readonly BACKUP_BIN="${TARGET_BIN}.bak"

# ---------- 日志 ----------
log()  { printf '\033[1;34m[%s]\033[0m %s\n' "$(date +%H:%M:%S)" "$*"; }
warn() { printf '\033[1;33m[WARN]\033[0m %s\n' "$*"; }
err()  { printf '\033[1;31m[ERR]\033[0m %s\n' "$*" >&2; }
fail() { err "$@"; exit 1; }

# ---------- 清理 ----------
cleanup() {
    log "清理临时文件"
    rm -rf "$TMP_DIR" "$STAGING_BIN"
}
trap cleanup EXIT

# ---------- 前置检查 ----------
precheck() {
    log "前置检查"
    [[ -f "$REPO_DIR/go.mod" ]] || fail "仓库目录 $REPO_DIR 不存在"
    command -v go       &>/dev/null || fail "未找到 go"
    command -v systemctl &>/dev/null || fail "未找到 systemctl"
    [[ $(id -u) -eq 0 ]] || fail "需要 root 权限"

    OLD_SIZE=$(stat -c%s "$TARGET_BIN" 2>/dev/null || echo 0)
    OLD_SHA256=$(sha256sum "$TARGET_BIN" 2>/dev/null | awk '{print $1}')

    log "旧二进制: $TARGET_BIN (${OLD_SIZE} bytes, sha256=${OLD_SHA256:0:16})"
    log "go: $(go version)"
}

# ---------- 备份旧二进制（用于回滚）----------
backup_old_binary() {
    log "备份旧二进制到 $BACKUP_BIN"
    if [[ -f "$TARGET_BIN" ]]; then
        cp -f "$TARGET_BIN" "$BACKUP_BIN"
        log "✓ 旧二进制已备份 (sha256=$(sha256sum "$BACKUP_BIN" | awk '{print $1}' | cut -c1-16))"
    else
        warn "旧二进制不存在，跳过备份（无回滚点）"
    fi
}

# ---------- 构建 ----------
build() {
    log "构建新二进制"
    rm -rf "$TMP_DIR"
    mkdir -p "$TMP_DIR/bin"

    cd "$REPO_DIR"
    CGO_ENABLED=1 go build -tags luajit \
        -trimpath \
        -ldflags="-s -w" \
        -o "$TMP_DIR/bin/$BIN_NAME" \
        ./cmd/llmsproxy

    NEW_BIN="$TMP_DIR/bin/$BIN_NAME"
    NEW_SIZE=$(stat -c%s "$NEW_BIN")

    if [[ "$NEW_SIZE" -eq "$OLD_SIZE" ]]; then
        NEW_SHA256=$(sha256sum "$NEW_BIN" | awk '{print $1}')
        if [[ "$NEW_SHA256" == "$OLD_SHA256" ]]; then
            warn "新二进制与旧版本完全一致（sha256 相同），跳过部署"
            exit 0
        fi
    fi

    log "新二进制: $NEW_BIN (${NEW_SIZE} bytes)"

    # 验证：确认包含关键修复
    local checks=(
        "mergeUsage"
        "prompt_tokens"
        "completion_tokens"
        "message_start"
    )
    for c in "${checks[@]}"; do
        if strings "$NEW_BIN" | grep -qF "$c"; then
            log "  ✓ 新二进制包含: $c"
        else
            warn "  ⚠ 新二进制不包含: $c（非预期但非致命）"
        fi
    done
}

# ---------- 原子替换二进制 ----------
# 原理：
#   /tmp 是 tmpfs，/usr/local/bin 是 ext4——不同文件系统 rename(2) 不原子。
#   因此先 copy 到目标文件系统上的 .staging 文件，再 rename(2) 到目标路径。
#   rename(2) 在同一文件系统内是元数据级操作，对正在读写的进程安全：
#   - 旧 inode 的 fd 继续有效（进程持有旧文件，不会读到截断内容）
#   - 新进程打开 path 看到新 inode
#   - 旧 inode 在最后一个 fd 关闭后释放
atomic_replace_binary() {
    log "原子替换二进制"

    local new_bin_tmp="$TMP_DIR/bin/$BIN_NAME"
    [[ -f "$new_bin_tmp" ]] || fail "新二进制文件不存在: $new_bin_tmp"

    # Step 1: copy 到目标文件系统（非原子，但目标文件唯一）
    log "  copy → $STAGING_BIN (同文件系统暂存)"
    cp -f "$new_bin_tmp" "$STAGING_BIN"
    chmod 0755 "$STAGING_BIN"

    # Step 2: 原子 rename（同一文件系统，元数据级操作）
    log "  rename → $TARGET_BIN (原子)"
    mv -f "$STAGING_BIN" "$TARGET_BIN"

    # 验证
    NEW_SHA256=$(sha256sum "$TARGET_BIN" | awk '{print $1}')
    log "✓ 二进制已原子替换 (sha256=${NEW_SHA256:0:16})"
}

# ---------- 同步适配器 ----------
sync_adapters() {
    log "同步适配器"

    # 只覆盖内置适配器（与 internal/lua/adapters 同名的文件），
    # 保留目录里其它运行时上传的 .lua —— WebUI 上传的自定义适配器
    # 必须跨部署存活，清空会静默移除线上源依赖的适配器。
    SRC_ADAPTERS="$REPO_DIR/internal/lua/adapters"
    [[ -d "$SRC_ADAPTERS" ]] || fail "内嵌适配器目录不存在: $SRC_ADAPTERS"

    # 备份旧的适配器文件（含任何运行时覆盖版本，如 opencode.lua）
    BACKUP_DIR="/etc/llmsproxy/adapters.bak.$(date +%Y%m%d%H%M%S)"
    if [[ -d "$TARGET_ADAPTERS" ]]; then
        mkdir -p "$BACKUP_DIR"
        cp -rf "$TARGET_ADAPTERS/"*.lua "$BACKUP_DIR/" 2>/dev/null || true
        log "  旧适配器已备份到 $BACKUP_DIR"
    fi

    cp -f "$SRC_ADAPTERS/"*.lua "$TARGET_ADAPTERS/"
    chmod 0644 "$TARGET_ADAPTERS/"*.lua

    # 校验：每个适配器必须包含 usage 透传修复（opencode/openai/deepseek 等应有 'uses'）
    local required_files=(openai deepseek anthropic gemini ollama opencode github groq kimicode mistral sensenova agentrouter)
    for f in "${required_files[@]}"; do
        [[ -f "$TARGET_ADAPTERS/$f.lua" ]] || warn "  缺少适配器: $f.lua"
    done
    local usage_fixed=0
    for f in openai deepseek github groq kimicode mistral opencode; do
        grep -qE '\buses\b' "$TARGET_ADAPTERS/$f.lua" 2>/dev/null && ((usage_fixed++)) || true
    done
    log "  usage 透传修复适配器: ${usage_fixed}/7 (openai/deepseek/github/groq/kimicode/mistral/opencode)"
    log "✓ 适配器已同步到 $TARGET_ADAPTERS/"
}

# ---------- 重启服务 ----------
restart_service() {
    log "重启服务（中断 <500ms）"
    systemctl restart llmsproxy.service || fail "systemd restart 失败"
    log "✓ llmsproxy.service restarted"
}

# ---------- 健康检查 ----------
healthcheck() {
    log "健康检查"
    local attempts=20
    local wait_sec=1
    local attempt=1

    while (( attempt <= attempts )); do
        # 用 curl 检查（200=有 key 未鉴权但服务正常，401=需要认证，都说明服务在运行）
        local code
        code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 3 "$HEALTH_URL" 2>/dev/null || echo "000")
        if [[ "$code" == "200" || "$code" == "401" ]]; then
            log "✓ 服务健康 (status=$code, ${attempt}/${attempts})"
            return 0
        fi
        log "  等待中 ($attempt/$attempts, status=$code)"
        sleep "$wait_sec"
        (( attempt++ ))
    done

    warn "健康检查未通过，请检查服务状态: systemctl status llmsproxy.service"
    return 1
}

# ---------- 回滚 ----------
# 健康检查失败时恢复旧二进制并重启，保证推理服务可用。
# 这是最后一道保险：即使新版本启动失败（编译错误、配置不兼容、
# panic on start 等），也能让旧版本继续服务，避免推理能力丢失。
rollback() {
    err "健康检查失败，启动回滚流程"
    if [[ ! -f "$BACKUP_BIN" ]]; then
        err "无回滚备份 ($BACKUP_BIN 不存在)，无法自动恢复"
        err "请手动检查: systemctl status llmsproxy.service && journalctl -u llmsproxy -n 50"
        return 1
    fi

    log "恢复旧二进制"
    cp -f "$BACKUP_BIN" "$TARGET_BIN"
    chmod 0755 "$TARGET_BIN"

    log "重启服务（旧版本）"
    systemctl restart llmsproxy.service || {
        err "重启失败，服务可能完全宕机"
        err "请手动检查: systemctl status llmsproxy.service"
        return 1
    }

    # 回滚后再健康检查（只查 1 次，失败就认了）
    sleep 2
    local code
    code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 "$HEALTH_URL" 2>/dev/null || echo "000")
    if [[ "$code" == "200" || "$code" == "401" ]]; then
        log "✓ 回滚成功，旧版本已恢复服务 (status=$code)"
        log "  新版本部署失败，请检查构建/配置后重试"
        return 0
    else
        err "回滚后服务仍未恢复 (status=$code)"
        err "服务可能已完全宕机，请手动介入"
        return 1
    fi
}

# ---------- 打印部署摘要 ----------
print_summary() {
    log "═════════════════════════════════════════════"
    log "部署完成"
    log "  二进制: $TARGET_BIN"
    log "  适配器: $TARGET_ADAPTERS"
    log "═════════════════════════════════════════════"
}

# ---------- 主流程 ----------
main() {
    log "═══ llmsproxy 原子部署开始 ═══"
    precheck
    build
    sync_adapters
    backup_old_binary
    atomic_replace_binary
    restart_service
    if ! healthcheck; then
        rollback || true
        exit 1
    fi
    print_summary
}

main "$@"