#!/usr/bin/env bash
set -Eeuo pipefail

# Couleurs
BLUE="\033[1;34m"
RED="\033[1;31m"
GREEN="\033[1;32m"
NC="\033[0m"

# ------------------------------------------------------------
# Fonctions d'affichage
# ------------------------------------------------------------
print_message() {
    echo -e "${BLUE}>>> $*${NC}"
}

print_error() {
    echo -e "${RED}>>> Erreur: $*${NC}" >&2
}

print_success() {
    echo -e "${GREEN}>>> Succès: $*${NC}"
}

# ------------------------------------------------------------
# Nettoyage en cas d'erreur ou d'interruption
# ------------------------------------------------------------
cleanup() {
    # Action à réaliser avant de quitter (si nécessaire)
    echo "Nettoyage en cours..."
    # Ex. : pkill -f "ollama" si besoin
}

trap cleanup SIGINT SIGTERM ERR

# ------------------------------------------------------------
# Vérifications système de base
# ------------------------------------------------------------
check_macos() {
    if [[ "$(uname)" != "Darwin" ]]; then
        print_error "Ce script est prévu pour macOS. Annulation."
        exit 1
    fi
}

# ------------------------------------------------------------
# Vérifier/install Homebrew
# ------------------------------------------------------------
check_homebrew() {
    if command -v brew >/dev/null 2>&1; then
        print_message "Homebrew est déjà installé"
        return 0
    fi
    return 1
}

install_homebrew() {
    print_message "Installation de Homebrew..."
    /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" || {
        print_error "Impossible d'installer Homebrew"
        exit 1
    }
    print_success "Homebrew a été installé avec succès"
}

# ------------------------------------------------------------
# Vérifier/install Docker
# ------------------------------------------------------------
check_docker() {
    if command -v docker >/dev/null 2>&1; then
        print_message "Docker est déjà installé"
        return 0
    fi
    return 1
}

check_docker_running() {
    if docker info >/dev/null 2>&1; then
        print_message "Docker est déjà en cours d'exécution"
        return 0
    fi
    return 1
}

start_docker() {
    print_message "Démarrage de Docker..."
    open -a Docker
    
    # Attendre que Docker démarre
    local max_attempts=30
    local attempt=1
    
    while ! docker info >/dev/null 2>&1; do
        if [ $attempt -gt $max_attempts ]; then
            print_error "Docker n'a pas démarré dans le temps imparti"
            return 1
        fi
        echo -n "."
        sleep 1
        ((attempt++))
    done
    
    echo
    print_success "Docker a démarré avec succès"
    return 0
}

install_docker() {
    print_message "Installation de Docker..."
    brew install --cask docker || {
        print_error "Impossible d'installer Docker"
        exit 1
    }
    print_success "Docker a été installé avec succès"
}

# ------------------------------------------------------------
# Vérifier/install Ollama
# ------------------------------------------------------------
check_ollama() {
    if command -v ollama >/dev/null 2>&1; then
        print_message "Ollama est déjà installé"
        return 0
    fi
    return 1
}

check_ollama_running() {
    if pgrep -x "ollama" >/dev/null; then
        print_message "Ollama est déjà en cours d'exécution"
        return 0
    fi
    return 1
}

start_ollama() {
    print_message "Démarrage d'Ollama..."
    
    # Lancer Ollama en arrière-plan
    ollama serve >/dev/null 2>&1 &
    
    # Attendre que le service soit prêt
    local max_attempts=30
    local attempt=1
    
    while ! curl -s http://localhost:11434/api/version >/dev/null 2>&1; do
        if [ $attempt -gt $max_attempts ]; then
            print_error "Ollama n'a pas démarré dans le temps imparti"
            return 1
        fi
        echo -n "."
        sleep 1
        ((attempt++))
    done
    
    echo
    print_success "Ollama a démarré avec succès"
    return 0
}

install_ollama() {
    print_message "Installation d'Ollama..."
    curl -fsSL https://ollama.com/install.sh | sh || {
        print_error "Impossible d'installer Ollama"
        exit 1
    }
    print_success "Ollama a été installé avec succès"
}

# ------------------------------------------------------------
# Installer Open Web UI (docker container)
# ------------------------------------------------------------
install_open_webui() {
    print_message "Installation de Open Web UI..."

    if ! check_docker_running; then
        print_message "Docker n'est pas en cours d'exécution"
        start_docker || {
            print_error "Impossible de démarrer Docker"
            exit 1
        }
    fi

    # Vérifier si le conteneur existe et le supprimer si nécessaire
    CONTAINER_ID=$(docker ps -aq -f name=open-webui)
    if [ -n "$CONTAINER_ID" ]; then
        print_message "Suppression de l'ancien conteneur Open Web UI..."
        docker stop "$CONTAINER_ID" >/dev/null 2>&1
        docker rm "$CONTAINER_ID" >/dev/null 2>&1
        if [ $? -ne 0 ]; then
            print_error "Impossible de supprimer l'ancien conteneur. Tentative de suppression forcée..."
            docker rm -f "$CONTAINER_ID" >/dev/null 2>&1 || {
                print_error "Impossible de supprimer le conteneur. Veuillez exécuter manuellement :"
                echo "docker rm -f $CONTAINER_ID"
                exit 1
            }
        fi
    fi

    print_message "Lancement d'un nouveau conteneur open-webui..."
    docker run -d -p 3000:8080 \
        --add-host=host.docker.internal:host-gateway \
        -v open-webui:/app/backend/data \
        --name open-webui \
        --restart always \
        ghcr.io/open-webui/open-webui:main \
    || { print_error "Erreur lors du lancement de Open Web UI"; exit 1; }

    # Attendre que le conteneur soit prêt
    print_message "Attente du démarrage de Open Web UI..."
    for i in {1..30}; do
        if curl -s http://localhost:3000 >/dev/null; then
            print_success "Open Web UI est accessible sur http://localhost:3000"
            print_message "Identifiants par défaut :"
            echo "  - Nom d'utilisateur : admin"
            echo "  - Mot de passe : admin"
            return 0
        fi
        echo -n "."
        sleep 1
    done
    
    print_error "Open Web UI n'a pas démarré dans le temps imparti"
    print_message "Vous pouvez vérifier les logs avec : docker logs open-webui"
}

# ------------------------------------------------------------
# Sélection du modèle
# ------------------------------------------------------------
choose_model() {
    # Imprimer le menu sur STDERR
    cat >&2 <<EOF

MODÈLES DISPONIBLES :
==================
1) 1.5B (~3GB)
2) 7B (~4.7GB) [Recommandé]
3) 8B (~5GB)
4) 14B (~9GB)
5) 32B (~19GB)
6) 70B (~40GB)
7) 671B (~350GB)

EOF

    # Demander le choix sur STDERR
    read -r -p "Votre choix (1-7, défaut: 2) : " choice >&2
    echo >&2 # Nouvelle ligne sur STDERR

    # Normaliser le choix et ne renvoyer QUE le modèle sur stdout
    case "$choice" in
        1)  echo "1.5b" ;;
        2|"") echo "7b" ;;  # défaut
        3)  echo "8b" ;;
        4)  echo "14b" ;;
        5)  echo "32b" ;;
        6)  echo "70b" ;;
        7)  echo "671b" ;;
        *)  echo "7b" ;;  # fallback
    esac
}

# ------------------------------------------------------------
# Téléchargement du modèle
# ------------------------------------------------------------
download_model() {
    local model_size="$1"
    print_message "Téléchargement du modèle DeepSeek-R1:$model_size..."
    
    # Debug: afficher la valeur exacte du modèle
    print_message "Debug: MODEL_SIZE='$model_size'" >&2
    
    ollama pull "deepseek-r1:$model_size" \
        || { print_error "Erreur lors du téléchargement du modèle DeepSeek-R1:$model_size"; exit 1; }
    
    print_success "Le modèle DeepSeek-R1:$model_size a été téléchargé avec succès"
}

# ------------------------------------------------------------
# Test du modèle
# ------------------------------------------------------------
test_model() {
    local model_size="$1"
    print_message "Test du modèle avec une requête simple..."
    echo "Ecris un poème de 2 lignes qui montre ton amour pour Korben de korben.info" | ollama run "deepseek-r1:$model_size"
}

# ------------------------------------------------------------
# Création de la documentation
# ------------------------------------------------------------
create_readme() {
    local model_size="$1"
    cat > README.md << EOF
# DeepSeek-R1 Installation

Ce projet a installé DeepSeek-R1:$model_size avec Ollama et Open Web UI.

## Utilisation

1. Interface graphique (Open Web UI):
   - Ouvrez http://localhost:3000 dans votre navigateur
   - Connectez-vous avec :
     * Nom d'utilisateur : admin
     * Mot de passe : admin
   - Sélectionnez le modèle deepseek-r1:$model_size

2. Ligne de commande:
   - Utilisez la commande: \`ollama run deepseek-r1:$model_size\`

## Intégration Python

\`\`\`python
import openai

client = openai.Client(
    base_url="http://localhost:11434/v1",
    api_key="ollama"
)

response = client.chat.completions.create(
    model="deepseek-r1:$model_size",
    messages=[{"role": "user", "content": "Votre question ici"}],
    temperature=0.7
)
\`\`\`

## Ressources utiles

- [Documentation Ollama](https://ollama.com/)
- [Documentation DeepSeek](https://api-docs.deepseek.com/)
- [Open Web UI](https://github.com/open-webui/open-webui)
EOF

    print_success "README.md créé avec les instructions d'utilisation"
}

# ------------------------------------------------------------
# Programme principal
# ------------------------------------------------------------
main() {
    print_message "Début de l'installation de DeepSeek-R1..."

    # Vérifier et installer les prérequis
    if ! check_homebrew; then
        install_homebrew
    fi
    
    if ! check_docker; then
        install_docker
    fi
    
    if ! check_docker_running; then
        print_message "Docker n'est pas en cours d'exécution"
        start_docker || { print_error "Impossible de démarrer Docker"; exit 1; }
    fi
    
    if ! check_ollama; then
        install_ollama
    fi
    
    if ! check_ollama_running; then
        print_message "Ollama n'est pas en cours d'exécution"
        start_ollama || { print_error "Impossible de démarrer Ollama"; exit 1; }
    fi

    # Sélectionner le modèle (une seule fois)
    print_message "Configuration du modèle..."
    MODEL_SIZE=$(choose_model)
    print_message "Debug: MODEL_SIZE sélectionné='$MODEL_SIZE'" >&2

    # Installer Open Web UI
    install_open_webui

    # Télécharger le modèle choisi avec la valeur capturée
    download_model "$MODEL_SIZE"

    # Créer la documentation
    test_model "$MODEL_SIZE"
    create_readme "$MODEL_SIZE"

    print_success "Installation terminée !"
    echo
    echo "Vous pouvez maintenant utiliser DeepSeek-R1 de deux façons :"
    echo "1. Interface graphique : http://localhost:3000"
    echo "2. Ligne de commande : ollama run deepseek-r1:$MODEL_SIZE"
    echo
    print_message "Consultez le fichier README.md pour plus d'informations sur l'intégration Python"
}

# Lancer le script
main