import whisper
from pytube import YouTube
from transformers import pipeline
import gradio as gr
import os
import re

# model = whisper.load_model("base")
# model = pipeline(model="AlexMo/FIFA_WC22_WINNER_LANGUAGE_MODEL")
model = pipeline(model="AlexMo/improved_whisper_model")
summarizer = pipeline("summarization")

def transcribe_inp(microphone, file_upload):
    warn_output = ""
    if (microphone is not None) and (file_upload is not None):
        warn_output = (
            "NOTE: The audio file will be discarded after this run.\n"
        )

    elif (microphone is None) and (file_upload is None):
        return "ERROR: You have to either use the microphone or upload an audio file"

    file = microphone if microphone is not None else file_upload

    text = model(file, batch_size=1024)["text"]
    
    return warn_output + text

def getAudio(url):
    link = YouTube(url)
    video = link.streams.filter(only_audio=True).first()
    file = video.download(output_path=".")
    base, ext = os.path.splitext(file)
    file_ext = base + '.mp3'
    os.rename(file, file_ext)
    return file_ext


def getText(url):
    if url != '':
        output_text_transcribe = ''
    res = model(getAudio(url))
    return res['text'].strip()


def getSummary(article):
    # header = ' '.join(re.split(r'(?<=[.:;])\s', article)[:5])
    b = summarizer(article, min_length=15, max_length=120, do_sample=False)
    b = b[0]['summary_text'].replace(' .', '.').strip()

    return b


with gr.Blocks() as demo:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 500px; margin: 0 auto;">
              <div>
                <h1>Dutch whisperer</h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%">
                Summarize audio files, mic input or Youtube videos using OpenAI's Whisper
              </p>
            </div>
        """
    )
    with gr.Tab('Get a summary from your own mic or audio file'):
            input_audio = [
                gr.inputs.Audio(source="microphone", type="filepath", optional=True),
                gr.inputs.Audio(source="upload", type="filepath", optional=True),
            ]
            result_button_transcribe_audio = gr.Button('1. Transcribe')
            output_text_transcribe_audio = gr.Textbox(placeholder='Transcript of the audio file.', label='Transcript')
    
            result_button_summary_audio = gr.Button('2. Get a summary')
            output_text_summary_audio = gr.Textbox(placeholder='Summary of the audio file.', label='Summary')
    
            result_button_transcribe_audio.click(transcribe_inp, inputs=input_audio, outputs=output_text_transcribe_audio)
            result_button_summary_audio.click(getSummary, inputs=output_text_transcribe_audio, outputs=output_text_summary_audio)
    with gr.Tab('Summary of Youtube video'):
            input_text_url = gr.Textbox(placeholder='Youtube video URL', label='URL')
            result_button_transcribe = gr.Button('1. Transcribe')
            output_text_transcribe = gr.Textbox(placeholder='Transcript of the YouTube video.', label='Transcript')
    
            result_button_summary = gr.Button('2. Create Summary')
            output_text_summary = gr.Textbox(placeholder='Summary of the YouTube video transcript.', label='Summary')
    
            result_button_transcribe.click(getText, inputs=input_text_url, outputs=output_text_transcribe)
            result_button_summary.click(getSummary, inputs=output_text_transcribe, outputs=output_text_summary)

demo.launch(debug=True)