# vosk_stt.py
import json
import yaml
import logging
from datetime import datetime
from vosk import Model, KaldiRecognizer

class VoskSTT:
    def __init__(self, model_path="/home/sirmrak/models/vosk-model-small-ru-0.22", commands_file="/mnt/Project/my_speaker/commands/commands.yaml", sample_rate=16000):
        """
        Инициализирует Vosk STT с моделью и словарем команд из YAML файла.
        
        :param model_path: Путь к модели Vosk (vosk-model-small-ru-0.22)
        :param commands_file: Путь к файлу commands.yaml
        :param sample_rate: Частота дискретизации аудио (по умолчанию 16000)
        """
        self.model_path = model_path
        self.commands_file = commands_file
        self.sample_rate = sample_rate
        self.model = None
        self.recognizer = None
        self._load_model()
        self._load_commands()

    def _load_model(self):
        """Загружает модель Vosk."""
        logger = logging.getLogger('VADApp')
        try:
            self.model = Model(self.model_path)
            logger.info(f"Модель Vosk загружена: {self.model_path}")
        except Exception as e:
            logger.error(f"Ошибка загрузки модели: {e}")
            raise

    def _load_commands(self):
        """Формирует словарь (список фраз) из commands.yaml и инициализирует recognizer."""
        logger = logging.getLogger('VADApp')
        try:
            with open(self.commands_file, 'r', encoding='utf-8') as f:
                commands = yaml.safe_load(f)
            
            phrases = []
            
            # Собираем все фразы из всех разделов
            for section, content in commands.items():
                if section in ['Wakewords', 'Session']:
                    # Обрабатываем списки с вложенными phrases
                    for item in content:
                        if isinstance(item, dict) and 'phrases' in item:
                            phrases.extend(item['phrases'])
                elif section == 'Volume':
                    # Обрабатываем новую структуру Volume с приоритетом
                    for item in content:
                        if isinstance(item, dict) and 'commands' in item:
                            phrases.extend(list(item['commands'].keys()))
                elif section == 'Text':
                    # Обрабатываем новую структуру Text с приоритетом
                    for item in content:
                        if isinstance(item, dict) and 'commands' in item:
                            phrases.extend(list(item['commands'].keys()))
            
            # Удаляем дубликаты и сортируем
            phrases = sorted(set(phrases))
            
            # Инициализируем recognizer с грамматикой (списком фраз)
            grammar = json.dumps(phrases, ensure_ascii=False)
            self.recognizer = KaldiRecognizer(self.model, self.sample_rate, grammar)
            
            logger.info(f"Словарь команд загружен: {len(phrases)} фраз")
            logger.debug(f"Фразы: {phrases}")
        except Exception as e:
            logger.error(f"Ошибка загрузки команд: {e}")
            raise

    def recognize(self, audio_data):
        """
        Распознает речь из аудио данных (bytes).
        
        :param audio_data: Bytes аудио (16-bit PCM, mono)
        :return: Распознанный текст или None если ничего не распознано
        """
        logger = logging.getLogger('VADApp')
        try:
            if not self.recognizer.AcceptWaveform(audio_data):
                # Частичный результат (если нужно, но для коротких команд используем Final)
                partial = json.loads(self.recognizer.PartialResult()).get("partial", "")
                if partial:
                    logger.debug(f"Частичный результат: {partial}")
            result = json.loads(self.recognizer.FinalResult()).get("text", "")
            if result:
                logger.info(f"Распознанная команда: {result}")
                return result
            else:
                logger.info("Ничего не распознано")
                return None
        except Exception as e:
            logger.error(f"Ошибка распознавания: {e}")
            return None

    def reload_commands(self):
        """Перезагружает словарь команд из файла."""
        self._load_commands()