#!/bin/bash

# === Конфигурация ===
PROJECT_PATH="/root/rl-swarm"
LOG_FILE="$PROJECT_PATH/rl_swarm.log"
LOG_INACTIVITY_THRESHOLD=1800     # 30 минут (в секундах)
CHECK_INTERVAL=60                 # Проверка каждые 60 сек
STARTUP_GRACE_PERIOD=180          # Ожидание после запуска ноды
SCREEN_SESSION="gensyn"

# === Проверка, что скрипт существует ===
APP_COMMAND="$PROJECT_PATH/run_rl_swarm.sh"
if [ ! -f "$APP_COMMAND" ]; then
    echo "❌ Ошибка: Скрипт $APP_COMMAND не найден!"
    exit 1
fi
if [ ! -x "$APP_COMMAND" ]; then
    echo "❌ Ошибка: Нет права на исполнение $APP_COMMAND"
    echo "➡️  Выполните: chmod +x $APP_COMMAND"
    exit 1
fi

# === Бесконечный цикл мониторинга ===
while true; do
    echo -e "\n▒▒ Запуск RL-сварм-ноды..."

    echo "▒▒ Создаём новую screen-сессию '$SCREEN_SESSION'..."
    cd "$PROJECT_PATH" || exit 1

    if [ ! -d ".venv" ]; then
        echo " Создаём виртуальное окружение .venv..."
        python3 -m venv .venv
    fi

    # Активируем виртуальное окружение
    source .venv/bin/activate
    
    # Поднимаем лимиты для текущей сессии (и для screen-процессов)
    ulimit -n 65535
    ulimit -u 65535

    # Настройка модели перед запуском (Python внутри bash)
    python3 - <<EOF
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

model_name = os.environ.get("MODEL_NAME", "Gensyn/Qwen2.5-0.5B-Instruct")
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Устанавливаем max_length для ограничения потребления памяти
model.config.max_length = 384
EOF

    # Запуск в screen с отключением uvloop, буферизацией Python и CPU-only
    screen -S "$SCREEN_SESSION" -dm bash -c \
    "export PYTHONASYNCIODEBUG=1 \
     UVLOOP_NO_EXTENSIONS=1 \
     PYTHONUNBUFFERED=1 \
     CUDA_VISIBLE_DEVICES=\"\" \
     OMP_NUM_THREADS=1 \
     MKL_NUM_THREADS=1; \
    cd $PROJECT_PATH && source .venv/bin/activate && \
    pip install --force-reinstall transformers==4.51.3 trl==0.19.1 && \
    bash run_rl_swarm.sh 2>&1 | tee -a \"$LOG_FILE\""

    # Проверка лимитов внутри screen
    screen -S "$SCREEN_SESSION" -X stuff $'echo "===== ulimit -n внутри screen ====="; ulimit -n; echo "============================\n"'

    echo "⏳ Ожидание $STARTUP_GRACE_PERIOD секунд перед началом мониторинга логов..."
    sleep "$STARTUP_GRACE_PERIOD"

    while true; do
        # Проверка активности screen-процесса
        SCREEN_PID=$(screen -ls | grep "\.${SCREEN_SESSION}" | awk -F. '{print $1}' | tr -d '[:space:]')
        if [ -z "$SCREEN_PID" ] || ! ps -p "$SCREEN_PID" > /dev/null 2>&1; then
            echo "❗️ screen-сессия '$SCREEN_SESSION' завершилась. Немедленный перезапуск..."
            screen -S "$SCREEN_SESSION" -X quit || true
            pkill -f run_rl_swarm.sh || true
            pkill -f swarm_launcher || true
            pkill -f p2pd || true
            sleep 5
            break
        fi

        if ! screen -list | grep -q "$SCREEN_SESSION"; then
            echo "⛔ screen-сессия '$SCREEN_SESSION' отсутствует. Перезапуск..."
            break
        fi

        if [ ! -f "$LOG_FILE" ]; then
            echo "⚠️ Лог-файл не найден: $LOG_FILE"
            break
        fi

        LAST_UPDATE=$(stat -c %Y "$LOG_FILE")
        NOW=$(date +%s)
        DIFF=$((NOW - LAST_UPDATE))

        echo "▒▒️ Проверка активности логов: последний апдейт был $DIFF сек назад."

        if [ "$DIFF" -ge "$LOG_INACTIVITY_THRESHOLD" ]; then
            echo "⛔ Лог не обновлялся более $((LOG_INACTIVITY_THRESHOLD/60)) минут. Считаем зависанием."

            echo "▒▒ Завершаем screen-сессию '$SCREEN_SESSION'..."
            screen -S "$SCREEN_SESSION" -X quit || true

            echo "▒▒ Завершаем фоновые процессы..."
            pkill -f run_rl_swarm.sh || true
            pkill -f swarm_launcher || true
            pkill -f p2pd || true

            echo "▒▒ Состояние процессов на момент зависания:" >> "$LOG_FILE"
            echo "===== TOP 10 по памяти =====" >> "$LOG_FILE"
            ps aux --sort=-%mem | head -n 10 >> "$LOG_FILE"
            echo "===== TOP 10 по CPU =====" >> "$LOG_FILE"
            ps aux --sort=-%cpu | head -n 10 >> "$LOG_FILE"
            free -h >> "$LOG_FILE"

            echo "♻️ Перезапуск через 5 секунд..."
            sleep 5
            break
        fi

        sleep "$CHECK_INTERVAL"
    done
done
