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

model = whisper.load_model("base")
summarizer = pipeline("summarization")

def get_audio(url):
    yt = YouTube(url)
    video = yt.streams.filter(only_audio=True).first()
    out_file = video.download(output_path=".")
    base, ext = os.path.splitext(out_file)
    new_file = base + '.mp3'
    os.rename(out_file, new_file)
    return new_file

def get_text_from_url(url):
    result = model.transcribe(get_audio(url))
    return result['text']

def get_text_from_file(file):
    # Assuming the uploaded file is already in MP3 format
    result = model.transcribe(file.name)
    return result['text']

def get_summary_from_url(url):
    article = get_text_from_url(url)
    b = summarizer(article)
    return b[0]['summary_text']

def get_summary_from_file(file):
    article = get_text_from_file(file)
    b = summarizer(article)
    return b[0]['summary_text']

def process_url(url):
    transcription = get_text_from_url(url)
    summary = get_summary_from_url(url)
    return summary, transcription

def process_file(file):
    transcription = get_text_from_file(file)
    summary = get_summary_from_file(file)
    return summary, transcription

with gr.Blocks() as demo:
    gr.HTML("<h1><center>Youtube and Video File Upload with Whisper Transcription and Summary</center></h1>")
    gr.HTML("<p>Note to Users:<br>Enter the link of any youtube video or upload an MP4 file to get the transcription and a summary in text form. Note: I'm using a git trick in the requirements file to run this without an openai API Key, if you want a little more speed (because this is really slow) and want to use an openai API Key check out the code base at <a href=\"https://huggingface.co/spaces/eaglelandsonce/ChatGPT_Enhanced\" target=\"_blank\">Chat Toolkit</a>, if you want to interact live with folks online check out my Meetup at <a href=\"https://www.meetup.com/florence-aws-user-group-meetup/\" target=\"_blank\">AWS Cloud Meetup</a></p>")

    with gr.Tab('Youtube Video'):
        with gr.Row():
            input_text = gr.Textbox(placeholder='Enter the Youtube video URL', label='URL')
            output_summary = gr.Textbox(placeholder='Summary text of the Youtube Video', label='Summary')
            output_transcription = gr.Textbox(placeholder='Transcription of the video', label='Transcription')
        result_button = gr.Button('Process Youtube Video')
        
    with gr.Tab('Uploaded MP4'):
        with gr.Row():
            input_file = gr.File(label='Upload MP4')
            output_file_summary = gr.Textbox(placeholder='Summary text of the video', label='Summary')
            output_file_transcription = gr.Textbox(placeholder='Transcription of the video', label='Transcription')
        result_button_file = gr.Button('Process Uploaded MP4')

    result_button.click(process_url, inputs=input_text, outputs=[output_summary, output_transcription])
    result_button_file.click(process_file, inputs=input_file, outputs=[output_file_summary, output_file_transcription])
    
demo.launch(debug=True)