#!/bin/bash

# 远程服务器列表，按需修改
SERVERS=("175.178.54.98" "1.12.218.143" "43.139.169.25" "129.204.180.174" "42.194.142.169")
REMOTE_USER="ubuntu"
REMOTE_SCRIPT_DIR="/home/ubuntu/pyscripts/stockapp/reports_em"
REMOTE_LOG_DIR="/home/ubuntu/pyscripts/stockapp/log"
DATA_DIR="/home/ubuntu/hostdir/stock_data/pdfs"
LOCAL_FILES=("config.py" "em_reports.py" "fetch.py")

# 远程任务参数配置（每台服务器的不同参数）
TASK_PARAMS=(
    "--mode=down --begin=2024-01-01 --end=2024-06-30"
    "--mode=down --begin=2023-01-01 --end=2023-06-30"
    "--mode=down --begin=2022-01-01 --end=2022-06-30"
    "--mode=down --begin=2021-01-01 --end=2021-06-30"
    "--mode=down --begin=2020-01-01 --end=2020-06-30"
)

# 推送代码到所有服务器
function push_code() {
    for SERVER in "${SERVERS[@]}"; do
        echo "Pushing code to $SERVER..."
        scp "${LOCAL_FILES[@]}" "$REMOTE_USER@$SERVER:$REMOTE_SCRIPT_DIR/"
    done
}

# 启动任务
function start_tasks() {
    for i in "${!SERVERS[@]}"; do
        SERVER="${SERVERS[$i]}"
        PARAMS="${TASK_PARAMS[$i]}"
        echo "Starting task on $SERVER with params: $PARAMS"
        #ssh "$REMOTE_USER@$SERVER" "cd $REMOTE_SCRIPT_DIR && nohup python3 ./fetch.py $PARAMS > ../log/nohup.log 2>&1 &"
        #ssh "$REMOTE_USER@$SERVER" "cd $REMOTE_SCRIPT_DIR && nohup python3 ./fetch.py $PARAMS > ../log/nohup.log 2>&1 < /dev/null & disown"
        # nohup ... < /dev/null：防止 nohup 等待 stdin，立即释放控制权。
        # & disown：确保进程与 SSH 彻底分离，避免 SIGHUP 信号影响，SSH 立即返回。
        # ssh -n：禁用 ssh 的 stdin，防止远程进程等待输入。
        ssh -n "$REMOTE_USER@$SERVER" "cd $REMOTE_SCRIPT_DIR && nohup python3 ./fetch.py $PARAMS > ../log/nohup.log 2>&1 < /dev/null & disown" 
    done
}

# 停止任务
function stop_tasks() {
    for SERVER in "${SERVERS[@]}"; do
        echo "Stopping task on $SERVER..."
        ssh "$REMOTE_USER@$SERVER" "pkill -f 'python3 ./fetch.py'"
    done
}

# 获取任务进度
function check_progress() {
    for SERVER in "${SERVERS[@]}"; do
        echo -e "\nChecking progress on $SERVER..."
        FILE_COUNT=$(ssh "$REMOTE_USER@$SERVER" "ls -lRh $DATA_DIR | grep pdf | wc -l")
        FILE_SIZE=$(ssh "$REMOTE_USER@$SERVER" "du -sh $DATA_DIR")
        PROCESS_COUNT=$(ssh "$REMOTE_USER@$SERVER" "ps aux | grep '[f]etch.py' | wc -l")
        
        if [ "$PROCESS_COUNT" -gt 0 ]; then
            echo "Process status: Running ($PROCESS_COUNT instances), if 2, include parent progress"
        else
            echo "Process status: Not running"
        fi

        echo "Total files: $FILE_COUNT"
        echo "Total size : $FILE_SIZE"

        ERROR_LINES=$(ssh "$REMOTE_USER@$SERVER" "grep -v INFO $REMOTE_LOG_DIR/fetch_202503* | wc -l")
        echo "Error lines: $ERROR_LINES"

        TASK_COUNT=$(ssh "$REMOTE_USER@$SERVER" "grep 'running task. id' ~/pyscripts/stockapp/log/fetch_20250316.log | wc -l")
        echo "Task count: $TASK_COUNT"

    done
}

# 获取任务进度
function check_progress_robot() {
    result=""
    for SERVER in "${SERVERS[@]}"; do
        result+="\nChecking progress on $SERVER...\n"
        FILE_COUNT=$(ssh "$REMOTE_USER@$SERVER" "ls -lRh $DATA_DIR | grep pdf | wc -l")
        FILE_SIZE=$(ssh "$REMOTE_USER@$SERVER" "du -sh $DATA_DIR")
        PROCESS_COUNT=$(ssh "$REMOTE_USER@$SERVER" "ps aux | grep '[f]etch.py' | wc -l")
        
        if [ "$PROCESS_COUNT" -gt 0 ]; then
            result+="Process status: Running ($PROCESS_COUNT instances), if 2, include parent progress\n"
        else
            result+="Process status: Not running\n"
        fi

        result+="Total files: $FILE_COUNT\n"
        result+="Total size : $FILE_SIZE\n"

        ERROR_LINES=$(ssh "$REMOTE_USER@$SERVER" "grep -v INFO $REMOTE_LOG_DIR/fetch_202503* | wc -l")
        result+="Error lines: $ERROR_LINES\n"

        TASK_COUNT=$(ssh "$REMOTE_USER@$SERVER" "grep 'running task. id' ~/pyscripts/stockapp/log/fetch_20250316.log | wc -l")
    done
    echo -e "$result"
    # 调用 Python 脚本发送消息
    python3 ./robot.py "$result"
}

# 脚本菜单
case "$1" in
    push) push_code ;;
    start) start_tasks ;;
    stop) stop_tasks ;;
    check) check_progress ;;
    *)
        echo "Usage: $0 {push|start|stop|check}"
        exit 1
        ;;
esac