# -*- coding: utf-8 -*-
import torch
import gradio as gr
import yt_dlp as youtube_dl
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
from huggingface_hub import login

import tempfile
import os
import time

# --- Конфигурация ---
HF_TOKEN = os.getenv("HF_ACCESS_TOKEN")
MODEL_NAME = "artyomboyko/whisper-small-ru-v4"
BATCH_SIZE = 8
# FILE_LIMIT_MB = 1000

# --- Инициализация ---
try:
    if HF_TOKEN:
        login(token=HF_TOKEN)
        print("Успешный вход в Hugging Face Hub.")
    else:
        print("HF_ACCESS_TOKEN не найден. Запуск в оффлайн-режиме или с публичными моделями.")
except Exception as e:
    print(f"Не удалось войти в Hugging Face Hub: {e}")

# Определение устройства для вычислений (GPU или CPU)
device = 0 if torch.cuda.is_available() else "cpu"
if device == 0:
    print("GPU доступен, вычисления будут производиться на GPU.")
else:
    print("GPU не доступен, вычисления будут производиться на CPU.")

# Загрузка модели распознавания речи
try:
    pipe = pipeline(
        task="automatic-speech-recognition",
        model=MODEL_NAME,
        chunk_length_s=30,
        device=device,
    )
    print(f"Модель {MODEL_NAME} успешно загружена.")
except Exception as e:
    print(f"Ошибка при загрузке модели {MODEL_NAME}: {e}")
    exit()


# --- Функции ---
def transcribe(filepath, task):
    """
    Транскрибирует аудиофайл с помощью загруженной модели Whisper.

    Args:
        filepath (str): Путь к аудиофайлу.
        task (str): Задача для модели ("transcribe" или "translate").

    Returns:
        str: Распознанный текст.
    """
    if filepath is None:
        # Используем gr.Warning вместо gr.Error для неблокирующего уведомления
        gr.Warning("Аудиофайл не предоставлен! Пожалуйста, загрузите или запишите аудио перед отправкой запроса.")
        return "Ошибка: Аудиофайл не предоставлен." # Возвращаем строку ошибки

    print(f"Начало транскрибации файла: {filepath}, Задача: {task}")
    try:
        result = pipe(filepath, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
        text = result["text"]
        print("Транскрибация завершена успешно.")
        return text
    except Exception as e:
        print(f"Ошибка во время транскрибации: {e}")
        # Используем gr.Error для вывода серьезной ошибки
        raise gr.Error(f"Произошла ошибка при обработке аудио: {e}")


# --- Создание интерфейса Gradio ---
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="neutral")) as demo:
    gr.Markdown(
        f"""
        # Whisper: Транскрибация аудио
        Транскрибируйте длинные аудиозаписи (с микрофона или из файла) одним нажатием кнопки!
        Демо использует модель OpenAI Whisper [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME})
        через библиотеку 🤗 Transformers для распознавания речи.
        """
    )

    # Создание вкладок
    with gr.Tabs():
        
        # Вкладка 1: Транскрибация с микрофона
        with gr.TabItem("Микрофон"):
            with gr.Row():
                mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Запись с микрофона")
                task_mic = gr.Radio(["transcribe", "translate"], label="Задача", value="transcribe")
            output_mic = gr.Textbox(label="Результат", lines=7) # Увеличили количество строк для текста
            mic_button = gr.Button("Транскрибировать с микрофона")

        # Вкладка 2: Транскрибация из файла
        with gr.TabItem("Аудиофайл"):
            with gr.Row():
                file_input = gr.Audio(sources=["upload"], type="filepath", label="Загрузить аудиофайл")
                task_file = gr.Radio(["transcribe", "translate"], label="Задача", value="transcribe")
            output_file = gr.Textbox(label="Результат", lines=7)
            file_button = gr.Button("Транскрибировать файл")


    # --- Привязка функций к кнопкам ---
    
    # Связываем кнопку "Транскрибировать с микрофона" с функцией transcrib
    mic_button.click(
        fn=transcribe,                     
        inputs=[mic_input, task_mic],      
        outputs=output_mic                 
    )

    # Связываем кнопку "Транскрибировать файл" с функцией transcribe
    file_button.click(
        fn=transcribe,
        inputs=[file_input, task_file],
        outputs=output_file
    )

# --- Запуск приложения ---
print("Запуск интерфейса Gradio...")
demo.queue().launch(debug=True, share=False)