#!/bin/bash

# AI Server Model Management Script
# Optimized for RTX 2060

# Colors
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
NC='\033[0m'

show_menu() {
    echo -e "${BLUE}🤖 AI Server - Управление моделями${NC}"
    echo "=================================="
    echo "1. Показать установленные модели"
    echo "2. Установить рекомендуемые модели для RTX 2060"
    echo "3. Установить конкретную модель"
    echo "4. Удалить модель"
    echo "5. Тестировать модель"
    echo "6. Показать использование VRAM"
    echo "7. Рекомендации по моделям"
    echo "8. Выход"
    echo ""
}

list_models() {
    echo -e "${PURPLE}📦 Установленные модели:${NC}"
    sudo -u aiserver OLLAMA_HOME=/opt/ai-server/.ollama ollama list
    echo ""
    
    if command -v nvidia-smi &> /dev/null; then
        echo -e "${PURPLE}💾 Использование VRAM:${NC}"
        nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader,nounits | \
            awk -F', ' '{printf "Используется: %d MB / %d MB (%.1f%%)\n", $1, $2, ($1/$2)*100}'
    fi
    echo ""
}

install_recommended() {
    echo -e "${YELLOW}🎯 Установка рекомендуемых моделей для RTX 2060...${NC}"
    echo ""
    
    models=("llama3.2:1b" "llama3.2:3b" "codellama:7b")
    descriptions=("Быстрая модель (1.3GB)" "Оптимальная модель (2GB)" "Модель для кода (3.8GB)")
    
    for i in "${!models[@]}"; do
        model="${models[$i]}"
        desc="${descriptions[$i]}"
        
        echo -e "${BLUE}Устанавливаю: $model - $desc${NC}"
        sudo -u aiserver OLLAMA_HOME=/opt/ai-server/.ollama ollama pull "$model"
        
        if [ $? -eq 0 ]; then
            echo -e "${GREEN}✅ $model установлена${NC}"
        else
            echo -e "${RED}❌ Ошибка установки $model${NC}"
        fi
        echo ""
    done
}

install_custom() {
    echo -e "${YELLOW}Введите название модели (например: llama3.2:3b):${NC}"
    read -r model_name
    
    if [[ -z "$model_name" ]]; then
        echo -e "${RED}❌ Название модели не может быть пустым${NC}"
        return
    fi
    
    echo -e "${BLUE}Установка модели: $model_name${NC}"
    echo -e "${YELLOW}Это может занять несколько минут...${NC}"
    
    sudo -u aiserver OLLAMA_HOME=/opt/ai-server/.ollama ollama pull "$model_name"
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✅ Модель $model_name успешно установлена${NC}"
    else
        echo -e "${RED}❌ Ошибка установки модели $model_name${NC}"
    fi
}

remove_model() {
    echo -e "${PURPLE}📦 Доступные модели для удаления:${NC}"
    sudo -u aiserver OLLAMA_HOME=/opt/ai-server/.ollama ollama list
    echo ""
    
    echo -e "${YELLOW}Введите название модели для удаления:${NC}"
    read -r model_name
    
    if [[ -z "$model_name" ]]; then
        echo -e "${RED}❌ Название модели не может быть пустым${NC}"
        return
    fi
    
    echo -e "${YELLOW}⚠️  Удалить модель $model_name? (y/N):${NC}"
    read -r confirm
    
    if [[ $confirm =~ ^[Yy]$ ]]; then
        sudo -u aiserver OLLAMA_HOME=/opt/ai-server/.ollama ollama rm "$model_name"
        
        if [ $? -eq 0 ]; then
            echo -e "${GREEN}✅ Модель $model_name удалена${NC}"
        else
            echo -e "${RED}❌ Ошибка удаления модели $model_name${NC}"
        fi
    else
        echo -e "${BLUE}Отменено${NC}"
    fi
}

test_model() {
    echo -e "${PURPLE}📦 Доступные модели:${NC}"
    sudo -u aiserver OLLAMA_HOME=/opt/ai-server/.ollama ollama list
    echo ""
    
    echo -e "${YELLOW}Введите название модели для тестирования:${NC}"
    read -r model_name
    
    if [[ -z "$model_name" ]]; then
        echo -e "${RED}❌ Название модели не может быть пустым${NC}"
        return
    fi
    
    echo -e "${BLUE}Тестирование модели: $model_name${NC}"
    echo -e "${YELLOW}Задаю тестовый вопрос...${NC}"
    
    sudo -u aiserver OLLAMA_HOME=/opt/ai-server/.ollama ollama run "$model_name" "Привет! Как дела? Ответь кратко."
}

show_vram_usage() {
    if command -v nvidia-smi &> /dev/null; then
        echo -e "${PURPLE}🎮 Информация о GPU:${NC}"
        nvidia-smi --query-gpu=name,memory.used,memory.total,utilization.gpu,temperature.gpu --format=csv,noheader
        echo ""
        
        echo -e "${PURPLE}📊 Детальное использование VRAM:${NC}"
        nvidia-smi
    else
        echo -e "${RED}❌ NVIDIA GPU не найден или драйверы не установлены${NC}"
    fi
}

show_recommendations() {
    echo -e "${BLUE}🎯 Рекомендации по моделям для RTX 2060 (6GB VRAM):${NC}"
    echo ""
    
    echo -e "${GREEN}🥇 ЛУЧШИЙ ВЫБОР:${NC}"
    echo "  • llama3.2:3b (2GB VRAM) - оптимальный баланс скорости и качества"
    echo ""
    
    echo -e "${YELLOW}🥈 БЫСТРЫЕ МОДЕЛИ:${NC}"
    echo "  • llama3.2:1b (1.3GB VRAM) - максимальная скорость"
    echo ""
    
    echo -e "${PURPLE}🥉 СПЕЦИАЛИЗИРОВАННЫЕ:${NC}"
    echo "  • codellama:7b (3.8GB VRAM) - для программирования"
    echo "  • mistral:7b (4.1GB VRAM) - высокое качество (почти предел)"
    echo ""
    
    echo -e "${RED}❌ НЕ РЕКОМЕНДУЕТСЯ:${NC}"
    echo "  • llama3:8b - требует 8GB+ VRAM"
    echo "  • llama3:70b - требует 40GB+ VRAM"
    echo "  • codellama:34b - требует 20GB+ VRAM"
    echo ""
    
    echo -e "${CYAN}💡 СОВЕТЫ:${NC}"
    echo "  • Не устанавливайте много больших моделей одновременно"
    echo "  • Мониторьте температуру GPU (ai-status)"
    echo "  • При перегреве используйте меньшие модели"
    echo "  • Для продакшена рекомендуется llama3.2:3b"
}

# Основной цикл
main() {
    while true; do
        show_menu
        echo -e "${YELLOW}Выберите опцию (1-8):${NC}"
        read -r choice
        
        case $choice in
            1)
                list_models
                ;;
            2)
                install_recommended
                ;;
            3)
                install_custom
                ;;
            4)
                remove_model
                ;;
            5)
                test_model
                ;;
            6)
                show_vram_usage
                ;;
            7)
                show_recommendations
                ;;
            8)
                echo -e "${GREEN}👋 До свидания!${NC}"
                exit 0
                ;;
            *)
                echo -e "${RED}❌ Неверный выбор. Попробуйте снова.${NC}"
                ;;
        esac
        
        echo ""
        echo -e "${CYAN}Нажмите Enter для продолжения...${NC}"
        read -r
        clear
    done
}

# Проверка прав
if [[ ! -d "/opt/ai-server" ]]; then
    echo -e "${RED}❌ AI Server не установлен. Запустите сначала install-ubuntu.sh${NC}"
    exit 1
fi

clear
main