import gradio as gr
from transformers import pipeline
from pydub import AudioSegment
from pydub.utils import make_chunks
import tempfile
import os

# Inicialize o modelo Whisper com um modelo otimizado para CPU
transcriber = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-tiny",  # Use modelos menores para maior velocidade
    device="cpu"  # Certifique-se de usar CPU, já que está na versão gratuita
)

# Função para dividir áudios longos em trechos menores (30 segundos)
def split_audio(audio_path, chunk_length=30_000):
    audio = AudioSegment.from_file(audio_path)
    chunks = make_chunks(audio, chunk_length)  # Divide em trechos de 30 segundos
    return chunks

# Função para comprimir áudio (ajustar taxa de amostragem, etc.)
def compress_audio(audio_path):
    audio = AudioSegment.from_file(audio_path)
    compressed_audio = audio.set_frame_rate(16000).set_channels(1).set_sample_width(2)
    compressed_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    compressed_audio.export(compressed_path, format="wav")
    return compressed_path

# Função para transcrever o áudio
def transcribe(audio_file):
    try:
        # Verifique o tamanho do arquivo (máx. 60 MB)
        if os.path.getsize(audio_file) > 60 * 1024 * 1024:
            return "Erro: O arquivo excede o limite de 60 MB. Por favor, envie um áudio menor."

        # Comprimir o áudio antes de processar
        compressed_audio = compress_audio(audio_file)

        # Divida o áudio em partes
        chunks = split_audio(compressed_audio)
        full_transcription = []

        # Processar cada parte separadamente
        for i, chunk in enumerate(chunks):
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as temp_chunk:
                chunk.export(temp_chunk.name, format="wav")
                transcription = transcriber(temp_chunk.name, return_timestamps=False)["text"]
                full_transcription.append(f"[Parte {i+1}]: {transcription}")
        
        return "\n".join(full_transcription)
    except Exception as e:
        return f"Erro ao processar áudio: {str(e)}"

# Interface gráfica com Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🎙️ Whisper Transcription - Suporte a Áudios Longos (60 Minutos / 60 MB)")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1️⃣ Envie seu áudio (máx. 60 MB e 60 minutos)")
            audio_input = gr.Audio(type="filepath", label="Envie um arquivo de áudio")
        
        with gr.Column(scale=1):
            gr.Markdown("### 2️⃣ Resultado da transcrição")
            transcription_output = gr.Textbox(label="Transcrição", lines=15, interactive=False)
    
    transcribe_button = gr.Button("🚀 Transcrever")
    
    # Vincular ação ao botão
    transcribe_button.click(transcribe, inputs=[audio_input], outputs=[transcription_output])

# Rodar a aplicação
demo.launch(share=True)