import gradio as gr
from transformers import pipeline
import torch

# Проверка доступности GPU
device = 0 if torch.cuda.is_available() else -1

print("Loading models...")

# 1. Speech-to-Text (Whisper)
transcriber = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-small",
    device=device
)

# 2. Summarization
summarizer = pipeline(
    "summarization",
    model="Falconsai/text_summarization",
    device=device
)

# 3. Text Classification (Zero-shot для определения темы)
classifier = pipeline(
    "zero-shot-classification",
    model="facebook/bart-large-mnli",
    device=device
)

print("Models loaded successfully!")

# Категории тем для классификации
CATEGORIES = [
    "Complaint / Problem",
    "Question / Information Request",
    "Gratitude / Thanks",
    "Order / Purchase",
    "Technical Support",
    "Delivery / Logistics",
    "Finance / Payment",
    "General Conversation",
    "Weather / Climate Information",
]


def analyze_audio(audio_file):
    """
    Анализирует аудио через три этапа:
    1. Транскрипция (Speech-to-Text)
    2. Краткая выжимка (Summarization)
    3. Определение темы (Classification)
    """
    if audio_file is None:
        return "Please upload an audio file", "", ""
    
    try:
        # 1. Speech-to-Text
        print("Transcribing audio...")
        transcription_result = transcriber(audio_file)
        transcription = transcription_result['text']
        
        if not transcription or len(transcription.strip()) < 10:
            return transcription, "Текст слишком короткий для суммаризации", "Не удалось определить тему"
        
        # 2. Summarization
        print("Summarizing text...")
        # Ограничиваем длину текста для суммаризации
        max_input_length = 1024
        text_for_summary = transcription[:max_input_length] if len(transcription) > max_input_length else transcription
        
        summary_result = summarizer(
            text_for_summary,
            max_length=100,
            min_length=20,
            do_sample=False
        )
        summary = summary_result[0]['summary_text']
        
        # 3. Text Classification (используем оригинальную транскрипцию)
        print("Classifying topic...")
        classification_result = classifier(
            transcription,
            candidate_labels=CATEGORIES,
            multi_label=False
        )
        
        # Форматируем результаты классификации
        top_category = classification_result['labels'][0]
        top_score = classification_result['scores'][0]
        
        # Показываем топ-3 категории
        category_text = f"**Основная тема:** {top_category} ({top_score:.1%})\n\n"
        category_text += "**Другие возможные темы:**\n"
        for label, score in zip(classification_result['labels'][1:3], 
                               classification_result['scores'][1:3]):
            category_text += f"- {label}: {score:.1%}\n"
        
        return transcription, summary, category_text
        
    except Exception as e:
        error_msg = f"Error during analysis: {str(e)}"
        print(error_msg)
        return error_msg, "", ""


# Создание интерфейса Gradio
with gr.Blocks(title="Voice Message Analyzer", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown(
        """
        # 🎙️ Voice Message Analyzer
        
        Анализатор голосовых сообщений для менеджеров и службы поддержки.
        
        **Возможности:**
        1. **Транскрипция** - преобразование аудио в текст
        2. **Краткая выжимка** - суть сообщения в паре предложений
        3. **Определение темы** - автоматическая категоризация обращения
        
        ---
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                type="filepath",
                label="Загрузите аудио файл",
                sources=["upload", "microphone"]
            )
            analyze_btn = gr.Button("🔍 Анализировать", variant="primary", size="lg")
            
            gr.Markdown(
                """
                ### 💡 Совет:
                - Используйте аудио на английском языке для лучших результатов
                - Длительность: 10 секунд - 5 минут
                - Форматы: WAV, MP3, M4A
                """
            )
        
        with gr.Column(scale=1):
            transcription_output = gr.Textbox(
                label="1️⃣ Транскрипция",
                placeholder="Полный текст голосового сообщения появится здесь...",
                lines=5
            )
            summary_output = gr.Textbox(
                label="2️⃣ Краткая выжимка",
                placeholder="Краткое содержание появится здесь...",
                lines=3
            )
            category_output = gr.Markdown(
                value="_Категория обращения появится здесь..._"
            )
    
    gr.Markdown(
        """
        ---
        ### 🤖 Используемые модели:
        - **Speech-to-Text**: OpenAI Whisper-small (3.4M загрузок)
        - **Summarization**: Falconsai/text_summarization
        - **Classification**: Facebook BART-large-mnli (Zero-shot)
        
        ### 🎯 Применение:
        - Обработка голосовых обращений клиентов
        - Быстрый анализ голосовых сообщений
        - Автоматическая категоризация запросов
        - Создание конспектов встреч и звонков
        """
    )
    
    # Обработчик события
    analyze_btn.click(
        fn=analyze_audio,
        inputs=audio_input,
        outputs=[transcription_output, summary_output, category_output]
    )

# Запуск приложения
if __name__ == "__main__":
    demo.launch()