import os
import gradio as gr
import numpy as np
import sherpa_onnx
from huggingface_hub import hf_hub_download

# ── 1. Загрузка приватных моделей ────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN")
REPO_ID = "exboogiman/streamingASR"  

print("Скачивание моделей...")
encoder_path = hf_hub_download(repo_id=REPO_ID, filename="encoder-uzbek_fc.onnx", token=HF_TOKEN)
decoder_path = hf_hub_download(repo_id=REPO_ID, filename="decoder-uzbek_fc.onnx", token=HF_TOKEN)
joiner_path  = hf_hub_download(repo_id=REPO_ID, filename="joiner-uzbek_fc.onnx", token=HF_TOKEN)
tokens_path  = hf_hub_download(repo_id=REPO_ID, filename="tokens.txt", token=HF_TOKEN)
vad_path     = hf_hub_download(repo_id=REPO_ID, filename="silero_vad.onnx", token=HF_TOKEN)

# ── 2. Инициализация компонентов ────────────────────────────
recognizer = sherpa_onnx.OfflineRecognizer.from_transducer(
    encoder=encoder_path,
    decoder=decoder_path,
    joiner=joiner_path,
    tokens=tokens_path,
    model_type="nemo_transducer",
    num_threads=2,
    sample_rate=16000,
    feature_dim=80,
    decoding_method="greedy_search",
)

# Настройки VAD
vad_config = sherpa_onnx.VadModelConfig()
vad_config.silero_vad.model = vad_path
vad_config.silero_vad.min_silence_duration = 0.3  # пауза 300мс = конец фразы
vad_config.silero_vad.threshold = 0.5
vad_config.sample_rate = 16000

# ── 3. Логика стриминга ──────────────────────────────────────
def initialize_state():
    return {
        "vad": sherpa_onnx.VoiceActivityDetector(vad_config, buffer_size_in_seconds=30),
        "buffer": np.array([], dtype=np.float32),
        "transcript": "",
    }

def transcribe_stream(audio, state):
    if state is None:
        state = initialize_state()
        
    if audio is None:
        return state["transcript"], state
        
    sr, y = audio
    
    # Приводим к моно
    if y.ndim > 1:
        y = y.mean(axis=-1)
        
    # Приводим к float32 и нормализуем
    samples = y.astype(np.float32)
    if np.max(np.abs(samples)) > 1.0:
        samples = samples / 32768.0
        
    # Добавляем новые семплы в буфер
    state["buffer"] = np.concatenate([state["buffer"], samples])
    
    vad = state["vad"]
    window_size = vad_config.silero_vad.window_size
    
    # Передаем в VAD фиксированными окнами
    while len(state["buffer"]) >= window_size:
        vad.accept_waveform(state["buffer"][:window_size])
        state["buffer"] = state["buffer"][window_size:]
        
    # Если VAD вырезал законченную фразу — распознаем её офлайн-моделью
    new_text_segments = []
    while not vad.empty():
        stream = recognizer.create_stream()
        stream.accept_waveform(16000, vad.front.samples)
        vad.pop()
        recognizer.decode_stream(stream)
        text = stream.result.text.strip()
        if text:
            new_text_segments.append(text)
            
    # Если есть новые распознанные слова, добавляем их к общей истории
    if new_text_segments:
        added_text = " ".join(new_text_segments)
        if state["transcript"]:
            state["transcript"] += " " + added_text
        else:
            state["transcript"] = added_text
            
    return state["transcript"], state

# ── 4. Описание интерфейса Gradio Blocks ─────────────────────
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# STT Real-time Streaming Demo (Uzbek)")
    
    # Состояние для сохранения VAD буфера между вызовами
    state = gr.State()
    
    with gr.Row():
        # Вход звука с микрофона с флагом streaming=True
        input_audio = gr.Audio(
            sources=["microphone"],
            type="numpy",
            streaming=True,
            label="Говорите в микрофон"
        )
        # Поле вывода текста
        output_text = gr.Textbox(
            label="Распознанный текст",
            interactive=False,
            placeholder="Ваша речь отобразится здесь..."
        )
        
    # Событие .stream срабатывает каждые ~0.5 сек по мере записи звука
    input_audio.stream(
        fn=transcribe_stream,
        inputs=[input_audio, state],
        outputs=[output_text, state],
        show_progress="hidden"
    )
    
    # Сброс состояния при очистке аудио
    input_audio.clear(
        fn=initialize_state,
        inputs=[],
        outputs=[state]
    )

if __name__ == "__main__":
    demo.launch()