#!/bin/bash
# =============================================================================
# Ollama Models Backup & Restore Menu
# =============================================================================
# Purpose: Backup and restore Ollama models to avoid re-downloading
# Usage:   Run this script directly: ./ollama_backup_menu.sh
# =============================================================================

# Source the backup/restore functions
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/backup_restore.sh"

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

# Default paths (modify these to match your setup)
OLLAMA_MODELS_DIR="/usr/share/ollama/.ollama/models"  # Common Ollama models location
BACKUP_DIR="/Data/Backup/OllamaModels"
RESTORE_DIR="/usr/share/ollama/.ollama/models"

# Function to show menu
show_menu() {
    clear
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║${NC}         ${GREEN}Ollama Models Backup & Restore Tool${NC}              ${CYAN}║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "${YELLOW}Current Settings:${NC}"
    echo -e "  Models location: ${BLUE}$OLLAMA_MODELS_DIR${NC}"
    echo -e "  Backup storage:   ${BLUE}$BACKUP_DIR${NC}"
    echo ""
    echo -e "${CYAN}┌────────────────────────────────────────────────────────┐${NC}"
    echo -e "${CYAN}│${NC}  ${GREEN}1)${NC} Backup Ollama Models                              ${CYAN}│${NC}"
    echo -e "${CYAN}│${NC}  ${GREEN}2)${NC} Restore Ollama Models                             ${CYAN}│${NC}"
    echo -e "${CYAN}│${NC}  ${RED}3)${NC} Quit                                                ${CYAN}│${NC}"
    echo -e "${CYAN}└────────────────────────────────────────────────────────┘${NC}"
    echo ""
}

# Function to restart Ollama (called by trap)
restart_ollama() {
    echo ""
    echo -e "${YELLOW}Restarting Ollama service...${NC}"
    if sudo systemctl start ollama; then
        echo -e "${GREEN}✓ Ollama service restarted successfully${NC}"
    else
        echo -e "${RED}⚠ Failed to restart Ollama service. Please start it manually:${NC}"
        echo "  sudo systemctl start ollama"
    fi
    echo ""
    read -p "Press Enter to return to menu..."
}

# Function to backup
do_backup() {
    echo ""
    echo -e "${YELLOW}Backup Ollama Models${NC}"
    echo "--------------------"
    
    # Check if directory exists
    if [ ! -d "$OLLAMA_MODELS_DIR" ]; then
        echo -e "${RED}⚠ Ollama models directory not found!${NC}"
        read -p "Enter correct models path [default: $OLLAMA_MODELS_DIR]: " NEW_PATH
        [ -n "$NEW_PATH" ] && OLLAMA_MODELS_DIR="$NEW_PATH"
    fi
    
    # Check if Ollama is running
    echo ""
    if systemctl is-active --quiet ollama 2>/dev/null; then
        echo -e "${YELLOW}⚠ Ollama service is currently running.${NC}"
        echo -e "${YELLOW}  The service will be stopped during backup and restarted after.${NC}"
    else
        echo -e "${BLUE}ℹ Ollama service is not currently running.${NC}"
    fi
    
    # Generate backup filename with timestamp
    TIMESTAMP=$(date +%Y_%m_%d_%H.%M.%S)
    BACKUP_FILE="$BACKUP_DIR/OllamaModelsBackup_$TIMESTAMP.tar.gz"
    
    # Confirm backup
    echo ""
    echo -e "${YELLOW}Ready to backup:${NC}"
    echo -e "  Source: ${BLUE}$OLLAMA_MODELS_DIR${NC}"
    echo -e "  Destination: ${BLUE}$BACKUP_FILE${NC}"
    echo ""
    read -p "Proceed? (y/N): " CONFIRM
    
    if [[ "$CONFIRM" == "y" || "$CONFIRM" == "Y" ]]; then
        # Set trap to restart Ollama even if backup fails or is cancelled
        trap restart_ollama EXIT
        
        mkdir -p "$BACKUP_DIR"
        
        # Stop Ollama
        echo ""
        echo -e "${YELLOW}Stopping Ollama service...${NC}"
        if sudo systemctl stop ollama; then
            echo -e "${GREEN}✓ Ollama service stopped${NC}"
        else
            echo -e "${RED}⚠ Failed to stop Ollama service (may not be running)${NC}"
        fi
        
        # Wait a moment to ensure files are released
        sleep 2
        
        # Run backup
        echo ""
        echo -e "${YELLOW}Starting backup process...${NC}"
        backup_function "$OLLAMA_MODELS_DIR" "$BACKUP_FILE"
        
        # Backup complete, clear trap and restart manually
        trap - EXIT
        restart_ollama
    else
        echo "Backup cancelled."
        echo ""
        read -p "Press Enter to return to menu..."
    fi
}

# Function to restore
do_restore() {
    echo ""
    echo -e "${YELLOW}Restore Ollama Models${NC}"
    echo "---------------------"
    
    # Check if models directory exists
    if [ ! -d "$RESTORE_DIR" ]; then
        echo -e "${RED}⚠ Ollama models directory not found!${NC}"
        read -p "Enter correct restore path [default: $RESTORE_DIR]: " NEW_PATH
        [ -n "$NEW_PATH" ] && RESTORE_DIR="$NEW_PATH"
    fi
    
    # Create directory if needed
    mkdir -p "$RESTORE_DIR"
    
    # Call restore function
    restore_function "$RESTORE_DIR"
    
    echo ""
    read -p "Press Enter to return to menu..."
}

# Main loop
while true; do
    show_menu
    read -p "Select option (1-3): " CHOICE
    
    case $CHOICE in
        1)
            do_backup
            ;;
        2)
            do_restore
            ;;
        3|q|Q)
            clear
            echo -e "${GREEN}Goodbye!${NC}"
            echo "Remember: No more re-downloading Ollama models! 😄"
            exit 0
            ;;
        *)
            echo -e "${RED}Invalid option!${NC}"
            sleep 1
            ;;
    esac
done