#!/bin/bash
# OpenClaw Dog 守护程序 - 优化重试版
# 功能：监控 Gateway 状态，自动修复故障
# 特性：智能重试机制（指数退避 + 错误类型适配）

# ==================== 配置 ====================
LOG_DIR="/root/.openclaw/logs"
STATUS_FILE="$LOG_DIR/guardian_status.json"
METRICS_FILE="$LOG_DIR/guardian_metrics.json"
BACKUP_DIR="/root/.openclaw/backup"
CONFIG_FILE="/root/.openclaw/openclaw.json"
LOG_FILE="$LOG_DIR/guardian.log"
PID_FILE="$LOG_DIR/guardian.pid"

GATEWAY_PORT=18789
GATEWAY_PROCESS="openclaw.*gateway"

# 锁清理配置
SESSION_LOCK_DIR="/root/.openclaw/agents/main/sessions"
LOCK_CLEAN_MIN=10

# 检查配置
FAILURE_THRESHOLD=3
MAX_RETRIES=3           # 增加重试次数
INITIAL_RETRY_DELAY=10   # 初始重试延迟（秒）
MAX_RETRY_DELAY=120      # 最大重试延迟（秒）
BACKOFF_FACTOR=2         # 指数退避因子

FAILURE_THRESHOLD=3
STARTUP_WAIT=30
MAX_BACKUPS=10

# ==================== 智能重试函数 ====================

# 计算下一次重试延迟（指数退避）
calculate_retry_delay() {
    local attempt=$1
    local delay=$INITIAL_RETRY_DELAY
    
    # 指数退避: delay = initial * (factor ^ (attempt-1))
    for i in $(seq 2 $attempt); do
        delay=$((delay * BACKOFF_FACTOR))
    done
    
    # 不超过最大延迟
    [ $delay -gt $MAX_RETRY_DELAY ] && delay=$MAX_RETRY_DELAY
    
    echo $delay
}

# 根据错误类型获取建议的重试延迟
get_error_based_delay() {
    local error_msg=$1
    
    # 网络错误 - 较短延迟
    if echo "$error_msg" | grep -qi "network\|connection\|timeout"; then
        echo 15
    # 权限错误 - 中等延迟
    elif echo "$error_msg" | grep -qi "permission\|denied\|auth"; then
        echo 30
    # 资源不足 - 较长延迟
    elif echo "$error_msg" | grep -qi "memory\|cpu\|resource"; then
        echo 60
    # 默认使用指数退避
    else
        echo 0  # 使用默认指数退避
    fi
}

# 智能重试函数
smart_retry() {
    local max_attempts=$1
    local attempt=1
    local success=0
    
    while [ $attempt -le $max_attempts ]; do
        # 执行检查
        if check_gateway; then
            success=1
            break
        fi
        
        # 获取错误信息
        local error_msg=$(check_gateway 2>&1)
        
        # 计算延迟
        local delay
        local specific_delay=$(get_error_based_delay "$error_msg")
        
        if [ $specific_delay -gt 0 ]; then
            delay=$specific_delay
        else
            delay=$(calculate_retry_delay $attempt)
        fi
        
        log "RETRY" "Attempt $attempt/$max_attempts failed, waiting ${delay}s..."
        sleep $delay
        
        attempt=$((attempt + 1))
    done
    
    return $success
}

# ==================== 会话锁清理 ====================

clean_stale_locks() {
    local cleaned=0
    local now=$(date +%s)
    
    find "$SESSION_LOCK_DIR" -name "*.lock" -type f 2>/dev/null | while read lockfile; do
        local age=$(($(date +%s) - $(stat -c %Y "$lockfile" 2>/dev/null || echo 0)))
        local age_min=$((age / 60))
        
        if [ $age_min -gt $LOCK_CLEAN_MIN ]; then
            local pid=$(grep -o '"pid":[0-9]*' "$lockfile" 2>/dev/null | cut -d':' -f2)
            
            if [ -n "$pid" ]; then
                if ! ps -p "$pid" > /dev/null 2>&1; then
                    rm -f "$lockfile"
                    log "CLEAN" "Removed stale lock: $(basename $lockfile) (age: ${age_min}min)"
                    cleaned=$((cleaned + 1))
                fi
            else
                rm -f "$lockfile"
                cleaned=$((cleaned + 1))
            fi
        fi
    done
    
    [ $cleaned -gt 0 ] && log "CLEAN" "Cleaned $cleaned stale locks"
    return 0
}

# ==================== 并发函数 ====================

check_single_instance() {
    if [ -f "$PID_FILE" ]; then
        local old_pid=$(cat "$PID_FILE" 2>/dev/null)
        if [ -n "$old_pid" ] && kill -0 "$old_pid" 2>/dev/null; then
            return 1
        fi
    fi
    echo $$ > "$PID_FILE"
    return 0
}

cleanup_pid() {
    rm -f "$PID_FILE"
}

# ==================== 日志 ====================
log() {
    local level=$1
    shift
    local msg="$@"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [$level] $msg" >> "$LOG_FILE"
    echo "[$timestamp] [$level] $msg"
}

log_header() {
    log "=================================================="
    log "  OpenClaw Dog - Gateway Guardian (Smart Retry)"
    log "=================================================="
}

log_status() {
    log "  Gateway: $1 | Fail: $2/$3 | Retry: $4/$5"
}

log_summary() {
    local runs=$1 rate=$2 fails=$3 restarts=$4
    log "--------------------------------------------------"
    log "  STATS: Runs=$runs | Success=$rate% | Fail=$fails | Restart=$restarts"
    log "--------------------------------------------------"
}

# ==================== 核心函数 ====================
get_timestamp_ms() { date +%s%3N; }
get_elapsed_ms() { echo $(($(date +%s%3N) - $1)); }
get_system_load() { uptime | awk -F'load average:' '{print $2}' | awk '{print $1}' | tr -d ','; }
get_gateway_pid() { pgrep -f "$GATEWAY_PROCESS" | head -1; }

check_port() { ss -tlnp 2>/dev/null | grep -q ":$GATEWAY_PORT"; }

check_http() {
    curl -s -o /dev/null -w "%{http_code}" --connect-timeout 5 "http://127.0.0.1:$GATEWAY_PORT/" 2>/dev/null
}

check_gateway() {
    local issues=0
    [ -z "$(get_gateway_pid)" ] && ((issues++))
    check_port || ((issues++))
    [ "$(check_http)" != "200" ] && ((issues++))
    return $issues
}

# ==================== 状态管理 ====================
read_status() { [ -f "$STATUS_FILE" ] && cat "$STATUS_FILE" || echo '{"failure_count":0,"retry_count":0}'; }
get_value() { echo "$1" | grep -o "\"$2\":[0-9]*" | cut -d':' -f2; }

save_status() {
    cat > "$STATUS_FILE" << EOF
{
    "failure_count": $1,
    "retry_count": $2,
    "last_check": "$(date '+%Y-%m-%d %H:%M:%S')"
}
EOF
}

record_metrics() {
    local time=$1 result=$2
    local runs=1 fails=0 restarts=0
    
    if [ -f "$METRICS_FILE" ]; then
        runs=$(($(grep -o '"total_runs":[0-9]*' "$METRICS_FILE" | cut -d':' -f2) + 1))
        fails=$(grep -o '"total_failures":[0-9]*' "$METRICS_FILE" | cut -d':' -f2)
        restarts=$(grep -o '"total_restarts":[0-9]*' "$METRICS_FILE" | cut -d':' -f2)
        [ -z "$fails" ] && fails=0
        [ -z "$restarts" ] && restarts=0
        [ "$result" != "success" ] && fails=$((fails + 1))
        [ "$result" == "restarted" ] && restarts=$((restarts + 1))
    fi
    
    local rate=$((100 - (fails * 100 / runs)))
    
    cat > "$METRICS_FILE" << EOF
{
    "total_runs": $runs,
    "total_failures": $fails,
    "total_restarts": $restarts,
    "success_rate": $rate,
    "avg_execution_ms": $time,
    "last_run": "$(date '+%Y-%m-%d %H:%M:%S')",
    "last_result": "$result"
}
EOF
    echo "$runs,$rate,$fails,$restarts"
}

# ==================== 维修模块 ====================
backup_config() {
    mkdir -p "$BACKUP_DIR"
    local bf="${BACKUP_DIR}/openclaw_$(date '+%Y%m%d_%H%M%S').json"
    [ -f "$CONFIG_FILE" ] && cp "$CONFIG_FILE" "$bf" && log "REPAIR" "Backup: $(basename $bf)"
    echo "$bf"
}

list_backups() { ls -t "${BACKUP_DIR}"/openclaw_*.json 2>/dev/null; }

restore_config() {
    [ -f "$1" ] && cp "$1" "$CONFIG_FILE" && log "REPAIR" "Restored: $(basename $1)"
}

restart_gateway() {
    log "REPAIR" "Restarting Gateway..."
    systemctl --user restart openclaw-gateway.service
}

repair_gateway() {
    log "REPAIR" "=== ENTER REPAIR MODE ==="
    
    local current=$(backup_config)
    local backups=$(list_backups)
    local count=$(echo "$backups" | wc -l)
    log "REPAIR" "Found $count backups"
    
    local attempt=1 success=0
    for backup in $backups; do
        [ $attempt -gt $MAX_BACKUPS ] && break
        log "REPAIR" "Try #$attempt: $(basename $backup)"
        restore_config "$backup"
        restart_gateway
        sleep 10
        if check_gateway; then
            success=1
            break
        fi
        attempt=$((attempt + 1))
    done
    
    if [ $success -eq 1 ]; then
        log "REPAIR" "SUCCESS: Gateway recovered"
        return 0
    else
        log "REPAIR" "FAILED: Restore original"
        restore_config "$current"
        restart_gateway
        log "REPAIR" "!!! REPAIR COMPLETELY FAILED !!!"
        return 1
    fi
}

# ==================== 主程序 ====================
main() {
    if ! check_single_instance; then
        log "WARN" "Another instance is running, exiting"
        exit 1
    fi
    
    trap cleanup_pid EXIT
    
    # 先清理过期会话锁
    log "CLEAN" "=== Session Lock Cleanup ==="
    clean_stale_locks
    
    local start=$(get_timestamp_ms)
    
    log_header
    
    local status=$(read_status)
    local failure=$(get_value "$status" "failure_count")
    local retry=$(get_value "$status" "retry_count")
    [ -z "$failure" ] && failure=0
    [ -z "$retry" ] && retry=0
    
    log_status "CHECK" "$failure" "$FAILURE_THRESHOLD" "$retry" "$MAX_RETRIES"
    log "  Load: $(get_system_load)"
    
    check_gateway
    issues=$?
    local elapsed=$(get_elapsed_ms $start)
    
    if [ $issues -eq 0 ]; then
        [ $failure -gt 0 ] && log "  [OK] Gateway recovered" || log "  [OK] Gateway healthy"
        
        local m=$(record_metrics $elapsed "success")
        local runs=$(echo $m | cut -d',' -f1)
        local rate=$(echo $m | cut -d',' -f2)
        local fails=$(echo $m | cut -d',' -f3)
        local restarts=$(echo $m | cut -d',' -f4)
        
        log_summary $runs $rate $fails $restarts
        log "  Time: ${elapsed}ms"
        save_status 0 0
        exit 0
    else
        failure=$((failure + 1))
        
        if [ $failure -lt $FAILURE_THRESHOLD ]; then
            save_status $failure $retry
            log "  [WARN] Check failed, need $((FAILURE_THRESHOLD - failure)) more"
            exit 1
        fi
        
        log "  [ERROR] $FAILURE_THRESHOLD consecutive failures"
        
        # 使用智能重试
        retry=0
        while [ $retry -lt $MAX_RETRIES ]; do
            retry=$((retry + 1))
            
            if [ $retry -eq 1 ]; then
                log "  [RETRY] Smart Retry #1 (initial delay: ${INITIAL_RETRY_DELAY}s)"
            else
                local delay=$(calculate_retry_delay $retry)
                log "  [RETRY] Smart Retry #$retry (backoff delay: ${delay}s)"
            fi
            
            restart_gateway
            sleep $STARTUP_WAIT
            
            if check_gateway; then
                log "  [OK] Gateway recovered after smart retry"
                
                local m=$(record_metrics $elapsed "restarted")
                local runs=$(echo $m | cut -d',' -f1)
                local rate=$(echo $m | cut -d',' -f2)
                local fails=$(echo $m | cut -d',' -f3)
                local restarts=$(echo $m | cut -d',' -f4)
                log_summary $runs $rate $fails $restarts
                save_status 0 0
                exit 0
            fi
        done
        
        # 维修
        log "  [ERROR] All smart retries failed, entering repair..."
        
        if repair_gateway; then
            record_metrics $elapsed "repaired"
            save_status 0 0
            exit 0
        else
            record_metrics $elapsed "failed"
            save_status $FAILURE_THRESHOLD $MAX_RETRIES
            exit 2
        fi
    fi
}

main