#!/bin/bash

# Script de sauvegarde Ollama + WebUI vers partage Samba via rsync
# Auteur: Script généré par Claude
# Date: $(date +%Y-%m-%d)

set -euo pipefail

# =============================================================================
# CONFIGURATION
# =============================================================================

# Configuration Ollama
OLLAMA_HOME="/usr/share/ollama"
OLLAMA_DATA_DIR="/var/lib/ollama"
OLLAMA_MODELS_DIR="/usr/share/ollama/.ollama/models"
OLLAMA_CONFIG_DIR="/etc/ollama"
OLLAMA_USER="ollama"
OLLAMA_SERVICE="ollama"

# Configuration WebUI (Open WebUI / Ollama WebUI)
WEBUI_DATA_DIR="/opt/open-webui"
WEBUI_CONFIG_DIR="/etc/open-webui"
WEBUI_DATABASE_DIR="/opt/open-webui/data"
WEBUI_UPLOADS_DIR="/opt/open-webui/data/uploads"
WEBUI_SERVICE="open-webui"
WEBUI_USER="webui"

# Chemins alternatifs courants pour WebUI (Docker)
DOCKER_WEBUI_VOLUME="open-webui"
DOCKER_COMPOSE_DIR="/opt/open-webui"

# Configuration du partage Samba
SAMBA_SERVER="192.168.77.230"
SAMBA_SHARE="/srv/samba/partage"
SAMBA_USER="backup_user"  # À adapter selon votre configuration

# Configuration SSH
SSH_KEY_PATH="$HOME/.ssh/ollama_backup_ed25519"
SSH_OPTIONS="-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR"

# Configuration de sauvegarde
BACKUP_BASE_DIR="/tmp/ollama_webui_backup_$(date +%Y%m%d_%H%M%S)"
REMOTE_BACKUP_DIR="ollama_webui_backups"
LOG_FILE="/var/log/ollama_webui_backup.log"
RETENTION_DAYS=30

# Options de sauvegarde
BACKUP_MODELS=true          # Sauvegarder les modèles (peut être volumineux)
BACKUP_CONVERSATIONS=true   # Sauvegarder les conversations WebUI
BACKUP_SETTINGS=true        # Sauvegarder les paramètres et configurations
COMPRESS_MODELS=true        # Compresser les modèles pour économiser l'espace

# =============================================================================
# FONCTIONS UTILITAIRES
# =============================================================================

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

error_exit() {
    log "ERREUR: $1"
    exit 1
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        error_exit "Ce script doit être exécuté en tant que root"
    fi
}

check_command() {
    if ! command -v "$1" &> /dev/null; then
        return 1
    fi
    return 0
}

# =============================================================================
# INSTALLATION DES DÉPENDANCES
# =============================================================================

install_dependencies() {
    log "Vérification et installation des dépendances..."
    
    # Détecter la distribution
    if [[ -f /etc/debian_version ]]; then
        DISTRO="debian"
        PKG_MANAGER="apt-get"
        PACKAGES="rsync samba-client cifs-utils openssh-client docker.io docker-compose jq sqlite3 gzip"
    elif [[ -f /etc/redhat-release ]]; then
        DISTRO="redhat"
        PKG_MANAGER="yum"
        PACKAGES="rsync cifs-utils openssh-clients docker docker-compose jq sqlite gzip"
    else
        error_exit "Distribution non supportée"
    fi
    
    log "Distribution détectée: $DISTRO"
    
    # Mise à jour du cache des paquets
    if [[ "$DISTRO" == "debian" ]]; then
        apt-get update -qq
    fi
    
    # Installation des paquets manquants
    for package in $PACKAGES; do
        if [[ "$DISTRO" == "debian" ]]; then
            if ! dpkg -l | grep -q "^ii  $package "; then
                log "Installation de $package..."
                apt-get install -y -qq "$package" || log "Impossible d'installer $package (peut-être optionnel)"
            else
                log "$package est déjà installé"
            fi
        elif [[ "$DISTRO" == "redhat" ]]; then
            if ! rpm -q "$package" >/dev/null 2>&1; then
                log "Installation de $package..."
                yum install -y "$package" || log "Impossible d'installer $package (peut-être optionnel)"
            else
                log "$package est déjà installé"
            fi
        fi
    done
}

# =============================================================================
# GESTION DES CERTIFICATS SSH
# =============================================================================

setup_ssh_key() {
    log "Configuration de la clé SSH ED25519..."
    
    if [[ ! -f "$SSH_KEY_PATH" ]]; then
        log "Création de la clé SSH ED25519..."
        ssh-keygen -t ed25519 -f "$SSH_KEY_PATH" -N "" -C "ollama-webui-backup-$(hostname)"
        chmod 600 "$SSH_KEY_PATH"
        chmod 644 "${SSH_KEY_PATH}.pub"
        
        log "Clé SSH créée: $SSH_KEY_PATH"
        log "Clé publique: ${SSH_KEY_PATH}.pub"
        echo
        log "IMPORTANT: Vous devez copier cette clé publique sur le serveur distant:"
        log "=========================================="
        cat "${SSH_KEY_PATH}.pub"
        log "=========================================="
        echo
        log "Commande pour copier la clé sur le serveur distant:"
        log "ssh-copy-id -i $SSH_KEY_PATH ${SAMBA_USER}@${SAMBA_SERVER}"
        echo
        read -p "Appuyez sur Entrée une fois la clé configurée sur le serveur distant..."
    else
        log "Clé SSH existante trouvée: $SSH_KEY_PATH"
    fi
}

test_ssh_connection() {
    log "Test de la connexion SSH..."
    if ssh -i "$SSH_KEY_PATH" $SSH_OPTIONS "${SAMBA_USER}@${SAMBA_SERVER}" "echo 'Connexion SSH réussie'" 2>/dev/null; then
        log "Connexion SSH: OK"
    else
        error_exit "Impossible de se connecter via SSH. Vérifiez la configuration de la clé."
    fi
}

# =============================================================================
# DÉTECTION DE L'ENVIRONNEMENT
# =============================================================================

detect_environment() {
    log "Détection de l'environnement Ollama et WebUI..."
    
    # Détection d'Ollama
    OLLAMA_INSTALL_TYPE="none"
    if systemctl is-enabled "$OLLAMA_SERVICE" >/dev/null 2>&1; then
        OLLAMA_INSTALL_TYPE="systemd"
        log "Ollama détecté: installation systemd"
    elif check_command ollama; then
        OLLAMA_INSTALL_TYPE="binary"
        log "Ollama détecté: installation binaire"
    elif docker ps --format "table {{.Names}}" | grep -q ollama; then
        OLLAMA_INSTALL_TYPE="docker"
        log "Ollama détecté: conteneur Docker"
    else
        log "ATTENTION: Ollama non détecté"
    fi
    
    # Détection du WebUI
    WEBUI_INSTALL_TYPE="none"
    if systemctl is-enabled "$WEBUI_SERVICE" >/dev/null 2>&1; then
        WEBUI_INSTALL_TYPE="systemd"
        log "WebUI détecté: installation systemd"
    elif docker ps --format "table {{.Names}}" | grep -q -E "(webui|open-webui)"; then
        WEBUI_INSTALL_TYPE="docker"
        log "WebUI détecté: conteneur Docker"
    elif [[ -d "$DOCKER_COMPOSE_DIR" ]]; then
        WEBUI_INSTALL_TYPE="docker-compose"
        log "WebUI détecté: Docker Compose"
    else
        log "ATTENTION: WebUI non détecté"
    fi
    
    # Adaptation des chemins selon l'environnement
    if [[ "$OLLAMA_INSTALL_TYPE" == "docker" ]]; then
        OLLAMA_DATA_DIR="/var/lib/docker/volumes/ollama/_data"
        OLLAMA_MODELS_DIR="$OLLAMA_DATA_DIR/models"
    fi
    
    if [[ "$WEBUI_INSTALL_TYPE" == "docker" ]]; then
        WEBUI_DATA_DIR="/var/lib/docker/volumes/${DOCKER_WEBUI_VOLUME}/_data"
        WEBUI_DATABASE_DIR="$WEBUI_DATA_DIR"
    fi
}

# =============================================================================
# GESTION DES SERVICES
# =============================================================================

stop_services() {
    log "Arrêt des services..."
    
    OLLAMA_WAS_RUNNING=false
    WEBUI_WAS_RUNNING=false
    
    # Arrêt d'Ollama
    case "$OLLAMA_INSTALL_TYPE" in
        "systemd")
            if systemctl is-active --quiet "$OLLAMA_SERVICE"; then
                systemctl stop "$OLLAMA_SERVICE"
                OLLAMA_WAS_RUNNING=true
                log "Service Ollama arrêté"
            fi
            ;;
        "docker")
            if docker ps -q -f name=ollama | grep -q .; then
                docker stop ollama >/dev/null 2>&1
                OLLAMA_WAS_RUNNING=true
                log "Conteneur Ollama arrêté"
            fi
            ;;
    esac
    
    # Arrêt du WebUI
    case "$WEBUI_INSTALL_TYPE" in
        "systemd")
            if systemctl is-active --quiet "$WEBUI_SERVICE"; then
                systemctl stop "$WEBUI_SERVICE"
                WEBUI_WAS_RUNNING=true
                log "Service WebUI arrêté"
            fi
            ;;
        "docker")
            if docker ps -q -f name=webui | grep -q .; then
                docker stop webui >/dev/null 2>&1 || docker stop open-webui >/dev/null 2>&1
                WEBUI_WAS_RUNNING=true
                log "Conteneur WebUI arrêté"
            fi
            ;;
        "docker-compose")
            if [[ -f "$DOCKER_COMPOSE_DIR/docker-compose.yml" ]]; then
                cd "$DOCKER_COMPOSE_DIR"
                docker-compose stop >/dev/null 2>&1
                WEBUI_WAS_RUNNING=true
                log "Services Docker Compose arrêtés"
            fi
            ;;
    esac
}

start_services() {
    log "Redémarrage des services..."
    
    # Redémarrage d'Ollama
    if [[ "$OLLAMA_WAS_RUNNING" == "true" ]]; then
        case "$OLLAMA_INSTALL_TYPE" in
            "systemd")
                systemctl start "$OLLAMA_SERVICE"
                log "Service Ollama redémarré"
                ;;
            "docker")
                docker start ollama >/dev/null 2>&1
                log "Conteneur Ollama redémarré"
                ;;
        esac
    fi
    
    # Redémarrage du WebUI
    if [[ "$WEBUI_WAS_RUNNING" == "true" ]]; then
        case "$WEBUI_INSTALL_TYPE" in
            "systemd")
                systemctl start "$WEBUI_SERVICE"
                log "Service WebUI redémarré"
                ;;
            "docker")
                docker start webui >/dev/null 2>&1 || docker start open-webui >/dev/null 2>&1
                log "Conteneur WebUI redémarré"
                ;;
            "docker-compose")
                if [[ -f "$DOCKER_COMPOSE_DIR/docker-compose.yml" ]]; then
                    cd "$DOCKER_COMPOSE_DIR"
                    docker-compose start >/dev/null 2>&1
                    log "Services Docker Compose redémarrés"
                fi
                ;;
        esac
    fi
}

# =============================================================================
# FONCTIONS DE SAUVEGARDE
# =============================================================================

create_backup_structure() {
    log "Création de la structure de sauvegarde: $BACKUP_BASE_DIR"
    mkdir -p "$BACKUP_BASE_DIR"/{ollama,webui,metadata}
}

backup_ollama() {
    log "Sauvegarde d'Ollama..."
    
    # Sauvegarde des données Ollama
    if [[ -d "$OLLAMA_DATA_DIR" && "$BACKUP_SETTINGS" == "true" ]]; then
        log "Sauvegarde des données Ollama..."
        rsync -av --progress "$OLLAMA_DATA_DIR/" "$BACKUP_BASE_DIR/ollama/data/" \
            --exclude="models" --exclude="*.log"
    fi
    
    # Sauvegarde des modèles
    if [[ "$BACKUP_MODELS" == "true" && -d "$OLLAMA_MODELS_DIR" ]]; then
        log "Sauvegarde des modèles Ollama..."
        if [[ "$COMPRESS_MODELS" == "true" ]]; then
            log "Compression des modèles en cours..."
            tar -czf "$BACKUP_BASE_DIR/ollama/models.tar.gz" -C "$(dirname "$OLLAMA_MODELS_DIR")" "$(basename "$OLLAMA_MODELS_DIR")"
        else
            rsync -av --progress "$OLLAMA_MODELS_DIR/" "$BACKUP_BASE_DIR/ollama/models/"
        fi
    fi
    
    # Sauvegarde de la configuration
    if [[ -d "$OLLAMA_CONFIG_DIR" && "$BACKUP_SETTINGS" == "true" ]]; then
        log "Sauvegarde de la configuration Ollama..."
        rsync -av --progress "$OLLAMA_CONFIG_DIR/" "$BACKUP_BASE_DIR/ollama/config/"
    fi
    
    # Liste des modèles installés
    if check_command ollama; then
        log "Création de la liste des modèles installés..."
        ollama list > "$BACKUP_BASE_DIR/ollama/installed_models.txt" 2>/dev/null || true
    fi
}

backup_webui() {
    log "Sauvegarde du WebUI..."
    
    # Sauvegarde des données WebUI
    if [[ -d "$WEBUI_DATA_DIR" && "$BACKUP_CONVERSATIONS" == "true" ]]; then
        log "Sauvegarde des données WebUI..."
        rsync -av --progress "$WEBUI_DATA_DIR/" "$BACKUP_BASE_DIR/webui/data/" \
            --exclude="*.log" --exclude="cache"
    fi
    
    # Sauvegarde de la base de données
    if [[ -f "$WEBUI_DATABASE_DIR/webui.db" ]]; then
        log "Sauvegarde de la base de données WebUI..."
        cp "$WEBUI_DATABASE_DIR/webui.db" "$BACKUP_BASE_DIR/webui/webui.db"
        
        # Export des conversations en JSON si sqlite3 est disponible
        if check_command sqlite3; then
            log "Export des conversations en JSON..."
            sqlite3 "$WEBUI_DATABASE_DIR/webui.db" \
                ".mode json" \
                ".output $BACKUP_BASE_DIR/webui/conversations_export.json" \
                "SELECT * FROM chat;" 2>/dev/null || true
        fi
    fi
    
    # Sauvegarde des uploads/fichiers
    if [[ -d "$WEBUI_UPLOADS_DIR" ]]; then
        log "Sauvegarde des fichiers uploadés..."
        rsync -av --progress "$WEBUI_UPLOADS_DIR/" "$BACKUP_BASE_DIR/webui/uploads/"
    fi
    
    # Sauvegarde de la configuration WebUI
    if [[ -d "$WEBUI_CONFIG_DIR" && "$BACKUP_SETTINGS" == "true" ]]; then
        log "Sauvegarde de la configuration WebUI..."
        rsync -av --progress "$WEBUI_CONFIG_DIR/" "$BACKUP_BASE_DIR/webui/config/"
    fi
    
    # Sauvegarde des fichiers Docker Compose si applicable
    if [[ "$WEBUI_INSTALL_TYPE" == "docker-compose" && -d "$DOCKER_COMPOSE_DIR" ]]; then
        log "Sauvegarde de la configuration Docker Compose..."
        rsync -av --progress "$DOCKER_COMPOSE_DIR/" "$BACKUP_BASE_DIR/webui/docker-compose/" \
            --include="*.yml" --include="*.yaml" --include="*.env" --exclude="*"
    fi
}

create_metadata() {
    log "Création des métadonnées de sauvegarde..."
    
    # Informations système
    cat > "$BACKUP_BASE_DIR/metadata/system_info.txt" << EOF
Sauvegarde Ollama + WebUI
=========================
Date: $(date)
Serveur: $(hostname)
Version script: 1.0
Utilisateur: $(whoami)

Environnement détecté:
- Ollama: $OLLAMA_INSTALL_TYPE
- WebUI: $WEBUI_INSTALL_TYPE

Répertoires sauvegardés:
- Ollama Data: $OLLAMA_DATA_DIR
- Ollama Models: $OLLAMA_MODELS_DIR
- WebUI Data: $WEBUI_DATA_DIR
- WebUI Database: $WEBUI_DATABASE_DIR

Options de sauvegarde:
- Modèles: $BACKUP_MODELS
- Conversations: $BACKUP_CONVERSATIONS
- Paramètres: $BACKUP_SETTINGS
- Compression modèles: $COMPRESS_MODELS
EOF

    # Informations sur les services Docker si applicable
    if check_command docker; then
        log "Sauvegarde des informations Docker..."
        docker ps -a --format "table {{.Names}}\t{{.Image}}\t{{.Status}}" > "$BACKUP_BASE_DIR/metadata/docker_containers.txt" 2>/dev/null || true
        docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}" > "$BACKUP_BASE_DIR/metadata/docker_images.txt" 2>/dev/null || true
    fi
    
    # Taille de la sauvegarde
    BACKUP_SIZE=$(du -sh "$BACKUP_BASE_DIR" | cut -f1)
    echo "Taille totale de la sauvegarde: $BACKUP_SIZE" >> "$BACKUP_BASE_DIR/metadata/system_info.txt"
    
    log "Sauvegarde locale créée dans: $BACKUP_BASE_DIR (Taille: $BACKUP_SIZE)"
}

sync_to_remote() {
    log "Synchronisation vers le serveur distant..."
    
    # Création du répertoire distant si nécessaire
    ssh -i "$SSH_KEY_PATH" $SSH_OPTIONS "${SAMBA_USER}@${SAMBA_SERVER}" \
        "mkdir -p ${SAMBA_SHARE}/${REMOTE_BACKUP_DIR}/$(date +%Y/%m)"
    
    # Synchronisation des fichiers
    REMOTE_PATH="${SAMBA_USER}@${SAMBA_SERVER}:${SAMBA_SHARE}/${REMOTE_BACKUP_DIR}/$(date +%Y/%m)/ollama_webui_backup_$(date +%Y%m%d_%H%M%S)/"
    
    log "Envoi vers: $REMOTE_PATH"
    rsync -avz --progress --delete \
          -e "ssh -i $SSH_KEY_PATH $SSH_OPTIONS" \
          "$BACKUP_BASE_DIR/" \
          "$REMOTE_PATH"
    
    log "Synchronisation terminée"
}

cleanup_local() {
    log "Nettoyage des fichiers temporaires..."
    if [[ -d "$BACKUP_BASE_DIR" ]]; then
        rm -rf "$BACKUP_BASE_DIR"
        log "Répertoire temporaire supprimé: $BACKUP_BASE_DIR"
    fi
}

cleanup_remote() {
    log "Nettoyage des anciennes sauvegardes (plus de $RETENTION_DAYS jours)..."
    
    ssh -i "$SSH_KEY_PATH" $SSH_OPTIONS "${SAMBA_USER}@${SAMBA_SERVER}" \
        "find ${SAMBA_SHARE}/${REMOTE_BACKUP_DIR} -type d -name 'ollama_webui_backup_*' -mtime +${RETENTION_DAYS} -exec rm -rf {} + 2>/dev/null || true"
    
    log "Nettoyage des anciennes sauvegardes terminé"
}

# =============================================================================
# FONCTION PRINCIPALE
# =============================================================================

main() {
    log "=== DÉBUT DE LA SAUVEGARDE OLLAMA + WEBUI ==="
    
    # Vérifications préliminaires
    check_root
    
    # Installation des dépendances
    install_dependencies
    
    # Configuration SSH
    setup_ssh_key
    test_ssh_connection
    
    # Détection de l'environnement
    detect_environment
    
    # Variables pour suivre l'état des services
    OLLAMA_WAS_RUNNING=false
    WEBUI_WAS_RUNNING=false
    
    # Gestion des erreurs avec nettoyage
    trap 'cleanup_local; start_services' EXIT ERR
    
    # Arrêt des services pour une sauvegarde cohérente
    stop_services
    
    # Création de la structure de sauvegarde
    create_backup_structure
    
    # Sauvegarde d'Ollama
    backup_ollama
    
    # Sauvegarde du WebUI
    backup_webui
    
    # Création des métadonnées
    create_metadata
    
    # Redémarrage des services dès que possible
    start_services
    
    # Synchronisation vers le serveur distant
    sync_to_remote
    
    # Nettoyage
    cleanup_remote
    cleanup_local
    
    log "=== SAUVEGARDE OLLAMA + WEBUI TERMINÉE AVEC SUCCÈS ==="
}

# =============================================================================
# POINT D'ENTRÉE
# =============================================================================

# Vérification des arguments pour les options avancées
while [[ $# -gt 0 ]]; do
    case $1 in
        --no-models)
            BACKUP_MODELS=false
            log "Option: Sauvegarde des modèles désactivée"
            shift
            ;;
        --no-conversations)
            BACKUP_CONVERSATIONS=false
            log "Option: Sauvegarde des conversations désactivée"
            shift
            ;;
        --no-compression)
            COMPRESS_MODELS=false
            log "Option: Compression des modèles désactivée"
            shift
            ;;
        --help)
            echo "Usage: $0 [OPTIONS]"
            echo "Options:"
            echo "  --no-models        Ne pas sauvegarder les modèles Ollama"
            echo "  --no-conversations Ne pas sauvegarder les conversations WebUI"
            echo "  --no-compression   Ne pas compresser les modèles"
            echo "  --help            Afficher cette aide"
            exit 0
            ;;
        *)
            error_exit "Option inconnue: $1"
            ;;
    esac
done

# Création du fichier de log si nécessaire
mkdir -p "$(dirname "$LOG_FILE")"
touch "$LOG_FILE"

# Lancement du script principal
main "$@"