import torch
import gradio as gr
from transformers import pipeline
import os
import pandas as pd  # Importa o Pandas
import json  # Importa o módulo json para ler o arquivo de palavras-chave


MODEL_NAME = "openai/whisper-large-v2"
BATCH_SIZE = 8

device = 0 if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)

# Função para carregar palavras-chave do arquivo JSON
def load_keywords(json_path):
    with open(json_path, 'r') as file:
        data = json.load(file)
    keywords = data.get('keywords', []) + data.get('key_words_Telus', [])
    return keywords

# Caminho para o arquivo JSON de palavras-chave
keywords_path = os.path.join(os.path.dirname(__file__), "keywords.json")
keywords = load_keywords(keywords_path)

def transcribe(inputs, task):
    if inputs is None:
        raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
    
    result = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
    text = result["text"]
    
    # Identifica as palavras-chave na transcrição
    found_keywords = [keyword for keyword in keywords if keyword.lower() in text.lower()]
    
    # Cria um DataFrame para armazenar os dados
    df = pd.DataFrame({'Transcription': [text], 'Found Keywords': [", ".join(found_keywords)]})
    
    # Define o caminho onde o arquivo CSV será salvo
    csv_save_path = os.path.join(os.path.dirname(__file__), "transcription.csv")
    # Salva o DataFrame como um arquivo CSV
    df.to_csv(csv_save_path, index=False)
    
    # Retorna o texto da transcrição, as palavras-chave encontradas e o caminho do arquivo CSV
    return text, ", ".join(found_keywords), csv_save_path

# Ajusta as interfaces Gradio para incluir a saída das palavras-chave encontradas
mf_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="filepath"),
        gr.Radio(choices=["transcribe", "translate"], label="Task"),
    ],
    outputs=[
        gr.Textbox(label="Transcription Text"),
        gr.Textbox(label="Found Keywords", visible=False),  # Oculta para a aba de microfone
        gr.File(label="Download CSV File", visible=False),  # Oculta o download do CSV para a aba de microfone
    ],
    title="GexelView.AI - Transcriber",
    description="Transcribe or translate audio inputs with a single click.",
    allow_flagging="never",
)

file_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="filepath", label="Audio file"),
        gr.Radio(choices=["transcribe", "translate"], label="Task"),
    ],
    outputs=[
        gr.Textbox(label="Transcription Text"),
        gr.Textbox(label="Found Keywords"),  # Exibe as palavras-chave encontradas
        gr.File(label="Download CSV File"),  # Torna o download do arquivo CSV visível para a aba de arquivo
    ],
    title="GexelView.AI - Transcriber",
    description="Transcribe or translate audio files with a single click.",
    allow_flagging="never",
)

demo = gr.Blocks()

with demo:
    gr.TabbedInterface([mf_transcribe, file_transcribe], ["Microphone", "Audio file"])

demo.launch()