#!/bin/bash

# Настройка цветов для вывода
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Конфигурация
REPO_URL="https://git.tvstart.ru/ychernenko/KHL.git"
TARGET_DIR="/root/KHL"
SERVICE_NAME="khl-data.service"

show_help() {
    echo "Использование: $0 -r <релиз> [-h]"
    echo "  -r  Релиз (тег или ветка в git)"
    echo "  -h  Показать эту справку"
    echo ""
    echo "Пример: $0 -r Barabanov_TEST"
    echo ""
    exit 0
}

# Функция для вывода цветных сообщений
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1" >&2
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1" >&2
}

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

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1" >&2
}

# Функция проверки зависимостей системы
check_dependencies() {
    local deps=("git" "python3" "pip3" "netstat" "systemctl")
    local missing=()
    
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            missing+=("$dep")
        fi
    done
    
    if [ ${#missing[@]} -ne 0 ]; then
        log_warn "Отсутствуют некоторые команды: ${missing[*]}"
        return 1
    fi
    return 0
}

# Функция проверки и установки пакетов
install_packages() {
    log_info "Обновление списка пакетов..."
    if ! apt-get update; then
        log_error "Ошибка при обновлении списка пакетов"
        exit 1
    fi
    
    log_info "Установка необходимых пакетов..."
    if ! apt-get install -y python3 python3-pip python3-venv git net-tools; then
        log_error "Ошибка при установке пакетов"
        exit 1
    fi
    
    # Проверка установки
    if ! command -v python3 &> /dev/null; then
        log_error "Python3 не установлен!"
        exit 1
    fi
    
    if ! command -v pip3 &> /dev/null; then
        log_error "pip3 не установлен!"
        exit 1
    fi
    
    log_info "Версия Python: $(python3 --version)"
    log_info "Версия pip: $(pip3 --version)"
}

# Функция загрузки кода
download_code() {
    local release="$1"
    local target_dir="$TARGET_DIR"
    
    log_info "Создание рабочей директории $target_dir..."
    mkdir -p "$target_dir"
    
    if [ -d "$target_dir/.git" ]; then
        log_info "Обновление существующего репозитория..."
        cd "$target_dir" || exit 1
        git fetch --all
        if ! git checkout "$release"; then
            log_error "Не удалось переключиться на релиз $release"
            log_info "Доступные ветки:"
            git branch -r
            exit 1
        fi
        git pull origin "$release"
    else
        log_info "Клонирование репозитория $REPO_URL..."
        if ! git clone "$REPO_URL" "$target_dir"; then
            log_error "Ошибка при клонировании репозитория"
            exit 1
        fi
        cd "$target_dir" || exit 1
        if ! git checkout "$release"; then
            log_error "Не удалось переключиться на релиз $release"
            log_info "Доступные ветки:"
            git branch -r
            exit 1
        fi
    fi
    
    log_info "Код успешно загружен (релиз: $release)"
    log_info "Текущая ветка: $(git branch --show-current)"
    log_info "Последний коммит: $(git log --oneline -1)"
}

# Функция настройки виртуального окружения
setup_venv() {
    local target_dir="$TARGET_DIR"
    cd "$target_dir" || exit 1
    
    # Удаляем существующее виртуальное окружение для чистоты
    if [ -d ".venv" ]; then
        log_info "Удаление существующего виртуального окружения..."
        rm -rf .venv
    fi
    
    log_info "Создание виртуального окружения..."
    if ! python3 -m venv .venv; then
        log_error "Не удалось создать виртуальное окружение"
        exit 1
    fi
    
    if [ ! -f ".venv/bin/activate" ]; then
        log_error "Не удалось создать виртуальное окружение"
        exit 1
    fi
    
    log_info "Активация виртуального окружения и установка зависимостей..."
    source .venv/bin/activate
    
    # Обновление pip
    if ! pip install --upgrade pip; then
        log_warn "Не удалось обновить pip"
    fi
    
    # Проверка наличия requirements.txt
    if [ -f "requirements.txt" ]; then
        log_info "Установка зависимостей из requirements.txt..."
        if pip install -r requirements.txt; then
            log_info "Все зависимости успешно установлены"
        else
            log_error "Ошибка при установке зависимостей из requirements.txt"
            exit 1
        fi
    else
        log_warn "Файл requirements.txt не найден, устанавливаем базовые зависимости..."
        if ! pip install requests pandas numpy fastapi uvicorn python-telegram-handler python-dotenv; then
            log_error "Ошибка при установке базовых зависимостей"
            exit 1
        fi
    fi

    # Проверка основных пакетов
    log_info "Проверка установки основных пакетов..."
    if ! python -c "
try:
    import requests, pandas, numpy
    print('✓ Все основные пакеты успешно импортируются')
except ImportError as e:
    print(f'✗ Ошибка импорта: {e}')
    exit(1)
"; then
        log_error "Ошибка при проверке пакетов"
        exit 1
    fi
    
    log_info "Установленные пакеты:"
    pip list --format=columns
}

# Функция создания systemd сервиса
create_systemd_service() {
    log_info "Создание systemd сервиса"

    # Останавливаем и отключаем старый сервис если он есть
    if systemctl is-active --quiet "$SERVICE_NAME"; then
        log_info "Остановка $SERVICE_NAME..."
        systemctl stop "$SERVICE_NAME"
    fi

    if systemctl is-enabled --quiet "$SERVICE_NAME"; then
        log_info "Отключение $SERVICE_NAME..."
        systemctl disable "$SERVICE_NAME"
    fi

    # Формируем команду для data сервиса
    local data_command="$TARGET_DIR/.venv/bin/python3 $TARGET_DIR/get_data.py"
    local data_service_file="/etc/systemd/system/$SERVICE_NAME"
    
    log_info "Создание файла сервиса: $data_service_file"
    cat > "$data_service_file" << EOF
[Unit]
Description=KHL Data Service
After=network.target

[Service]
Type=simple
User=root
WorkingDirectory=$TARGET_DIR
Environment=PATH=$TARGET_DIR/.venv/bin
ExecStart=$data_command
Restart=always
RestartSec=30
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF

    # Настройка прав
    chmod 644 "$data_service_file"
    
    log_info "Перезагрузка systemd демона..."
    if ! systemctl daemon-reload; then
        log_error "Ошибка при перезагрузке systemd демона"
        exit 1
    fi
    
    log_info "Включение сервиса $SERVICE_NAME..."
    if ! systemctl enable "$SERVICE_NAME"; then
        log_error "Ошибка при включении сервиса"
        exit 1
    fi
}

# Функция проверки файлов
check_required_files() {
    local target_dir="$TARGET_DIR"
    cd "$target_dir" || exit 1
    
    local missing_files=()
    local required_files=("get_data.py")
    
    for file in "${required_files[@]}"; do
        if [ ! -f "$file" ]; then
            missing_files+=("$file")
        fi
    done
    
    if [ ${#missing_files[@]} -ne 0 ]; then
        log_error "Отсутствуют необходимые файлы: ${missing_files[*]}"
        log_error "Содержимое директории $target_dir:"
        ls -la
        exit 1
    fi
    
    log_info "Все необходимые файлы присутствуют"
    
    # Дополнительная проверка на исполняемость Python файла
    if ! python3 -c "import ast; ast.parse(open('get_data.py').read())" 2>/dev/null; then
        log_warn "Файл get_data.py содержит синтаксические ошибки Python"
    fi
}

# Функция настройки firewall
setup_firewall() {
    log_info "Настройка firewall (открытие порта 8000)..."
    
    # Проверяем наличие ufw
    if command -v ufw &> /dev/null && systemctl is-active --quiet ufw; then
        if ufw allow 8000/tcp; then
            log_info "Порт 8000 открыт в ufw"
        else
            log_warn "Не удалось открыть порт 8000 в ufw"
        fi
    fi
    
    # Для firewalld (CentOS/RHEL)
    if command -v firewall-cmd &> /dev/null && systemctl is-active --quiet firewalld; then
        if firewall-cmd --permanent --add-port=8000/tcp && firewall-cmd --reload; then
            log_info "Порт 8000 открыт в firewalld"
        else
            log_warn "Не удалось открыть порт 8000 в firewalld"
        fi
    fi
}

# Функция проверки порта
check_port() {
    local port=8000
    local occupied=false
    
    if command -v netstat &> /dev/null && netstat -tuln 2>/dev/null | grep -q ":$port "; then
        occupied=true
    elif command -v ss &> /dev/null && ss -tuln 2>/dev/null | grep -q ":$port "; then
        occupied=true
    fi
    
    if [ "$occupied" = true ]; then
        log_warn "Порт $port уже занят. Возможно, приложение уже запущено."
        return 1
    fi
    return 0
}

# Функция управления сервисами
manage_services() {
    local action="$1"
    
    case $action in
        "start")
            log_info "Запуск сервиса $SERVICE_NAME..."
            if ! systemctl start "$SERVICE_NAME"; then
                log_error "Ошибка при запуске сервиса"
                return 1
            fi
            ;;
        "restart")
            log_info "Перезапуск сервиса $SERVICE_NAME..."
            if ! systemctl restart "$SERVICE_NAME"; then
                log_error "Ошибка при перезапуске сервиса"
                return 1
            fi
            ;;
        "stop")
            log_info "Остановка сервиса $SERVICE_NAME..."
            if ! systemctl stop "$SERVICE_NAME"; then
                log_error "Ошибка при остановке сервиса"
                return 1
            fi
            ;;
    esac
    
    # Даем время сервису на запуск/остановку
    sleep 3
    return 0
}

# Функция проверки статуса сервиса
check_service_status() {
    log_info "Статус сервиса $SERVICE_NAME:"
    if systemctl is-active "$SERVICE_NAME" &>/dev/null; then
        systemctl status "$SERVICE_NAME" --no-pager -l
        return 0
    else
        log_warn "Сервис $SERVICE_NAME не запущен"
        return 1
    fi
}

# Функция проверки системы
check_system() {
    log_info "Проверка системы..."
    
    # Проверка дистрибутива
    if [ -f /etc/os-release ]; then
        source /etc/os-release
        log_info "Дистрибутив: $PRETTY_NAME"
    fi
    
    # Проверка свободного места
    local available_space
    available_space=$(df "$TARGET_DIR" --output=avail 2>/dev/null | tail -1)
    if [ -z "$available_space" ]; then
        available_space=$(df / --output=avail | tail -1)
    fi
    
    if [ "$available_space" -lt 1048576 ]; then  # Меньше 1GB
        log_warn "Мало свободного места: $((available_space / 1024)) MB"
    else
        log_info "Свободное место: $((available_space / 1024 / 1024)) GB"
    fi
    
    # Проверка памяти
    if [ -f /proc/meminfo ]; then
        local total_mem
        total_mem=$(grep MemTotal /proc/meminfo | awk '{print $2}')
        log_info "Оперативная память: $((total_mem / 1024)) MB"
    fi
    
    # Проверка архитектуры
    log_info "Архитектура: $(uname -m)"
}

# Функция очистки при прерывании
cleanup() {
    log_info "Очистка..."
    if [ -f "/etc/systemd/system/$SERVICE_NAME" ]; then
        manage_services "stop"
        systemctl disable "$SERVICE_NAME" 2>/dev/null || true
        rm -f "/etc/systemd/system/$SERVICE_NAME"
        systemctl daemon-reload
    fi
}

# Основная функция
main() {
    local release="main"  # значение по умолчанию
    
    # Обработка аргументов командной строки
    while getopts "r:h" opt; do
        case $opt in
            r) release="$OPTARG" ;;
            h) show_help ;;
            *) log_error "Неверный аргумент"; exit 1 ;;
        esac
    done
     
    log_info "Начало установки KHL Data Service..."
    if [ -n "$release" ]; then
        log_info "Релиз: $release"
    else
        log_info "Релиз: не указан (будет использовано значение по умолчанию: main)"
        release="main"
    fi
    
    # Проверка прав root
    if [[ $EUID -ne 0 ]]; then
        log_error "Этот скрипт должен запускаться с правами root"
        exit 1
    fi
    
    # Установка обработчика прерывания
    trap 'log_error "Прервано пользователем"; cleanup; exit 1' INT TERM
    
    # Проверка системы
    check_system
    
    # Проверка зависимостей
    check_dependencies
    
    # Установка пакетов
    install_packages
    
    # Загрузка кода
    download_code "$release"
    
    # Проверка файлов
    check_required_files
    
    # Настройка виртуального окружения
    setup_venv
    
    # Настройка firewall
    setup_firewall
    
    # Проверка порта
    check_port
    
    # Создание systemd сервиса
    create_systemd_service
    
    log_info "Настройка завершена!"
    
    # Запуск сервиса
    if manage_services "start"; then
        log_info "Сервис успешно запущен"
    else
        log_error "Не удалось запустить сервис"
        exit 1
    fi
    
    # Проверка статуса
    check_service_status
    
    # Получаем IP для вывода
    local ip_address
    ip_address=$(ip route get 1 2>/dev/null | awk '{print $7; exit}' || 
                 ip a 2>/dev/null | grep -Eo 'inet (addr:)?([0-9]*\.){3}[0-9]*' | 
                 grep -Eo '([0-9]*\.){3}[0-9]*' | grep -v '127.0.0.1' | head -n1)
    
    log_info "=================================================="
    log_info "Установка завершена успешно!"
    log_info "Релиз: $release"
    log_info "Рабочая директория: $TARGET_DIR"
    log_info ""
    log_info "Для просмотра логов:"
    log_info "  journalctl -u $SERVICE_NAME -f"
    log_info ""
    log_info "Управление сервисом:"
    log_info "  Перезапуск: systemctl restart $SERVICE_NAME"
    log_info "  Остановка: systemctl stop $SERVICE_NAME"
    log_info "  Статус: systemctl status $SERVICE_NAME"
    log_info "  Логи: journalctl -u $SERVICE_NAME"
    log_info ""
    log_info "Проверка работы:"
    log_info "  Проверить процессы: ps aux | grep get_data.py"
    log_info "  Проверить логи: tail -f $TARGET_DIR/logs/*.log 2>/dev/null || echo 'Директория логов не найдена'"
    log_info "=================================================="
}

# Запуск основной функции
main "$@"