#!/bin/bash

# 注意：不使用 set -e，以便更好地处理错误和重试

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

# 状态文件路径
STATUS_FILE="/root/.init-debian-status"

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查任务是否已完成
task_completed() {
    local task_id="$1"
    if [ -f "$STATUS_FILE" ]; then
        grep -q "^${task_id}$" "$STATUS_FILE"
        return $?
    fi
    return 1
}

# 标记任务为已完成
mark_task_completed() {
    local task_id="$1"
    mkdir -p "$(dirname "$STATUS_FILE")"
    if ! task_completed "$task_id"; then
        echo "$task_id" >> "$STATUS_FILE"
    fi
}

# 执行任务（带状态检查）
execute_task() {
    local task_id="$1"
    local task_name="$2"
    local task_func="$3"
    
    if task_completed "$task_id"; then
        log_info "任务 $task_id: $task_name - 已完成，跳过"
        return 0
    fi
    
    log_info "执行任务 $task_id: $task_name..."
    if $task_func; then
        mark_task_completed "$task_id"
        log_success "任务 $task_id 完成"
        return 0
    else
        log_error "任务 $task_id 失败"
        return 1
    fi
}

# 自动接受 SSH host key
accept_ssh_host_key() {
    local host="$1"
    local port="$2"
    local user_home="$3"  # 用户家目录路径
    local known_hosts_file="${user_home}/.ssh/known_hosts"
    
    # 确保 .ssh 目录存在并设置正确的权限
    mkdir -p "${user_home}/.ssh"
    chmod 700 "${user_home}/.ssh"
    chown kairee:kairee "${user_home}/.ssh"
    
    # 检查是否已存在该 host 的 key
    if ssh-keygen -F "[${host}]:${port}" -f "$known_hosts_file" &>/dev/null; then
        log_info "SSH host key 已存在于 known_hosts"
        return 0
    fi
    
    log_info "自动接受 SSH host key..."
    # 使用 ssh-keyscan 获取并添加 host key
    ssh-keyscan -p "$port" "$host" >> "$known_hosts_file" 2>/dev/null
    if [ $? -eq 0 ]; then
        chown kairee:kairee "$known_hosts_file"
        chmod 600 "$known_hosts_file"
        log_success "SSH host key 已添加"
        return 0
    else
        log_warning "无法自动获取 SSH host key，将在连接时手动确认"
        return 1
    fi
}

# 检查是否为 root
if [[ $EUID -ne 0 ]]; then
   log_error "此脚本必须以 root 权限运行"
   exit 1
fi

# 提示输入 kairee 用户密码
echo "========================================"
echo "Debian 系统初始化脚本"
echo "========================================"
echo
echo "请输入要为用户 'kairee' 设置的密码:"
read -s -p "密码: " USER_PASSWORD
echo
read -s -p "确认密码: " USER_PASSWORD_CONFIRM
echo

# 验证密码
if [ -z "$USER_PASSWORD" ]; then
    log_error "密码不能为空"
    exit 1
fi

if [ "$USER_PASSWORD" != "$USER_PASSWORD_CONFIRM" ]; then
    log_error "两次输入的密码不匹配"
    exit 1
fi

log_success "密码验证通过"
echo

# 任务函数定义
task_1_create_user() {
    if id "kairee" &>/dev/null; then
        log_warning "用户 'kairee' 已存在，跳过创建"
        return 0
    fi
    
    mkdir -p /data/home
    useradd -m -d /data/home/kairee -s /bin/bash -u 1000 -U kairee
    echo "kairee:$USER_PASSWORD" | chpasswd
    log_success "用户 'kairee' 创建成功 (uid=1000, 家目录: /data/home/kairee)"
    return 0
}

task_2_download_data() {
    DATA_URL="https://example.com/path/to/data.tar.gz"
    
    if [ -n "$1" ]; then
        DATA_URL="$1"
        log_info "使用提供的 URL: $DATA_URL"
    fi
    
    if [ "$DATA_URL" = "https://example.com/path/to/data.tar.gz" ]; then
        log_warning "未设置 DATA_URL，跳过下载步骤"
        log_warning "你可以手动将 data 目录复制到 /data"
        return 0
    fi
    
    TEMP_FILE=$(mktemp)
    log_info "正在下载..."
    if ! curl -L -o "$TEMP_FILE" "$DATA_URL"; then
        log_error "下载失败"
        rm -f "$TEMP_FILE"
        return 1
    fi
    
    log_info "解压到 /data..."
    mkdir -p /data
    if ! tar -xzf "$TEMP_FILE" -C /; then
        log_error "解压失败"
        rm -f "$TEMP_FILE"
        return 1
    fi
    
    rm -f "$TEMP_FILE"
    log_success "数据包下载并解压完成"
    return 0
}

task_3_install_packages() {
    apt-get update || return 1
    apt-get install -y sudo vim rsync curl sshpass || return 1
    apt-get reinstall -y openssl openssh-server || return 1
    log_success "基础软件包安装完成"
    return 0
}

task_4_install_docker() {
    if command -v docker &> /dev/null; then
        log_warning "Docker 已安装，跳过安装步骤"
        return 0
    fi
    
    if curl -Lo- https://get.docker.com | bash; then
        log_success "Docker CE 安装完成"
        return 0
    else
        log_error "Docker 安装失败"
        return 1
    fi
}

task_5_configure_apt() {
    if [ ! -d /data/etc/apt/sources.list.d ]; then
        log_warning "未找到 /data/etc/apt/sources.list.d 目录，跳过 apt 源配置"
        return 0
    fi
    
    if [ -f /etc/apt/sources.list ]; then
        mv /etc/apt/sources.list /etc/apt/sources.list.bak
        log_info "备份 /etc/apt/sources.list -> /etc/apt/sources.list.bak"
    fi
    
    if [ -f /etc/apt/sources.list.d/docker.list ]; then
        mv /etc/apt/sources.list.d/docker.list /etc/apt/sources.list.d/docker.list.bak
        log_info "备份 /etc/apt/sources.list.d/docker.list -> docker.list.bak"
    fi
    
    mkdir -p /etc/apt/sources.list.d
    
    for src_file in /data/etc/apt/sources.list.d/*; do
        if [ -f "$src_file" ]; then
            filename=$(basename "$src_file")
            target_file="/etc/apt/sources.list.d/$filename"
            
            if [ -L "$target_file" ]; then
                rm -f "$target_file"
            fi
            
            ln -s "$src_file" "$target_file"
            log_success "创建软链接: $src_file -> $target_file"
        fi
    done
    return 0
}

task_6_configure_docker() {
    if [ ! -f /data/etc/docker/daemon.json ]; then
        log_warning "未找到 /data/etc/docker/daemon.json，跳过 Docker 配置"
        return 0
    fi
    
    mkdir -p /etc/docker
    
    if [ -f /etc/docker/daemon.json ]; then
        cp /etc/docker/daemon.json /etc/docker/daemon.json.bak
        log_info "备份 /etc/docker/daemon.json -> daemon.json.bak"
    fi
    
    ln -sf /data/etc/docker/daemon.json /etc/docker/daemon.json
    log_success "创建 Docker daemon 配置软链接"
    return 0
}

task_7_configure_ssh() {
    if [ ! -f /data/etc/ssh/sshd_config ]; then
        log_warning "未找到 /data/etc/ssh/sshd_config，跳过 SSH 配置"
        return 0
    fi
    
    if [ -f /etc/ssh/sshd_config ]; then
        cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak
        log_info "备份 /etc/ssh/sshd_config -> sshd_config.bak"
    fi
    
    ln -sf /data/etc/ssh/sshd_config /etc/ssh/sshd_config
    log_success "创建 SSH 配置软链接"
    
    systemctl restart sshd || systemctl restart ssh
    log_success "SSH 服务已重启"
    return 0
}

task_8_configure_sudoers() {
    echo "kairee    ALL=(ALL)    NOPASSWD: ALL" > /etc/sudoers.d/kairee
    chmod 0440 /etc/sudoers.d/kairee
    log_success "创建 /etc/sudoers.d/kairee"
    return 0
}

task_9_add_docker_group() {
    if getent group docker > /dev/null; then
        usermod -aG docker kairee
        log_success "已将 kairee 添加到 docker 组"
        return 0
    else
        log_warning "docker 组不存在，请检查 Docker 是否正确安装"
        return 0  # 不算错误，只是警告
    fi
}

task_10_create_home_link() {
    if [ ! -d /data/home/kairee ]; then
        log_warning "未找到 /data/home/kairee，无法创建软链接"
        return 0
    fi
    
    if [ -e "/home/kairee" ] && [ ! -L "/home/kairee" ]; then
        log_warning "/home/kairee 已存在但不是软链接，将被删除并重新创建"
        rm -rf /home/kairee
    fi
    
    if [ -L "/home/kairee" ] && [ "$(readlink /home/kairee)" != "/data/home/kairee" ]; then
        log_warning "/home/kairee 软链接指向错误的位置，将被重新创建"
        rm -f /home/kairee
    fi
    
    if [ ! -e "/home/kairee" ]; then
        ln -s /data/home/kairee /home/kairee
        log_success "创建家目录软链接: /data/home/kairee -> /home/kairee"
    else
        log_success "家目录软链接已存在且正确"
    fi
    return 0
}

task_11_generate_ssh_key() {
    # 使用 kairee 用户的家目录
    KAIREE_HOME="/data/home/kairee"
    SSH_KEY_DIR="${KAIREE_HOME}/.ssh"
    SSH_PRIVATE_KEY="$SSH_KEY_DIR/id_ed25519"
    SSH_PUBLIC_KEY="$SSH_KEY_DIR/id_ed25519.pub"
    
    # 确保 kairee 用户存在
    if ! id "kairee" &>/dev/null; then
        log_error "用户 kairee 不存在，请先执行任务1"
        return 1
    fi
    
    # 创建 .ssh 目录并设置权限
    mkdir -p "$SSH_KEY_DIR"
    chmod 700 "$SSH_KEY_DIR"
    chown kairee:kairee "$SSH_KEY_DIR"
    
    if [ -f "$SSH_PRIVATE_KEY" ]; then
        log_warning "SSH 私钥已存在，跳过生成: $SSH_PRIVATE_KEY"
        return 0
    fi
    
    # 使用 kairee 用户身份生成密钥
    sudo -u kairee ssh-keygen -m PEM -t ed25519 -f "$SSH_PRIVATE_KEY" -N ""
    if [ $? -eq 0 ]; then
        log_success "生成 SSH 密钥对完成（kairee 用户）"
        log_info "私钥: $SSH_PRIVATE_KEY"
        log_info "公钥: $SSH_PUBLIC_KEY"
        return 0
    else
        log_error "SSH 密钥生成失败"
        return 1
    fi
}

task_12_copy_ssh_key() {
    REMOTE_HOST="qc5.ofcdn.cn"
    REMOTE_PORT="46571"
    REMOTE_USER="kairee"
    REMOTE_SSH_COPY_ID_USER="kairee"
    
    # 使用 kairee 用户的家目录
    KAIREE_HOME="/data/home/kairee"
    SSH_PUBLIC_KEY="${KAIREE_HOME}/.ssh/id_ed25519.pub"
    SSH_KNOWN_HOSTS="${KAIREE_HOME}/.ssh/known_hosts"
    
    if [ ! -f "$SSH_PUBLIC_KEY" ]; then
        log_error "公钥文件不存在: $SSH_PUBLIC_KEY"
        return 1
    fi
    
    # 自动接受 SSH host key（使用 kairee 用户的 known_hosts）
    accept_ssh_host_key "$REMOTE_HOST" "$REMOTE_PORT" "$KAIREE_HOME"
    
    # 密码重试机制
    MAX_RETRIES=3
    RETRY_COUNT=0
    
    while [ $RETRY_COUNT -lt $MAX_RETRIES ]; do
        echo
        echo "========================================"
        echo "输入远程服务器 SSH 密码"
        echo "========================================"
        echo "目标: $REMOTE_SSH_COPY_ID_USER@${REMOTE_HOST}:${REMOTE_PORT}"
        if [ $RETRY_COUNT -gt 0 ]; then
            log_warning "这是第 $((RETRY_COUNT + 1)) 次尝试（最多 $MAX_RETRIES 次）"
        fi
        echo
        echo "请输入 $REMOTE_SSH_COPY_ID_USER 在 $REMOTE_HOST 上的 SSH 密码:"
        read -s -p "SSH 密码: " REMOTE_SSH_PASSWORD
        echo
        
        if [ -z "$REMOTE_SSH_PASSWORD" ]; then
            log_warning "未输入 SSH 密码"
            if [ $RETRY_COUNT -lt $((MAX_RETRIES - 1)) ]; then
                read -p "是否继续重试？(y/n): " CONTINUE
                if [ "$CONTINUE" != "y" ] && [ "$CONTINUE" != "Y" ]; then
                    log_warning "用户取消，跳过 ssh-copy-id"
                    return 0
                fi
            fi
            RETRY_COUNT=$((RETRY_COUNT + 1))
            continue
        fi
        
        # 使用 sshpass 执行 ssh-copy-id，添加 SSH 选项以处理 host key
        log_info "正在尝试复制公钥到远程服务器（使用 kairee 用户的密钥）..."
        
        # 创建临时脚本文件来执行 SSH 命令（避免引号嵌套问题）
        TEMP_SCRIPT=$(mktemp)
        chmod 700 "$TEMP_SCRIPT"
        
        # 先测试连接并接受 host key（这会自动将 host key 添加到 kairee 用户的 known_hosts）
        cat > "$TEMP_SCRIPT" << EOF
#!/bin/bash
sshpass -p '$REMOTE_SSH_PASSWORD' ssh -p $REMOTE_PORT \
    -o StrictHostKeyChecking=accept-new \
    -o PasswordAuthentication=yes \
    $REMOTE_SSH_COPY_ID_USER@$REMOTE_HOST exit
EOF
        
        if sudo -u kairee bash "$TEMP_SCRIPT" >/dev/null 2>&1; then
            # 连接成功，现在复制公钥（此时 host key 已确认，使用 StrictHostKeyChecking=no）
            cat > "$TEMP_SCRIPT" << EOF
#!/bin/bash
sshpass -p '$REMOTE_SSH_PASSWORD' ssh-copy-id -p $REMOTE_PORT \
    -o StrictHostKeyChecking=no \
    $REMOTE_SSH_COPY_ID_USER@$REMOTE_HOST
EOF
            
            if sudo -u kairee bash "$TEMP_SCRIPT" >/dev/null 2>&1; then
                rm -f "$TEMP_SCRIPT"
                log_success "公钥已复制到远程服务器（kairee 用户）"
                rm -f "$TEMP_SCRIPT"
                unset REMOTE_SSH_PASSWORD
                return 0
            else
                rm -f "$TEMP_SCRIPT"
                RETRY_COUNT=$((RETRY_COUNT + 1))
                if [ $RETRY_COUNT -lt $MAX_RETRIES ]; then
                    log_error "公钥复制失败，请重试"
                    unset REMOTE_SSH_PASSWORD
                else
                    log_error "已达到最大重试次数，跳过 ssh-copy-id"
                    unset REMOTE_SSH_PASSWORD
                    return 1
                fi
            fi
        else
            rm -f "$TEMP_SCRIPT"
            RETRY_COUNT=$((RETRY_COUNT + 1))
            if [ $RETRY_COUNT -lt $MAX_RETRIES ]; then
                log_error "密码错误或连接失败，请重试"
                unset REMOTE_SSH_PASSWORD
            else
                log_error "已达到最大重试次数，跳过 ssh-copy-id"
                unset REMOTE_SSH_PASSWORD
                return 1
            fi
        fi
    done
    
    return 1
}

task_13_add_cron() {
    if [ ! -f /data/scripts/certs.sh ]; then
        log_warning "未找到 /data/scripts/certs.sh，跳过定时任务配置"
        return 0
    fi
    
    # 确保 kairee 用户存在
    if ! id "kairee" &>/dev/null; then
        log_error "用户 kairee 不存在，请先执行任务1"
        return 1
    fi
    
    # 确保 /data/certs 目录存在且权限正确
    mkdir -p /data/certs
    chown kairee:kairee /data/certs
    
    RANDOM_HOUR=$((1 + RANDOM % 7))
    RANDOM_MINUTE=$((RANDOM % 60))
    
    CRON_CMD="$RANDOM_MINUTE $RANDOM_HOUR * * * /data/scripts/certs.sh >> /data/certs/sync.log 2>&1"
    
    # 检查是否已存在相同的定时任务（kairee 用户的 crontab）
    if sudo -u kairee crontab -l 2>/dev/null | grep -q "/data/scripts/certs.sh"; then
        log_warning "定时任务已存在于 kairee 用户的 crontab，跳过添加"
        return 0
    fi
    
    # 添加到 kairee 用户的 crontab
    (sudo -u kairee crontab -l 2>/dev/null; echo "$CRON_CMD") | sudo -u kairee crontab -
    if [ $? -eq 0 ]; then
        log_success "定时任务已添加到 kairee 用户的 crontab"
        log_info "执行时间: 每天 $RANDOM_HOUR:$RANDOM_MINUTE"
        log_info "命令: /data/scripts/certs.sh >> /data/certs/sync.log 2>&1"
        return 0
    else
        log_error "添加定时任务失败"
        return 1
    fi
}

# 执行所有任务
echo
log_info "开始执行初始化任务..."
echo

# 保存原始参数（用于任务2）
ORIGINAL_ARGS="$@"

# 任务1: 创建用户 kairee
execute_task "task_1" "创建用户 kairee" task_1_create_user || log_warning "任务1失败，继续执行后续任务"

# 任务2: 下载并解压数据包（需要特殊处理参数）
if task_completed "task_2"; then
    log_info "任务 task_2: 下载并解压数据包 - 已完成，跳过"
else
    log_info "执行任务 task_2: 下载并解压数据包..."
    if task_2_download_data $ORIGINAL_ARGS; then
        mark_task_completed "task_2"
        log_success "任务 task_2 完成"
    else
        log_warning "任务2失败，继续执行后续任务"
    fi
fi

# 任务3: 安装软件包
execute_task "task_3" "安装基础软件包" task_3_install_packages || log_warning "任务3失败，继续执行后续任务"

# 任务4: 安装 Docker
execute_task "task_4" "安装 Docker CE" task_4_install_docker || log_warning "任务4失败，继续执行后续任务"

# 任务5: 配置 apt 源
execute_task "task_5" "配置 apt 源" task_5_configure_apt || log_warning "任务5失败，继续执行后续任务"

# 任务6: 配置 Docker daemon
execute_task "task_6" "配置 Docker daemon" task_6_configure_docker || log_warning "任务6失败，继续执行后续任务"

# 任务7: 配置 SSH
execute_task "task_7" "配置 SSH" task_7_configure_ssh || log_warning "任务7失败，继续执行后续任务"

# 任务8: 配置 sudoers
execute_task "task_8" "配置 sudoers" task_8_configure_sudoers || log_warning "任务8失败，继续执行后续任务"

# 任务9: 添加 kairee 到 docker 组
execute_task "task_9" "添加 kairee 到 docker 组" task_9_add_docker_group || log_warning "任务9失败，继续执行后续任务"

# 任务10: 创建家目录软链接
execute_task "task_10" "创建家目录软链接" task_10_create_home_link || log_warning "任务10失败，继续执行后续任务"

# 任务11: 生成 SSH 密钥对
execute_task "task_11" "生成 SSH 密钥对" task_11_generate_ssh_key || log_warning "任务11失败，继续执行后续任务"

# 任务12: 复制公钥到远程服务器
execute_task "task_12" "复制公钥到远程服务器" task_12_copy_ssh_key || log_warning "任务12失败，继续执行后续任务"

# 任务13: 添加定时任务
execute_task "task_13" "添加定时任务" task_13_add_cron || log_warning "任务13失败，继续执行后续任务"

echo
echo "========================================"
echo -e "${GREEN}初始化完成！${NC}"
echo "========================================"
echo
echo "重要提示："
echo "1. 请重新登录或执行 'su - kairee' 以使 Docker 组权限生效"
echo "2. 如果下载了数据包，请检查 /data 目录内容是否正确"
echo "3. SSH 配置已更新，请测试远程连接是否正常"
echo "4. 可以使用 'sudo -u kairee -i' 切换到 kairee 用户"
echo "5. SSH 密钥已生成（kairee 用户）: /data/home/kairee/.ssh/id_ed25519"
if sudo -u kairee crontab -l 2>/dev/null | grep -q "/data/scripts/certs.sh"; then
    CRON_TIME=$(sudo -u kairee crontab -l 2>/dev/null | grep "/data/scripts/certs.sh" | awk '{print $2":"$1}')
    echo "6. 定时任务（kairee 用户）: 每天 $CRON_TIME 执行 certs.sh"
else
    echo "6. 定时任务未配置（可能因为 /data/scripts/certs.sh 不存在）"
fi
echo
echo "状态文件: $STATUS_FILE"
echo "如需重新执行某个任务，可以编辑状态文件删除对应的任务ID"
echo

# 清理密码变量
unset USER_PASSWORD
unset USER_PASSWORD_CONFIRM
unset REMOTE_SSH_PASSWORD