#!/bin/bash
# 安全模式切换脚本
# 用法: safe-mode.sh [enter|exit|status]

set -e

# 配置
HERMES_HOME="${HERMES_HOME:-$HOME/.hermes}"
CONFIG_FILE="$HERMES_HOME/config.yaml"
BACKUP_FILE="$HERMES_HOME/safe_mode_backup.json"

# 安全模式模型
SAFE_MODEL="huihui-gemma-4-26b-a4b-it-abliterated-mlx"
SAFE_PROVIDER="lm-studik-hi6k"

# 颜色
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

print_line() {
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
}

print_box() {
    local content="$1"
    echo "┌────────────────────────────────────────┐"
    echo "$content" | while IFS= read -r line; do
        printf "│ %-38s │\n" "$line"
    done
    echo "└────────────────────────────────────────┘"
}

# 获取当前配置的模型
get_current_model() {
    if [ -f "$BACKUP_FILE" ] && is_in_safe_mode; then
        grep '"original_model"' "$BACKUP_FILE" | cut -d'"' -f4
    elif [ -f "$CONFIG_FILE" ]; then
        python3 << EOF
import yaml
try:
    with open("$CONFIG_FILE", 'r') as f:
        config = yaml.safe_load(f)
    model_config = config.get('model', {})
    print(model_config.get('default', 'unknown'))
except:
    print('unknown')
EOF
    else
        echo "unknown"
    fi
}

# 检查是否在安全模式
is_in_safe_mode() {
    if [ -f "$BACKUP_FILE" ]; then
        local in_safe_mode=$(grep -o '"in_safe_mode": *true' "$BACKUP_FILE" 2>/dev/null)
        if [ -n "$in_safe_mode" ]; then
            return 0
        fi
    fi
    return 1
}

# 创建备份
create_backup() {
    local backup_time=$(date -u +"%Y-%m-%dT%H:%M:%S")
    
    python3 << EOF
import yaml
import json
import sys

config_path = "$CONFIG_FILE"
backup_path = "$BACKUP_FILE"

try:
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    model_config = config.get('model', {})
    
    backup = {
        "original_model": model_config.get('default', 'unknown'),
        "original_base_url": model_config.get('base_url', ''),
        "original_api_key": model_config.get('api_key', ''),
        "original_api_mode": model_config.get('api_mode', 'openai'),
        "original_provider": model_config.get('provider', 'custom'),
        "backup_time": "$backup_time",
        "in_safe_mode": True
    }
    
    with open(backup_path, 'w') as f:
        json.dump(backup, f, indent=2)
    
    print("💾 备份已创建")
    print(f"   原模型: {backup['original_model']}")
    print(f"   备份时间: $backup_time")
except Exception as e:
    print(f"❌ 备份失败: {e}", file=sys.stderr)
    sys.exit(1)
EOF
}

# 读取备份
read_backup() {
    if [ -f "$BACKUP_FILE" ]; then
        python3 << EOF
import json
try:
    with open("$BACKUP_FILE", 'r') as f:
        backup = json.load(f)
    print(f"{backup.get('original_model', '')}|{backup.get('original_base_url', '')}|{backup.get('original_api_key', '')}|{backup.get('original_api_mode', 'openai')}|{backup.get('original_provider', 'custom')}")
except:
    print("|||")
EOF
    else
        echo "|||"
    fi
}

# 切换模型
switch_model() {
    local new_model="$1"
    local new_provider="$2"
    
    cp "$CONFIG_FILE" "${CONFIG_FILE}.bak"
    
    python3 << EOF
import yaml
import sys

config_path = "$CONFIG_FILE"
target_provider = "$new_provider"

try:
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    target_config = None
    if 'custom_providers' in config:
        for provider in config['custom_providers']:
            if provider.get('name') == target_provider:
                target_config = provider
                break
    
    if target_config:
        if 'model' not in config:
            config['model'] = {}
        
        config['model']['default'] = target_config.get('model', '')
        config['model']['base_url'] = target_config.get('base_url', '')
        config['model']['api_key'] = target_config.get('api_key', '')
        config['model']['provider'] = 'custom'
        
        if 'api_mode' in target_config:
            config['model']['api_mode'] = target_config['api_mode']
        else:
            config['model']['api_mode'] = 'openai'
        
        print(f"✅ 已切换到 {target_config.get('model', '')}")
    else:
        print(f"❌ 未找到provider: {target_provider}", file=sys.stderr)
        sys.exit(1)
    
    with open(config_path, 'w') as f:
        yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
    
    print("✅ 配置已更新")
except Exception as e:
    print(f"❌ 更新失败: {e}", file=sys.stderr)
    sys.exit(1)
EOF
}

# 进入安全模式
enter_safe_mode() {
    echo ""
    echo -e "${BLUE}🔄 正在切换到安全模式...${NC}"
    print_line
    
    if is_in_safe_mode; then
        echo ""
        echo -e "${YELLOW}⚠️  已经在安全模式中${NC}"
        echo ""
        return 0
    fi
    
    create_backup
    
    echo ""
    echo -e "${YELLOW}📝 正在更新配置...${NC}"
    switch_model "$SAFE_MODEL" "$SAFE_PROVIDER"
    
    echo ""
    print_line
    echo -e "${GREEN}✅ 已进入安全模式${NC}"
    echo ""
    print_box "🔓 当前模型
$SAFE_MODEL

💾 原模型已备份
$(get_current_model)

💡 说\"退出安全模式\"可恢复"
    echo ""
}

# 退出安全模式
exit_safe_mode() {
    echo ""
    echo -e "${BLUE}🔄 正在退出安全模式...${NC}"
    print_line
    
    if ! is_in_safe_mode; then
        echo ""
        echo -e "${RED}❌ 未找到安全模式备份${NC}"
        echo "   可能之前没有进入过安全模式"
        echo ""
        return 1
    fi
    
    local backup=$(read_backup)
    local original_model=$(echo "$backup" | cut -d'|' -f1)
    local original_base_url=$(echo "$backup" | cut -d'|' -f2)
    local original_api_key=$(echo "$backup" | cut -d'|' -f3)
    local original_api_mode=$(echo "$backup" | cut -d'|' -f4)
    local original_provider=$(echo "$backup" | cut -d'|' -f5)
    
    if [ -z "$original_model" ]; then
        echo ""
        echo -e "${RED}❌ 备份文件损坏，无法恢复${NC}"
        echo ""
        return 1
    fi
    
    echo ""
    echo -e "${YELLOW}📝 正在恢复配置...${NC}"
    
    python3 << EOF
import yaml
import sys

config_path = "$CONFIG_FILE"

try:
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    config['model'] = {
        'default': "$original_model",
        'base_url': "$original_base_url",
        'api_key': "$original_api_key",
        'api_mode': "$original_api_mode",
        'provider': "$original_provider"
    }
    
    with open(config_path, 'w') as f:
        yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
    
    print("✅ 配置已恢复")
except Exception as e:
    print(f"❌ 恢复失败: {e}", file=sys.stderr)
    sys.exit(1)
EOF
    
    # 更新备份状态
    python3 << EOF
import json

backup_file = "$BACKUP_FILE"
try:
    with open(backup_file, 'r') as f:
        backup = json.load(f)
    backup['in_safe_mode'] = False
    with open(backup_file, 'w') as f:
        json.dump(backup, f, indent=2)
except:
    pass
EOF
    
    echo ""
    print_line
    echo -e "${GREEN}✅ 已退出安全模式${NC}"
    echo ""
    print_box "🔒 当前模型
$original_model

📝 已恢复到切换前的配置"
    echo ""
}

# 查看状态
show_status() {
    echo ""
    echo -e "${BLUE}📊 安全模式状态${NC}"
    print_line
    
    if is_in_safe_mode; then
        local backup=$(read_backup)
        local original_model=$(echo "$backup" | cut -d'|' -f1)
        local backup_time=$(grep '"backup_time"' "$BACKUP_FILE" | cut -d'"' -f4)
        
        echo ""
        print_box "🔓 当前处于安全模式

当前模型: $SAFE_MODEL
原模型: $original_model
进入时间: $backup_time"
    else
        echo ""
        print_box "🔒 当前处于正常模式

当前模型: $(get_current_model)"
    fi
    
    echo ""
}

# 主函数
main() {
    local action="$1"
    
    case "$action" in
        enter)
            enter_safe_mode
            ;;
        exit)
            exit_safe_mode
            ;;
        status)
            show_status
            ;;
        *)
            echo "用法: $0 [enter|exit|status]"
            echo ""
            echo "  enter   - 进入安全模式"
            echo "  exit    - 退出安全模式"
            echo "  status  - 查看状态"
            exit 1
            ;;
    esac
}

main "$@"