#!/bin/bash
# =============================================================================
# descargar-modelos.sh — Descarga de modelos según nivel de hardware
# =============================================================================
#
# Uso:
#   chmod +x descargar-modelos.sh
#   ./descargar-modelos.sh          # Modo interactivo (pregunta el nivel)
#   ./descargar-modelos.sh --nivel1 # Descarga solo modelos nivel 1
#   ./descargar-modelos.sh --nivel2 # Descarga modelos nivel 2 (recomendado)
#   ./descargar-modelos.sh --nivel3 # Descarga modelos nivel 3
#   ./descargar-modelos.sh --todos  # Descarga todos los modelos
#
# Descripción:
#   Descarga los modelos recomendados según tu hardware. Cada nivel corresponde
#   a los rangos de RAM mencionados en el episodio.
#
# Requisitos:
#   - Ollama instalado (ver 02-instalacion/setup-ollama.sh)
#   - Conexión a internet
#   - Espacio en disco suficiente (ver tabla abajo)
#
# Tamaños aproximados (cuantizados):
#   Nivel 1: phi3.5 (~2.5 GB) + qwen2.5-coder:1.5b (~1 GB) = ~3.5 GB
#   Nivel 2: llama3.2:3b (~2 GB) + qwen2.5:7b (~4.5 GB) + deepseek-coder (~9 GB)
#   Nivel 3: llama3.1:8b (~5 GB) + mistral:7b (~4.5 GB)
#   Embeddings: nomic-embed-text (~0.3 GB)
#
# =============================================================================

set -euo pipefail

# Colores para output
VERDE='\033[0;32m'
AMARILLO='\033[1;33m'
AZUL='\033[0;34m'
RESET='\033[0m'

echo -e "${AZUL}==============================================${RESET}"
echo -e "${AZUL}  Descarga de modelos — Laboratorio IA Local${RESET}"
echo -e "${AZUL}==============================================${RESET}"
echo ""

# --- Comprobar Ollama ----------------------------------------------------
if ! command -v ollama &>/dev/null; then
    echo "ERROR: Ollama no está instalado."
    echo "  Ejecuta primero: curl -fsSL https://ollama.com/install.sh | sh"
    exit 1
fi

echo -e "  ${VERDE}✓${RESET} Ollama detectado: $(ollama --version 2>/dev/null || echo 'versión desconocida')"
echo ""

# --- Selección de nivel --------------------------------------------------
NIVEL="${1:-}"

if [[ -z "$NIVEL" ]]; then
    echo "Selecciona tu nivel de hardware:"
    echo ""
    echo "  Nivel 1 — Hardware humilde (8 GB RAM / Raspberry Pi 5)"
    echo "    Modelos: phi3.5, qwen2.5-coder:1.5b"
    echo "    Espacio: ~3.5 GB"
    echo ""
    echo "  Nivel 2 — Punto dulce (16 GB RAM) ⭐ RECOMENDADO"
    echo "    Modelos: llama3.2:3b, qwen2.5:7b, deepseek-coder-v2:16b-lite-instruct-q4_K_M"
    echo "    Espacio: ~15 GB"
    echo ""
    echo "  Nivel 3 — Hardware potente (32 GB+ RAM / GPU)"
    echo "    Modelos: llama3.1:8b, mistral:7b"
    echo "    Espacio: ~10 GB"
    echo ""
    echo "  Todos — Descarga todos los modelos"
    echo "    Espacio: ~28 GB"
    echo ""
    echo "  Embeddings — Solo modelo de embeddings (para RAG)"
    echo "    Modelo: nomic-embed-text (~0.3 GB)"
    echo ""
    read -r -p "  ¿Qué nivel quieres descargar? [1/2/3/todos/emb]: " NIVEL
    echo ""
fi

# --- Mapear nivel a modelos ---------------------------------------------
case "$NIVEL" in
    --nivel1|1)
        echo -e "${AZUL}Nivel 1: Hardware humilde${RESET}"
        MODELOS=(
            "phi3.5"
            "qwen2.5-coder:1.5b"
        )
        ;;
    --nivel2|2)
        echo -e "${AZUL}Nivel 2: Punto dulce (RECOMENDADO)${RESET}"
        MODELOS=(
            "llama3.2:3b"
            "qwen2.5:7b"
            "deepseek-coder-v2:16b-lite-instruct-q4_K_M"
        )
        ;;
    --nivel3|3)
        echo -e "${AZUL}Nivel 3: Hardware potente${RESET}"
        MODELOS=(
            "llama3.1:8b"
            "mistral:7b"
        )
        ;;
    --todos|todos)
        echo -e "${AZUL}Todos los modelos${RESET}"
        MODELOS=(
            "phi3.5"
            "qwen2.5-coder:1.5b"
            "llama3.2:3b"
            "qwen2.5:7b"
            "deepseek-coder-v2:16b-lite-instruct-q4_K_M"
            "llama3.1:8b"
            "mistral:7b"
        )
        ;;
    --emb|emb)
        echo -e "${AZUL}Solo modelo de embeddings${RESET}"
        MODELOS=(
            "nomic-embed-text"
        )
        ;;
    *)
        echo "ERROR: Nivel no reconocido: $NIVEL"
        echo "  Usa: 1, 2, 3, todos, o emb"
        exit 1
        ;;
esac

echo ""
echo "  Modelos a descargar:"
for modelo in "${MODELOS[@]}"; do
    echo "    - $modelo"
done
echo ""

# --- Confirmar -----------------------------------------------------------
read -r -p "  ¿Continuar con la descarga? [s/N]: " respuesta
if [[ ! "$respuesta" =~ ^[sS]$ ]]; then
    echo "  Descarga cancelada."
    exit 0
fi

echo ""

# --- Descargar modelos ---------------------------------------------------
for modelo in "${MODELOS[@]}"; do
    echo -e "${AMARILLO}Descargando:${RESET} ollama pull $modelo"
    echo "  (Esto puede tardar varios minutos según tu conexión)"
    echo ""
    
    # Ejecutar ollama pull
    ollama pull "$modelo"
    
    echo -e "${VERDE}  ✓${RESET} $modelo descargado"
    echo ""
done

# --- Descargar modelo de embeddings (si no está ya incluido) -------------
if [[ "$NIVEL" != "emb" && "$NIVEL" != "--emb" ]]; then
    echo -e "${AMARILLO}¿Descargar también modelo de embeddings?${RESET}"
    echo "  nomic-embed-text (~0.3 GB) — necesario para RAG"
    read -r -p "  [s/N]: " respuesta
    if [[ "$respuesta" =~ ^[sS]$ ]]; then
        echo ""
        echo -e "${AMARILLO}Descargando:${RESET} ollama pull nomic-embed-text"
        ollama pull nomic-embed-text
        echo -e "${VERDE}  ✓${RESET} nomic-embed-text descargado"
    fi
fi

# --- Resumen final -------------------------------------------------------
echo ""
echo -e "${VERDE}==============================================${RESET}"
echo -e "${VERDE}  Descarga completada${RESET}"
echo -e "${VERDE}==============================================${RESET}"
echo ""
echo "  Modelos disponibles localmente:"
echo ""
ollama list
echo ""
echo "  Para probar un modelo:"
echo "    ollama run qwen2.5:7b \"Explícame qué es Linux\""
echo ""
echo "  Para eliminar un modelo:"
echo "    ollama rm qwen2.5:7b"
echo ""