import gradio as gr
import os
from pytubefix import YouTube
from moviepy.editor import VideoFileClip
from transformers import pipeline

# Load Whisper model for transcription
asr = pipeline("automatic-speech-recognition", model="distil-whisper/distil-small.en")

# Load Summarization model
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

def process_youtube_link(youtube_url):
    try:
        # Download YouTube Video
        yt = YouTube(youtube_url)
        title = yt.title
        print(f"Downloading: {title}")

        video_stream = yt.streams.get_highest_resolution()
        if not video_stream:
            return "Error: No available video stream", ""

        video_path = f"{title}.mp4"
        video_stream.download(filename=video_path)

        # Extract Audio
        audio_path = f"{title}.wav"
        video = VideoFileClip(video_path)
        video.audio.write_audiofile(audio_path, codec="pcm_s16le")

        # Transcribe Audio
        transcription = asr(audio_path, return_timestamps=True)
        transcribed_text = transcription["text"]

        # Summarize Transcription
        summary = summarizer(transcribed_text, max_length=150, min_length=50, do_sample=False)[0]["summary_text"]

        # Clean up files after processing
        os.remove(video_path)
        os.remove(audio_path)

        return transcribed_text, summary

    except Exception as e:
        return f"Error: {str(e)}", ""

# Create Gradio Interface
iface = gr.Interface(
    fn=process_youtube_link,
    inputs=gr.Textbox(label="Enter YouTube URL"),
    outputs=[gr.Textbox(label="Transcription"), gr.Textbox(label="Summary")],
    title="YouTube Video Transcriber & Summarizer",
    description="Enter a YouTube link, and this app will transcribe and summarize the audio.",
)

iface.launch()