#!/bin/bash
set -e

# ============================================
# Kiro Shell 启动脚本 v1.0
# ============================================

RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

check_env() {
    local missing=0

    for var in WORKER_URL AUTH_TOKEN ENCRYPTION_KEY API_KEY; do
        if [ -z "${!var}" ]; then
            log_error "缺少环境变量: $var"
            missing=1
        fi
    done

    if [ $missing -eq 1 ]; then
        exit 1
    fi

    log_success "环境变量检查通过"
}

download_credentials() {
    log_info "从 Worker 下载凭据..."

    local response http_code body
    response=$(curl -sS --tlsv1.2 --connect-timeout 10 --max-time 20 \
        -H "Authorization: Bearer $AUTH_TOKEN" \
        -w "\nHTTP_CODE:%{http_code}\n" \
        "${WORKER_URL}/api/credentials")

    http_code=$(echo "$response" | awk -F'HTTP_CODE:' 'END{print $2}' | tr -d '\r')
    body=$(echo "$response" | sed '/HTTP_CODE:/d')

    if [ "$http_code" != "200" ]; then
        log_error "下载凭据失败 (HTTP $http_code): $body"
        exit 1
    fi

    ENCRYPTED=$(echo "$body" | jq -r '.encrypted')
    IV=$(echo "$body" | jq -r '.iv')
    COUNT=$(echo "$body" | jq -r '.count')

    if [ "$ENCRYPTED" = "null" ] || [ "$IV" = "null" ]; then
        log_error "凭据格式错误: $body"
        exit 1
    fi

    log_success "凭据下载成功 (共 $COUNT 个)"
}

decrypt_credentials() {
    log_info "解密凭据..."

    local iv_hex decrypted
    iv_hex=$(echo "$IV" | base64 -d | xxd -p | tr -d '\n')

    decrypted=$(echo "$ENCRYPTED" | base64 -d | \
        openssl enc -aes-256-cbc -d \
        -K "$ENCRYPTION_KEY" \
        -iv "$iv_hex" 2>/dev/null)

    if [ $? -ne 0 ] || [ -z "$decrypted" ]; then
        log_error "解密失败，请检查 ENCRYPTION_KEY"
        exit 1
    fi

    if ! echo "$decrypted" | jq . > /dev/null 2>&1; then
        log_error "解密后的数据不是有效 JSON"
        exit 1
    fi

    echo "$decrypted" > /app/credentials.json
    log_success "凭据解密成功"
}

generate_config() {
    log_info "生成配置文件..."

    cat > /app/config.json << EOF
{
    "host": "0.0.0.0",
    "port": 7860,
    "apiKey": "$API_KEY",
    "region": "${REGION:-us-east-1}",
    "adminApiKey": "${ADMIN_API_KEY:-admin12345}"
}
EOF

    log_success "配置文件生成完成"
}

start_kiro() {
    log_info "启动 kiro-rs..."
    echo ""
    echo "========================================"
    echo "  Kiro Proxy 已启动"
    echo "  监听地址: 0.0.0.0:7860"
    echo "========================================"
    echo ""

    exec /app/kiro-rs -c /app/config.json --credentials /app/credentials.json
}

main() {
    echo ""
    echo "========================================"
    echo "  Kiro Shell v1.0"
    echo "========================================"
    echo ""

    check_env
    download_credentials
    decrypt_credentials
    generate_config
    start_kiro
}

main