#!/bin/bash

# Script de démarrage pour Private-GPT avec Ollama personnalisé
# Permet de lancer facilement private-gpt avec des modèles d'IA personnalisés

set -e

# Répertoire courant où se trouve le script
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Couleurs pour les messages
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Afficher les modèles disponibles
list_models() {
    echo -e "${BLUE}Modèles Ollama disponibles:${NC}"
    docker exec -it pgpt-ollama-cpu-1 ollama list
}

# Télécharger un nouveau modèle
download_model() {
    if [ -z "$1" ]; then
        echo -e "${RED}Erreur: Veuillez spécifier un nom de modèle${NC}"
        echo "Exemple: $0 download llama3"
        exit 1
    fi

    echo -e "${YELLOW}Téléchargement du modèle $1...${NC}"
    docker exec -it pgpt-ollama-cpu-1 ollama pull "$1"
    echo -e "${GREEN}Modèle $1 téléchargé avec succès${NC}"
}

# Arrêter les services
stop_services() {
    echo -e "${YELLOW}Arrêt des services Private-GPT...${NC}"
    docker compose -f docker-compose.yaml down
    echo -e "${GREEN}Services arrêtés${NC}"
}

# Démarrer les services
start_services() {
    local profile="$1"

    if [ -z "$profile" ]; then
        profile="ollama-cpu"
    fi

    echo -e "${YELLOW}Démarrage des services Private-GPT avec le profil $profile...${NC}"
    docker compose -f docker-compose.yaml --profile "$profile" up -d
    echo -e "${GREEN}Services démarrés avec succès${NC}"
    echo -e "${BLUE}Interface web disponible sur: ${GREEN}http://localhost:8001${NC}"
}

# Afficher l'aide
show_help() {
    echo -e "${BLUE}Script de démarrage pour Private-GPT avec Ollama${NC}"
    echo ""
    echo -e "Usage: $0 [commande]"
    echo ""
    echo "Commandes:"
    echo "  start [profile]     Démarrer les services (profile: ollama-cpu (défaut), ollama-cuda)"
    echo "  stop                Arrêter les services"
    echo "  list                Lister les modèles Ollama disponibles"
    echo "  download <modèle>   Télécharger un nouveau modèle Ollama"
    echo "  help                Afficher cette aide"
    echo ""
    echo "Exemples:"
    echo "  $0 start            Démarrer avec CPU (par défaut)"
    echo "  $0 start ollama-cuda Démarrer avec GPU (CUDA)"
    echo "  $0 download mistral  Télécharger le modèle Mistral"
    echo ""
}

# Traitement des commandes
case "$1" in
    start)
        start_services "$2"
        ;;
    stop)
        stop_services
        ;;
    list)
        list_models
        ;;
    download)
        download_model "$2"
        ;;
    *)
        show_help
        ;;
esac