#!/usr/bin/env bash
set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
MODELS_DIR="$PROJECT_DIR/models"
SAMPLES_DIR="$SCRIPT_DIR/samples"
DATA_DIR="$SCRIPT_DIR/docker_data"

echo "============================================"
echo "  Cosmo Wake Word — обучение модели"
echo "============================================"
echo ""

# Проверяем Docker
if ! command -v docker &>/dev/null; then
    echo "ОШИБКА: Docker не найден."
    echo "Установи Docker Desktop: https://www.docker.com/products/docker-desktop/"
    exit 1
fi

if ! docker info &>/dev/null; then
    echo "ОШИБКА: Docker не запущен. Запусти Docker Desktop и попробуй снова."
    exit 1
fi

# Проверяем что есть записанные примеры
POSITIVE_DIR="$SAMPLES_DIR/positive"
if [ ! -d "$POSITIVE_DIR" ] || [ -z "$(ls "$POSITIVE_DIR"/*.wav 2>/dev/null)" ]; then
    echo "Записанные примеры не найдены."
    echo "Сначала запусти запись голоса:"
    echo "  python train_wakeword/record_samples.py"
    echo ""
    read -p "Продолжить без записанных примеров? (используются только TTS) [y/N]: " yn
    case "$yn" in
        [Yy]*) echo "Продолжаем с только TTS примерами..." ;;
        *) exit 0 ;;
    esac
fi

mkdir -p "$MODELS_DIR" "$DATA_DIR"

# Собираем Docker образ
echo "[1/4] Собираю Docker образ (первый раз ~5-10 мин)..."
docker build -t cosmo-wakeword-trainer "$SCRIPT_DIR" --quiet
echo "      Образ готов."
echo ""

# Скачиваем датасет негативных примеров если нет
NEGATIVE_FEATURES="$DATA_DIR/openwakeword_features_ACAV100M_2000_hrs_16bit.npy"
VALIDATION_FEATURES="$DATA_DIR/validation_set_features.npy"

if [ ! -f "$NEGATIVE_FEATURES" ]; then
    echo "[2/4] Скачиваю негативный датасет (~17 GB + ~500 MB, один раз)..."
    echo "      Это займёт время в зависимости от скорости интернета."
    echo ""
    echo "      Скачиваю ACAV100M features (~17 GB)..."
    wget -q --show-progress \
        -O "$NEGATIVE_FEATURES" \
        "https://huggingface.co/datasets/davidscripka/openwakeword_features/resolve/main/openwakeword_features_ACAV100M_2000_hrs_16bit.npy"
    echo "      Скачиваю validation features (~500 MB)..."
    wget -q --show-progress \
        -O "$VALIDATION_FEATURES" \
        "https://huggingface.co/datasets/davidscripka/openwakeword_features/resolve/main/validation_set_features.npy"
    echo "      Датасет готов."
else
    echo "[2/4] Негативный датасет уже скачан. Пропускаю."
fi
echo ""

# Запускаем обучение
echo "[3/4] Запускаю обучение в Docker..."
echo "      Это займёт ~30-60 минут."
echo ""

SAMPLES_MOUNT=""
if [ -d "$POSITIVE_DIR" ] && [ -n "$(ls "$POSITIVE_DIR"/*.wav 2>/dev/null)" ]; then
    SAMPLES_MOUNT="-v $POSITIVE_DIR:/samples/positive"
fi

docker run --rm \
    --shm-size=2g \
    -v "$SCRIPT_DIR/cosmo_config.yaml:/app/cosmo_config.yaml" \
    -v "$DATA_DIR:/data" \
    -v "$MODELS_DIR:/output" \
    $SAMPLES_MOUNT \
    cosmo-wakeword-trainer

echo ""
echo "[4/4] Копирую модель в проект..."

# Ищем готовую модель
ONNX_FILE=$(ls "$MODELS_DIR"/*.onnx 2>/dev/null | head -1)

if [ -n "$ONNX_FILE" ]; then
    echo ""
    echo "============================================"
    echo "  Готово! Модель сохранена:"
    echo "  $ONNX_FILE"
    echo ""
    echo "  Обновляю wake_word детектор..."
    # Обновляем путь в конфиге
    MODEL_FILENAME=$(basename "$ONNX_FILE")
    sed -i "s|wakeword_models=\[\"hey_jarvis\"\]|wakeword_models=[\"models/$MODEL_FILENAME\"]|g" \
        "$PROJECT_DIR/cosmo/wake_word.py" 2>/dev/null || true
    echo "  Теперь запускай: bash run.sh"
    echo "  и говори 'Hey Cosmo' для активации!"
    echo "============================================"
else
    echo "ОШИБКА: .onnx файл не найден в $MODELS_DIR"
    echo "Проверь логи Docker выше."
    exit 1
fi