import re
from youtube_transcript_api import YouTubeTranscriptApi
from youtube_transcript_api.formatters import TextFormatter
import torch
import gradio as gr
from transformers import pipeline

text_summary = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16)
text_translate = pipeline("translation", model="google-t5/t5-small", torch_dtype=torch.bfloat16)

def summary (input):
    max_length = 2048  # adjust this value as needed
    if len(input) > max_length:
        input = input[:max_length]
    output = text_summary(input)
    return output[0]['summary_text']

def extract_video_id(url):
    # Regex to extract the video ID from various YouTube URL formats
    regex = r"(?:youtube\.com\/(?:[^\/\n\s]+\/\S+\/|(?:v|e(?:mbed)?)\/|\S*?[?&]v=)|youtu\.be\/)([a-zA-Z0-9_-]{11})"
    match = re.search(regex, url)
    if match:
        return match.group(1)
    return None


def get_german_youtube_transcript(video_url):
    """Lädt das Transkript eines deutschen YouTube-Videos herunter.
    Args:
        video_url (str): Die URL des YouTube-Videos.
    Returns:
        str: Das Transkript des Videos in deutscher Sprache.
    """

    # Extrahiere die Video-ID aus der URL
    video_id = extract_video_id(video_url)
    
    if not video_id:
        return "Die Video-ID konnte nicht extrahiert werden."
    
    try:
        # Hole das Transkript in deutscher/engl. Sprache
        transcript = YouTubeTranscriptApi.get_transcript(video_id, languages=['de', 'en'])
        
        # Formatiere das Transkript in Klartext
        formatter = TextFormatter()
        text_transcript = formatter.format_transcript(transcript)
        summary_text = summary(text_transcript)
        
        summary_text = summary_text[:300] 
        #text_transcript = text_transcript[:300]  # truncate to 300 characters

        #transcript_output = text_translate(text_transcript)
        #transcript_output = transcript_output[0]['translation_text']
        summary_output = text_translate(summary_text)
        return summary_output[0]['translation_text'], text_transcript
    except Exception as e:
        return f"Es ist ein Fehler aufgetreten: {e}"

gr.close_all()

demo = gr.Interface(fn=get_german_youtube_transcript,
                    inputs=[gr.Textbox(label="YouTube Url zum Zusammenfassen eingeben",lines=1)],
                    outputs=[
                        gr.Textbox(label="Zusammengefasster Text (DE)",lines=4),
                        gr.Textbox(label="Video-Transkript (EN)", lines=10) 
                    ],
                    title="Projekt 2: YouTube-Video(EN)-Zusammenfassung (DE)",
                    description="Diese Anwendung wird genutzt, um den Inhalt eines YouTube-Videos auf Englisch in eine deutsche Zusammenfassung zu übersetzen.",
                    allow_flagging="never",
                    submit_btn="Übermitteln",
                    clear_btn="Bereinigen",
                    )
demo.launch()