#!/bin/bash
# ============================================================
# backup_api.sh
# 自动从 conf/config.json 读取数据库配置，扫描所有模块的
# models 目录提取表名，导出为 SQL 文件。
#
# 运行目录: sage.py 所在目录 (即 repos/sage/)
#
# 用法:
#   cd /path/to/sage && bash bin/backup_api.sh [options]
#
# 选项:
#   -p PASSWORD    直接指定密码 (覆盖config.json中的加密密码)
#   -h HOST        覆盖数据库主机
#   -o OUTPUT_DIR  输出目录 (默认: ./sql_dumps)
#   --no-data      只导出表结构，不导出数据
#   --help         显示帮助
# ============================================================

set -euo pipefail

# --- 定位工作目录 (sage.py 所在目录) ---
SAGE_DIR="$(cd "$(dirname "$0")/.." && pwd)"
REPOS_DIR="$(cd "$SAGE_DIR/.." && pwd)"
CONFIG_FILE="$SAGE_DIR/conf/config.json"

if [[ ! -f "$CONFIG_FILE" ]]; then
    echo "错误: 找不到配置文件 $CONFIG_FILE"
    echo "请在 sage.py 所在目录下运行此脚本"
    exit 1
fi

# 默认参数
OUTPUT_DIR="./sql_dumps"
NO_DATA=""
PASSWORD_OVERRIDE=""
HOST_OVERRIDE=""

usage() {
    sed -n '2,/^# ===/p' "$0" | grep '^#' | sed 's/^# \?//'
    exit 0
}

# 解析参数
while [[ $# -gt 0 ]]; do
    case "$1" in
        -p) PASSWORD_OVERRIDE="$2"; shift 2 ;;
        -h) HOST_OVERRIDE="$2"; shift 2 ;;
        -o) OUTPUT_DIR="$2"; shift 2 ;;
        --no-data) NO_DATA="--no-data"; shift ;;
        --help) usage ;;
        *) echo "未知参数: $1"; exit 1 ;;
    esac
done

# --- 从 config.json 提取数据库配置并解密密码 ---
read_config() {
    # Use the sage venv if available, otherwise system python
    local PYTHON="python3"
    if [[ -x "$SAGE_DIR/py3/bin/python3" ]]; then
        PYTHON="$SAGE_DIR/py3/bin/python3"
    fi

    $PYTHON -c "
import re, json, sys, os

# 读取并修复 config.json (可能有非标准 JSON)
text = open('$CONFIG_FILE').read()
# 移除无 key 的裸对象 (如 hot_reload)
text = re.sub(r',\s*\{[^{}]*\"hot_reload\"[^{}]*\{[^{}]*\}[^{}]*\}', '', text)
config = json.loads(text)

# 获取数据库配置 (取第一个数据库)
db_name = list(config['databases'].keys())[0]
db_cfg = config['databases'][db_name]['kwargs']
password_key = config.get('password_key', 'QRIVSRHrthhwyjy176556332')

# 尝试 RC4 解密
try:
    from appPublic.rc4 import unpassword
    decrypted = unpassword(db_cfg['password'], key=password_key)
except ImportError:
    # Fallback: inline RC4 implementation
    import base64
    from hashlib import sha1
    class RC4:
        def __init__(self):
            self.bcoding = 'iso-8859-1'
            self.dcoding = 'utf8'
            self.salt = b'AFUqx9WZuI32lnHk'
        def _crypt(self, data, key):
            x = 0; box = list(range(256))
            for i in range(256):
                x = (x + box[i] + key[i % len(key)]) % 256
                box[i], box[x] = box[x], box[i]
            x = y = 0; out = []
            for char in data:
                x = (x + 1) % 256; y = (y + box[x]) % 256
                box[x], box[y] = box[y], box[x]
                out.append(chr(char ^ box[(box[x] + box[y]) % 256]))
            return ''.join(out).encode(self.bcoding)
        def decode(self, data, key):
            if isinstance(data, str): data = data.encode(self.dcoding)
            key = key.encode(self.bcoding)
            data = base64.b64decode(data)
            a = sha1(key + self.salt); k = a.digest()
            return self._crypt(data[16:], k).decode(self.dcoding)
    rc4 = RC4()
    decrypted = rc4.decode(db_cfg['password'], password_key)
except Exception as e:
    decrypted = ''

# 如果解密失败或为空，使用原始值
if not decrypted:
    decrypted = db_cfg['password']

host = '$HOST_OVERRIDE' if '$HOST_OVERRIDE' else db_cfg.get('host', 'localhost')
port = str(db_cfg.get('port', '3306'))
user = db_cfg.get('user', 'root')
database = db_cfg.get('db', 'sage')

# Shell-safe output
print(f'HOST={host}')
print(f'PORT={port}')
print(f'USER={user}')
print(f'DATABASE={database}')
# Password needs quoting for special chars
print(f\"PASSWORD='{decrypted}'\")
"
}

echo "正在读取配置..."
eval "$(read_config)"

# 命令行密码覆盖
if [[ -n "$PASSWORD_OVERRIDE" ]]; then
    PASSWORD="$PASSWORD_OVERRIDE"
fi

if [[ -z "$PASSWORD" ]]; then
    echo "错误: 密码为空。请使用 -p 参数指定密码，或检查 config.json 中的密码配置"
    exit 1
fi

# 构建 mysqldump 基础命令 (密码通过环境变量传递，避免命令行暴露)
export MYSQL_PWD="$PASSWORD"
MYSQLDUMP_CMD="mysqldump -h${HOST} -P${PORT} -u${USER}"

# --- 各模块排除的表 (大数据量/日志表) ---
declare -A MODULE_EXCLUDES
MODULE_EXCLUDES["llmage"]="llmusage llmusage_accounting_failed llmusage_history"
MODULE_EXCLUDES["uapi"]="uapiset uptask"
MODULE_EXCLUDES["accounting"]="accountingdetail accountinghistory"
MODULE_EXCLUDES["harnessed_agent"]="harnessed_agent_log"
MODULE_EXCLUDES["harnessed_reasoning"]="harnessed_reasoning_log"

# --- 从 models/*.json 提取表名 ---
get_tables() {
    local module_dir="$1"
    shift
    local excludes=("$@")
    local tables=()

    if [[ ! -d "$module_dir/models" ]]; then
        return
    fi

    for f in "$module_dir/models"/*.json; do
        [[ -f "$f" ]] || continue
        local tbl
        tbl=$(python3 -c "
import json, sys
try:
    d = json.load(open('$f'))
    s = d.get('summary', [])
    if isinstance(s, list) and len(s) > 0:
        print(s[0].get('name', ''))
    elif isinstance(s, dict):
        print(s.get('name', ''))
except:
    pass
" 2>/dev/null)
        if [[ -z "$tbl" ]]; then
            continue
        fi
        # 检查排除
        local excluded=false
        for ex in "${excludes[@]+"${excludes[@]}"}"; do
            if [[ "$tbl" == "$ex" ]]; then
                excluded=true
                break
            fi
        done
        if [[ "$excluded" == "false" ]]; then
            tables+=("$tbl")
        fi
    done
    echo "${tables[@]}"
}

# --- 自动发现所有模块 ---
discover_modules() {
    for dir in "$REPOS_DIR"/*/; do
        [[ -d "$dir/models" ]] || continue
        local name
        name=$(basename "$dir")
        # 跳过无模型文件的模块
        if ls "$dir/models/"*.json &>/dev/null; then
            echo "$name"
        fi
    done
}

# --- 创建输出目录 ---
mkdir -p "$OUTPUT_DIR"

TIMESTAMP=$(date +%Y%m%d_%H%M%S)

echo "============================================================"
echo "  Sage 全模块表数据导出"
echo "  时间: $(date '+%Y-%m-%d %H:%M:%S')"
echo "  数据库: ${DATABASE}@${HOST}:${PORT}"
echo "  用户: ${USER}"
echo "  模块目录: ${REPOS_DIR}"
echo "  输出目录: ${OUTPUT_DIR}"
echo "============================================================"

# 发现所有模块
MODULES=($(discover_modules))
echo ""
echo "发现 ${#MODULES[@]} 个模块: ${MODULES[*]}"

TOTAL_TABLES=0
TOTAL_FILES=0
FAILED_MODULES=()

for module in "${MODULES[@]}"; do
    module_dir="$REPOS_DIR/$module"
    echo ""
    echo "--- 模块: $module ---"

    # 获取排除列表
    excludes_str="${MODULE_EXCLUDES[$module]:-}"
    excludes_arr=()
    if [[ -n "$excludes_str" ]]; then
        read -ra excludes_arr <<< "$excludes_str"
    fi

    tables=$(get_tables "$module_dir" "${excludes_arr[@]+"${excludes_arr[@]}"}")

    if [[ -z "$tables" ]]; then
        echo "  跳过: 未找到表定义"
        continue
    fi

    echo "  表: $tables"

    # 生成 SQL 文件
    outfile="${OUTPUT_DIR}/${module}_${TIMESTAMP}.sql"
    table_count=0

    for tbl in $tables; do
        echo -n "  导出 $tbl ... "
        if ${MYSQLDUMP_CMD} \
            --single-transaction \
            --routines \
            --triggers \
            --set-gtid-purged=OFF \
            --column-names \
            --complete-insert \
            $NO_DATA \
            "$DATABASE" "$tbl" >> "$outfile" 2>/dev/null; then
            echo "OK"
            ((table_count++))
        else
            echo "跳过(表不存在或无权限)"
        fi
    done

    if [[ $table_count -gt 0 ]]; then
        echo "  => $outfile ($table_count 个表)"
        ((TOTAL_TABLES += table_count))
        ((TOTAL_FILES++))
    else
        rm -f "$outfile"
        echo "  => 无有效表数据"
        FAILED_MODULES+=("$module")
    fi
done

# 清理密码环境变量
unset MYSQL_PWD

echo ""
echo "============================================================"
echo "  导出完成"
echo "  模块数: ${#MODULES[@]}"
echo "  文件数: $TOTAL_FILES"
echo "  表总数: $TOTAL_TABLES"
if [[ ${#FAILED_MODULES[@]} -gt 0 ]]; then
    echo "  无数据模块: ${FAILED_MODULES[*]}"
fi
echo "  输出目录: $(cd "$OUTPUT_DIR" && pwd)"
echo "============================================================"