#!/bin/bash
# 后台启动批量分析任务的包装脚本
# 支持 nohup、screen、tmux 三种方式

set -e

BASE_DIR="$HOME/bin/p_anlz"
SCRIPT_DIR="$BASE_DIR/scripts"
BATCH_SCRIPT="$SCRIPT_DIR/batch_analyze_all_folders.sh"
LOG_DIR="$BASE_DIR/logs"
NOHUP_LOG="$LOG_DIR/nohup.out"

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

# 显示使用说明
show_usage() {
    cat << EOF
批量图片分析后台运行工具

用法:
    $0 [方式]

方式选项:
    nohup       使用 nohup 后台运行（默认，最简单）
    screen      使用 screen 会话运行（可重新连接）
    tmux        使用 tmux 会话运行（推荐，功能最强）
    status      查看当前运行状态
    logs        实时查看主日志
    progress    查看处理进度
    live        实时查看当前 Python 程序输出（推荐）

示例:
    $0 nohup          # 使用 nohup 启动
    $0 screen         # 使用 screen 启动
    $0 tmux           # 使用 tmux 启动
    $0 status         # 查看状态
    $0 logs           # 查看主日志
    $0 progress       # 查看进度
    $0 live           # 实时查看 Python 输出（含进度条）

后台运行后如何查看:
    - nohup 方式: tail -f $NOHUP_LOG
    - screen 方式: screen -r batch_analyze
    - tmux 方式: tmux attach -t batch_analyze
    - 实时输出: $0 live
EOF
}

# 检查脚本是否存在
check_script() {
    if [ ! -f "$BATCH_SCRIPT" ]; then
        echo -e "${RED}错误: 批处理脚本不存在: $BATCH_SCRIPT${NC}"
        exit 1
    fi
    chmod +x "$BATCH_SCRIPT"
}

# 检查是否已经在运行
check_running() {
    if pgrep -f "batch_analyze_all_folders.sh" > /dev/null; then
        echo -e "${YELLOW}警告: 批量分析任务已在运行中${NC}"
        echo "进程信息:"
        ps aux | grep "batch_analyze_all_folders.sh" | grep -v grep
        echo ""
        read -p "是否继续启动新任务? (y/N): " confirm
        if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then
            echo "已取消"
            exit 0
        fi
    fi
}

# 使用 nohup 启动
start_nohup() {
    echo -e "${BLUE}使用 nohup 启动批量分析任务...${NC}"
    mkdir -p "$LOG_DIR"

    nohup bash "$BATCH_SCRIPT" > "$NOHUP_LOG" 2>&1 &
    local pid=$!

    echo -e "${GREEN}✓ 任务已在后台启动 (PID: $pid)${NC}"
    echo ""
    echo "查看日志:"
    echo "  tail -f $NOHUP_LOG"
    echo ""
    echo "查看进度:"
    echo "  $0 progress"
    echo ""
    echo "停止任务:"
    echo "  kill $pid"
}

# 使用 screen 启动
start_screen() {
    if ! command -v screen &> /dev/null; then
        echo -e "${YELLOW}screen 未安装，请先安装: sudo apt-get install screen${NC}"
        exit 1
    fi

    echo -e "${BLUE}使用 screen 启动批量分析任务...${NC}"

    # 检查是否已有同名 session
    if screen -list | grep -q "batch_analyze"; then
        echo -e "${YELLOW}警告: screen 会话 'batch_analyze' 已存在${NC}"
        read -p "是否删除旧会话并创建新会话? (y/N): " confirm
        if [ "$confirm" = "y" ] || [ "$confirm" = "Y" ]; then
            screen -S batch_analyze -X quit 2>/dev/null || true
        else
            echo "已取消"
            exit 0
        fi
    fi

    screen -dmS batch_analyze bash "$BATCH_SCRIPT"

    echo -e "${GREEN}✓ 任务已在 screen 会话中启动${NC}"
    echo ""
    echo "重新连接到会话:"
    echo "  screen -r batch_analyze"
    echo ""
    echo "分离会话: Ctrl+A, 然后按 D"
    echo ""
    echo "查看所有会话:"
    echo "  screen -ls"
    echo ""
    echo "终止会话:"
    echo "  screen -S batch_analyze -X quit"
}

# 使用 tmux 启动
start_tmux() {
    if ! command -v tmux &> /dev/null; then
        echo -e "${YELLOW}tmux 未安装，请先安装: sudo apt-get install tmux${NC}"
        exit 1
    fi

    echo -e "${BLUE}使用 tmux 启动批量分析任务...${NC}"

    # 检查是否已有同名 session
    if tmux has-session -t batch_analyze 2>/dev/null; then
        echo -e "${YELLOW}警告: tmux 会话 'batch_analyze' 已存在${NC}"
        read -p "是否删除旧会话并创建新会话? (y/N): " confirm
        if [ "$confirm" = "y" ] || [ "$confirm" = "Y" ]; then
            tmux kill-session -t batch_analyze 2>/dev/null || true
        else
            echo "已取消"
            exit 0
        fi
    fi

    tmux new-session -d -s batch_analyze "bash $BATCH_SCRIPT"

    echo -e "${GREEN}✓ 任务已在 tmux 会话中启动${NC}"
    echo ""
    echo "重新连接到会话:"
    echo "  tmux attach -t batch_analyze"
    echo ""
    echo "分离会话: Ctrl+B, 然后按 D"
    echo ""
    echo "查看所有会话:"
    echo "  tmux ls"
    echo ""
    echo "终止会话:"
    echo "  tmux kill-session -t batch_analyze"
}

# 查看运行状态
show_status() {
    echo -e "${BLUE}批量分析任务状态:${NC}"
    echo ""

    # 检查进程
    if pgrep -f "batch_analyze_all_folders.sh" > /dev/null; then
        echo -e "${GREEN}✓ 任务正在运行${NC}"
        echo ""
        echo "进程信息:"
        ps aux | grep "batch_analyze_all_folders.sh" | grep -v grep
    else
        echo -e "${YELLOW}○ 任务未运行${NC}"
    fi

    echo ""

    # 检查 screen 会话
    if command -v screen &> /dev/null; then
        if screen -list | grep -q "batch_analyze"; then
            echo -e "${GREEN}✓ Screen 会话存在${NC}"
            screen -list | grep "batch_analyze"
        fi
    fi

    echo ""

    # 检查 tmux 会话
    if command -v tmux &> /dev/null; then
        if tmux has-session -t batch_analyze 2>/dev/null; then
            echo -e "${GREEN}✓ Tmux 会话存在${NC}"
            tmux list-sessions | grep "batch_analyze"
        fi
    fi
}

# 查看日志
show_logs() {
    local main_log="$LOG_DIR/batch_analyze.log"

    if [ -f "$main_log" ]; then
        echo -e "${BLUE}实时查看主日志 (Ctrl+C 退出):${NC}"
        echo ""
        tail -f "$main_log"
    elif [ -f "$NOHUP_LOG" ]; then
        echo -e "${BLUE}实时查看 nohup 日志 (Ctrl+C 退出):${NC}"
        echo ""
        tail -f "$NOHUP_LOG"
    else
        echo -e "${YELLOW}未找到日志文件${NC}"
        echo "日志目录: $LOG_DIR"
    fi
}

# 查看进度
show_progress() {
    local progress_file="$LOG_DIR/progress.txt"

    if [ ! -f "$progress_file" ]; then
        echo -e "${YELLOW}进度文件不存在: $progress_file${NC}"
        exit 1
    fi

    echo -e "${BLUE}批量分析进度:${NC}"
    echo ""

    # 显示进度文件内容
    cat "$progress_file"

    echo ""
    echo -e "${BLUE}统计信息:${NC}"

    # 统计各状态数量
    local total=$(grep -c "SUCCESS\|FAILED\|SKIPPED" "$progress_file" || echo "0")
    local success=$(grep -c "SUCCESS" "$progress_file" || echo "0")
    local failed=$(grep -c "FAILED" "$progress_file" || echo "0")
    local skipped=$(grep -c "SKIPPED" "$progress_file" || echo "0")

    echo "总计: $total"
    echo -e "${GREEN}成功: $success${NC}"
    echo -e "${YELLOW}失败: $failed${NC}"
    echo -e "${BLUE}跳过: $skipped${NC}"

    # 显示最后更新时间
    if [ -f "$progress_file" ]; then
        local last_update=$(stat -c %y "$progress_file" 2>/dev/null || stat -f "%Sm" "$progress_file" 2>/dev/null)
        echo ""
        echo "最后更新: $last_update"
    fi
}

# 实时查看当前 Python 程序输出
show_live() {
    local main_log="$LOG_DIR/batch_analyze.log"

    echo -e "${BLUE}实时查看 Python 程序输出 (Ctrl+C 退出)${NC}"
    echo ""

    # 检查任务是否在运行
    if ! pgrep -f "batch_analyze_all_folders.sh" > /dev/null; then
        echo -e "${YELLOW}任务未运行${NC}"
        exit 1
    fi

    # 查找当前正在处理的文件夹
    if [ -f "$main_log" ]; then
        local current_folder=$(grep -oP "处理文件夹: \K.*" "$main_log" | tail -1)
        if [ -n "$current_folder" ]; then
            local folder_log="$LOG_DIR/${current_folder}.log"

            if [ -f "$folder_log" ]; then
                echo -e "${GREEN}当前处理: $current_folder${NC}"
                echo -e "${BLUE}日志文件: $folder_log${NC}"
                echo ""
                echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

                # 实时跟踪文件夹日志（显示 Python 程序的完整输出）
                tail -f "$folder_log"
            else
                echo -e "${YELLOW}等待日志文件生成...${NC}"
                sleep 2
                show_live
            fi
        else
            echo -e "${YELLOW}等待任务开始...${NC}"
            sleep 2
            show_live
        fi
    else
        echo -e "${YELLOW}主日志文件不存在${NC}"
        exit 1
    fi
}

# 主函数
main() {
    local mode="${1:-nohup}"

    case "$mode" in
        nohup)
            check_script
            check_running
            start_nohup
            ;;
        screen)
            check_script
            check_running
            start_screen
            ;;
        tmux)
            check_script
            check_running
            start_tmux
            ;;
        status)
            show_status
            ;;
        logs)
            show_logs
            ;;
        progress)
            show_progress
            ;;
        live)
            show_live
            ;;
        -h|--help|help)
            show_usage
            ;;
        *)
            echo "未知选项: $mode"
            echo ""
            show_usage
            exit 1
            ;;
    esac
}

main "$@"