#!/bin/bash

# ======================================
# AI Server Ubuntu Installation Script
# ======================================
# Автоматическое развертывание Ollama + Open WebUI на Ubuntu
# Оптимизировано для RTX 2060 и других NVIDIA GPU
# Версия: 1.0

set -e  # Остановка при ошибках

# Цвета для вывода
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# Конфигурация
DEFAULT_MODEL="llama3.2:3b"
WEB_PORT=3000
OLLAMA_PORT=11434
INSTALL_DIR="/opt/ai-server"
SERVICE_USER="aiserver"

# Функции вывода
print_banner() {
    echo -e "${GREEN}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                    AI SERVER UBUNTU                         ║"
    echo "║            Ollama + Open WebUI + RTX 2060                   ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

log_info() {
    echo -e "${CYAN}[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"
}

log_step() {
    echo -e "${PURPLE}[STEP]${NC} $1"
}

# Проверка прав root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "Этот скрипт должен запускаться с правами root"
        log_info "Используйте: sudo $0"
        exit 1
    fi
}

# Определение дистрибутива
detect_distro() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        DISTRO=$ID
        VERSION=$VERSION_ID
        log_info "Обнаружен дистрибутив: $PRETTY_NAME"
    else
        log_error "Не удается определить дистрибутив Linux"
        exit 1
    fi
}

# Обновление системы
update_system() {
    log_step "Обновление системы..."
    apt update && apt upgrade -y
    apt install -y curl wget git htop neofetch lsb-release gnupg software-properties-common
    log_success "Система обновлена"
}

# Проверка и установка NVIDIA драйверов
install_nvidia_drivers() {
    log_step "Проверка NVIDIA GPU..."
    
    if ! command -v nvidia-smi &> /dev/null; then
        log_warning "NVIDIA драйверы не найдены, устанавливаю..."
        
        # Определение рекомендуемого драйвера
        ubuntu-drivers devices
        
        # Автоматическая установка рекомендуемого драйвера
        ubuntu-drivers autoinstall
        
        log_info "NVIDIA драйверы установлены. Требуется перезагрузка!"
        log_warning "Выполните 'sudo reboot' и запустите скрипт снова"
        exit 0
    else
        log_success "NVIDIA драйверы уже установлены"
        nvidia-smi
    fi
}

# Установка Docker
install_docker() {
    log_step "Установка Docker..."
    
    if ! command -v docker &> /dev/null; then
        # Удаление старых версий
        apt remove -y docker.io docker-doc docker-compose podman-docker containerd runc 2>/dev/null || true
        
        # Установка репозитория Docker
        curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
        echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
        
        # Установка Docker
        apt update
        apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
        
        # Настройка пользователя
        usermod -aG docker $SUDO_USER 2>/dev/null || true
        systemctl enable docker
        systemctl start docker
        
        log_success "Docker установлен"
    else
        log_success "Docker уже установлен"
    fi
    
    # Проверка Docker
    docker --version
}

# Установка NVIDIA Container Toolkit
install_nvidia_docker() {
    log_step "Установка NVIDIA Container Toolkit..."
    
    if ! docker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi &>/dev/null; then
        # Добавление репозитория NVIDIA
        curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
        curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
            tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
        
        apt update
        apt install -y nvidia-container-toolkit
        
        # Настройка Docker для NVIDIA
        nvidia-ctk runtime configure --runtime=docker
        systemctl restart docker
        
        log_success "NVIDIA Container Toolkit установлен"
    else
        log_success "NVIDIA Container Toolkit уже настроен"
    fi
    
    # Тест GPU в Docker
    log_info "Тестирование GPU в Docker..."
    docker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi
}

# Установка Ollama
install_ollama() {
    log_step "Установка Ollama..."
    
    if ! command -v ollama &> /dev/null; then
        curl -fsSL https://ollama.ai/install.sh | sh
        log_success "Ollama установлен"
    else
        log_success "Ollama уже установлен"
    fi
    
    # Создание пользователя для сервиса
    if ! id "$SERVICE_USER" &>/dev/null; then
        useradd -r -s /bin/false -d /nonexistent $SERVICE_USER
        usermod -aG docker $SERVICE_USER
        log_success "Пользователь $SERVICE_USER создан"
    fi
}

# Настройка systemd сервиса для Ollama
setup_ollama_service() {
    log_step "Настройка Ollama сервиса..."
    
    cat > /etc/systemd/system/ollama.service << EOF
[Unit]
Description=Ollama Server
After=network-online.target

[Service]
ExecStart=/usr/local/bin/ollama serve
User=$SERVICE_USER
Group=$SERVICE_USER
Restart=always
RestartSec=3
Environment="OLLAMA_HOST=0.0.0.0:$OLLAMA_PORT"
Environment="OLLAMA_KEEP_ALIVE=24h"
Environment="OLLAMA_NUM_PARALLEL=2"
Environment="OLLAMA_MAX_LOADED_MODELS=2"
Environment="OLLAMA_FLASH_ATTENTION=1"

[Install]
WantedBy=multi-user.target
EOF
    
    systemctl daemon-reload
    systemctl enable ollama
    systemctl start ollama
    
    log_success "Ollama сервис настроен"
    
    # Ожидание запуска
    sleep 10
    
    # Проверка статуса
    if systemctl is-active --quiet ollama; then
        log_success "Ollama сервис запущен"
    else
        log_error "Ошибка запуска Ollama сервиса"
        journalctl -u ollama --no-pager -l
        exit 1
    fi
}

# Установка моделей
install_models() {
    log_step "Установка модели $DEFAULT_MODEL..."
    log_info "Это может занять несколько минут в зависимости от скорости интернета..."
    
    # Установка модели от имени пользователя сервиса
    sudo -u $SERVICE_USER ollama pull $DEFAULT_MODEL
    
    log_success "Модель $DEFAULT_MODEL установлена"
    
    # Список доступных моделей
    log_info "Доступные модели:"
    sudo -u $SERVICE_USER ollama list
}

# Создание рабочей директории
setup_directories() {
    log_step "Создание рабочих директорий..."
    
    mkdir -p $INSTALL_DIR
    mkdir -p $INSTALL_DIR/config
    mkdir -p $INSTALL_DIR/logs
    mkdir -p $INSTALL_DIR/scripts
    
    # Копирование файлов конфигурации
    if [[ -f docker-compose.yml ]]; then
        cp docker-compose.yml $INSTALL_DIR/
    fi
    
    chown -R $SERVICE_USER:$SERVICE_USER $INSTALL_DIR
    
    log_success "Директории созданы: $INSTALL_DIR"
}

# Создание Docker Compose файла
create_docker_compose() {
    log_step "Создание Docker Compose конфигурации..."
    
    cat > $INSTALL_DIR/docker-compose.yml << EOF
version: '3.8'

services:
  open-webui:
    image: ghcr.io/open-webui/open-webui:main
    container_name: ai-server-webui
    ports:
      - "$WEB_PORT:8080"
    environment:
      - OLLAMA_BASE_URL=http://host.docker.internal:$OLLAMA_PORT
      - WEBUI_SECRET_KEY=ai-server-secret-key-ubuntu-2024
      - WEBUI_NAME=AI Server Ubuntu
      - DEFAULT_LOCALE=ru-RU
      - DEFAULT_MODELS=$DEFAULT_MODEL
      - ENABLE_SIGNUP=true
      - ENABLE_LOGIN_FORM=true
      - SHOW_ADMIN_DETAILS=false
      - ENABLE_COMMUNITY_SHARING=false
    volumes:
      - open-webui-data:/app/backend/data
      - $INSTALL_DIR/config:/app/backend/config:ro
    extra_hosts:
      - "host.docker.internal:host-gateway"
    restart: unless-stopped
    networks:
      - ai-network
    deploy:
      resources:
        limits:
          cpus: '2.0'
          memory: 4G
        reservations:
          cpus: '0.5'
          memory: 1G

  # Мониторинг (опционально)
  portainer:
    image: portainer/portainer-ce:latest
    container_name: ai-server-portainer
    ports:
      - "9000:9000"
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock
      - portainer-data:/data
    restart: unless-stopped
    networks:
      - ai-network
    profiles:
      - monitoring

networks:
  ai-network:
    driver: bridge

volumes:
  open-webui-data:
    driver: local
  portainer-data:
    driver: local
EOF
    
    log_success "Docker Compose файл создан"
}

# Запуск веб-интерфейса
start_webui() {
    log_step "Запуск веб-интерфейса..."
    
    cd $INSTALL_DIR
    docker compose up -d
    
    # Проверка запуска
    sleep 15
    
    if docker compose ps | grep -q "running"; then
        log_success "Веб-интерфейс запущен"
        log_info "Доступен по адресу: http://$(hostname -I | awk '{print $1}'):$WEB_PORT"
    else
        log_error "Ошибка запуска веб-интерфейса"
        docker compose logs
        exit 1
    fi
}

# Создание скриптов управления
create_management_scripts() {
    log_step "Создание скриптов управления..."
    
    # Скрипт запуска
    cat > $INSTALL_DIR/scripts/start.sh << 'EOF'
#!/bin/bash
echo "🚀 Запуск AI Server..."

cd /opt/ai-server

# Проверка Ollama
if ! systemctl is-active --quiet ollama; then
    echo "Запуск Ollama..."
    sudo systemctl start ollama
    sleep 5
fi

# Запуск веб-интерфейса
echo "Запуск веб-интерфейса..."
docker compose up -d

echo "✅ AI Server запущен!"
echo "🌐 Веб-интерфейс: http://$(hostname -I | awk '{print $1}'):3000"
EOF

    # Скрипт остановки
    cat > $INSTALL_DIR/scripts/stop.sh << 'EOF'
#!/bin/bash
echo "⏹️  Остановка AI Server..."

cd /opt/ai-server

# Остановка веб-интерфейса
docker compose down

# Остановка Ollama (опционально)
# sudo systemctl stop ollama

echo "✅ AI Server остановлен!"
EOF

    # Скрипт статуса
    cat > $INSTALL_DIR/scripts/status.sh << 'EOF'
#!/bin/bash
echo "📊 Статус AI Server"
echo "==================="

# Статус Ollama
echo "🤖 Ollama:"
if systemctl is-active --quiet ollama; then
    echo "  ✅ Запущен"
    echo "  📍 http://localhost:11434"
else
    echo "  ❌ Остановлен"
fi

# Статус веб-интерфейса
echo ""
echo "🌐 Веб-интерфейс:"
cd /opt/ai-server
if docker compose ps | grep -q "running"; then
    echo "  ✅ Запущен"
    echo "  📍 http://$(hostname -I | awk '{print $1}'):3000"
else
    echo "  ❌ Остановлен"
fi

# Модели
echo ""
echo "📦 Установленные модели:"
sudo -u aiserver ollama list 2>/dev/null || echo "  ❌ Ollama недоступен"

# Ресурсы
echo ""
echo "💻 Использование ресурсов:"
echo "GPU:"
nvidia-smi --query-gpu=utilization.gpu,memory.used,memory.total --format=csv,noheader,nounits 2>/dev/null || echo "  ❌ NVIDIA GPU недоступен"

echo ""
echo "CPU/RAM:"
docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}" 2>/dev/null | grep ai-server || echo "  ❌ Контейнеры не запущены"
EOF

    # Установка прав
    chmod +x $INSTALL_DIR/scripts/*.sh
    
    # Создание символических ссылок
    ln -sf $INSTALL_DIR/scripts/start.sh /usr/local/bin/ai-start
    ln -sf $INSTALL_DIR/scripts/stop.sh /usr/local/bin/ai-stop
    ln -sf $INSTALL_DIR/scripts/status.sh /usr/local/bin/ai-status
    
    log_success "Скрипты управления созданы"
    log_info "Доступные команды: ai-start, ai-stop, ai-status"
}

# Настройка файрвола
setup_firewall() {
    log_step "Настройка файрвола..."
    
    if command -v ufw &> /dev/null; then
        ufw allow $WEB_PORT/tcp comment "AI Server Web UI"
        ufw allow $OLLAMA_PORT/tcp comment "Ollama API"
        ufw --force enable
        log_success "Файрвол настроен"
    else
        log_warning "UFW не найден, настройте файрвол вручную"
        log_info "Откройте порты: $WEB_PORT (веб-интерфейс), $OLLAMA_PORT (Ollama API)"
    fi
}

# Создание конфигурационного файла
create_config() {
    log_step "Создание конфигурационного файла..."
    
    cat > $INSTALL_DIR/config/server.conf << EOF
# AI Server Configuration
# RTX 2060 Optimized Settings

[ollama]
host = 0.0.0.0
port = $OLLAMA_PORT
keep_alive = 24h
num_parallel = 2
max_loaded_models = 2
flash_attention = true

[webui]
port = $WEB_PORT
default_model = $DEFAULT_MODEL
locale = ru-RU
enable_signup = true

[gpu]
# RTX 2060 имеет 6GB VRAM
# Рекомендуемые модели для RTX 2060:
# - llama3.2:1b (быстрая, ~1.3GB VRAM)
# - llama3.2:3b (оптимальная, ~2.0GB VRAM)  
# - codellama:7b (для кода, ~3.8GB VRAM)
# - mistral:7b (универсальная, ~4.1GB VRAM)

recommended_models = llama3.2:1b,llama3.2:3b,codellama:7b,mistral:7b
max_context_length = 4096
gpu_memory_fraction = 0.8

[performance]
# Оптимизация для RTX 2060
batch_size = 8
num_threads = 4
use_mlock = true
use_mmap = true
EOF
    
    log_success "Конфигурационный файл создан"
}

# Финальные настройки
finalize_installation() {
    log_step "Финальные настройки..."
    
    # Создание алиасов
    cat >> /home/$SUDO_USER/.bashrc << EOF

# AI Server aliases
alias ai-start='sudo /usr/local/bin/ai-start'
alias ai-stop='sudo /usr/local/bin/ai-stop'
alias ai-status='/usr/local/bin/ai-status'
alias ai-logs='cd /opt/ai-server && docker compose logs -f'
alias ai-models='sudo -u aiserver ollama list'
EOF
    
    # Информация о системе
    log_success "Установка завершена!"
    
    echo -e "${GREEN}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                   УСТАНОВКА ЗАВЕРШЕНА!                      ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    
    echo -e "${YELLOW}🌐 Веб-интерфейс:${NC} http://$(hostname -I | awk '{print $1}'):$WEB_PORT"
    echo -e "${YELLOW}🤖 Ollama API:${NC} http://$(hostname -I | awk '{print $1}'):$OLLAMA_PORT"
    echo -e "${YELLOW}📁 Рабочая директория:${NC} $INSTALL_DIR"
    echo ""
    echo -e "${CYAN}📋 Команды управления:${NC}"
    echo "  • ai-start    - запуск AI сервера"
    echo "  • ai-stop     - остановка AI сервера"
    echo "  • ai-status   - проверка статуса"
    echo "  • ai-logs     - просмотр логов"
    echo "  • ai-models   - список моделей"
    echo ""
    echo -e "${CYAN}📦 Установленные модели:${NC}"
    sudo -u $SERVICE_USER ollama list 2>/dev/null || echo "  Модели будут загружены при первом запуске"
    echo ""
    echo -e "${CYAN}🎯 Рекомендуемые модели для RTX 2060:${NC}"
    echo "  • ollama pull llama3.2:1b     (быстрая, 1.3GB)"
    echo "  • ollama pull llama3.2:3b     (сбалансированная, 2GB)"
    echo "  • ollama pull codellama:7b    (для кода, 3.8GB)"
    echo "  • ollama pull mistral:7b      (универсальная, 4.1GB)"
    echo ""
    echo -e "${GREEN}Перезайдите в терминал для применения алиасов!${NC}"
}

# ======================================
# ОСНОВНАЯ ЛОГИКА СКРИПТА
# ======================================

main() {
    print_banner
    
    # Проверки
    check_root
    detect_distro
    
    # Установка компонентов
    update_system
    install_nvidia_drivers
    install_docker
    install_nvidia_docker
    install_ollama
    
    # Настройка сервисов
    setup_directories
    create_docker_compose
    create_config
    setup_ollama_service
    install_models
    
    # Запуск сервисов
    start_webui
    create_management_scripts
    setup_firewall
    
    # Завершение
    finalize_installation
}

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