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

# Load environment variables from .env file
load_dotenv()

# Get the user ID and password from environment variables
USER_ID = os.getenv('USER_ID')
PASSWORD = os.getenv('PASSWORD')

# Initialize the summarization pipeline
text_summary = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16)

def summary(text_youtube):
    output = text_summary(text_youtube)
    return output[0]['summary_text']

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

def get_youtube_transcript(video_url):
    video_id = extract_video_id(video_url)
    if not video_id:
        return "Video ID could not be extracted.", ""

    try:
        # Fetch the transcript
        transcript = YouTubeTranscriptApi.get_transcript(video_id)

        # Format the transcript into plain text
        formatter = TextFormatter()
        text_transcript = formatter.format_transcript(transcript)
        summary_text = summary(text_transcript)

        return text_transcript, summary_text
    except Exception as e:
        return f"An error occurred: {e}", ""

# Authentication function
def authenticate(username, password):
    if username == USER_ID and password == PASSWORD:
        return True
    return False

def secure_get_youtube_transcript(username, password, video_url):
    if authenticate(username, password):
        return get_youtube_transcript(video_url)
    else:
        return "Authentication failed", ""

# Remove any existing Gradio interfaces
gr.close_all()

# Define the Gradio interface
demo = gr.Interface(
    fn=secure_get_youtube_transcript,
    inputs=[
        gr.Textbox(label="Username", lines=1),
        gr.Textbox(label="Password", lines=1, type="password"),
        gr.Textbox(label="Input YouTube URL to summarize", lines=1)
    ],
    outputs=[
        gr.Textbox(label="Text extracted from video", lines=4),
        gr.Textbox(label="Summarized text", lines=4)
    ],
    title="YouTube Script for Text Extract & Summarizer"
)

# Launch the Gradio interface
demo.launch()