"""
Запись голосовых примеров для обучения wake word модели.
Запускай: python train_wakeword/record_samples.py

Скрипт записывает N примеров слова "Hey Cosmo" с паузами между ними.
Файлы сохраняются в train_wakeword/samples/positive/
"""

import os
import sys
import time
import wave
import struct
import threading

try:
    import pyaudio
except ImportError:
    print("Установи pyaudio: pip install pyaudio")
    sys.exit(1)

# --- Настройки ---
WAKE_WORD       = "Hey Cosmo"
N_SAMPLES       = 30          # сколько примеров записать
RECORD_SECS     = 2.0         # длина одной записи (сек)
PAUSE_SECS      = 2.0         # пауза между записями (сек)
SAMPLE_RATE     = 16000
CHANNELS        = 1
CHUNK           = 512
OUTPUT_DIR      = os.path.join(os.path.dirname(__file__), "samples", "positive")

os.makedirs(OUTPUT_DIR, exist_ok=True)

def record_clip(pa: pyaudio.PyAudio, filename: str, duration: float):
    stream = pa.open(
        format=pyaudio.paInt16,
        channels=CHANNELS,
        rate=SAMPLE_RATE,
        input=True,
        frames_per_buffer=CHUNK,
    )
    frames = []
    n_chunks = int(SAMPLE_RATE / CHUNK * duration)
    for _ in range(n_chunks):
        frames.append(stream.read(CHUNK, exception_on_overflow=False))
    stream.stop_stream()
    stream.close()

    with wave.open(filename, "wb") as wf:
        wf.setnchannels(CHANNELS)
        wf.setsampwidth(pa.get_sample_size(pyaudio.paInt16))
        wf.setframerate(SAMPLE_RATE)
        wf.writeframes(b"".join(frames))

def countdown(seconds: int):
    for i in range(seconds, 0, -1):
        print(f"\r  {i}...", end="", flush=True)
        time.sleep(1)
    print("\r  Говори! ", end="", flush=True)

def main():
    # Считаем уже записанные файлы
    existing = [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".wav")]
    start_idx = len(existing)

    if start_idx >= N_SAMPLES:
        print(f"Уже записано {start_idx} примеров. Для перезаписи удали папку {OUTPUT_DIR}")
        return

    pa = pyaudio.PyAudio()

    print("=" * 50)
    print(f"  Запись примеров wake word: \"{WAKE_WORD}\"")
    print(f"  Нужно записать: {N_SAMPLES} примеров")
    print(f"  Уже есть: {start_idx}")
    print(f"  Длина каждой записи: {RECORD_SECS} сек")
    print("=" * 50)
    print()
    print("Инструкция:")
    print("  - Говори чётко и естественно")

    print("  - Меняй интонацию, темп, громкость")
    print("  - Можно говорить чуть тише / громче / быстрее")
    print("  - Представь что реально обращаешься к ассистенту")
    print()
    input("  Нажми Enter когда готов начать...")
    print()

    for i in range(start_idx, N_SAMPLES):
        num = i + 1
        filename = os.path.join(OUTPUT_DIR, f"hey_cosmo_{num:03d}.wav")

        print(f"[{num:2d}/{N_SAMPLES}] Приготовься... ", end="", flush=True)
        countdown(2)

        record_clip(pa, filename, RECORD_SECS)
        print(f" ✓ записано")

        if num < N_SAMPLES:
            time.sleep(PAUSE_SECS)

    pa.terminate()

    print()
    print("=" * 50)
    print(f"  Готово! Записано {N_SAMPLES} примеров.")
    print(f"  Папка: {OUTPUT_DIR}")
    print()
    print("  Следующий шаг:")
    print("    bash train_wakeword/train.sh")
    print("=" * 50)

if __name__ == "__main__":
    main()