from youtube_transcript_api import YouTubeTranscriptApi
from youtube_transcript_api._errors import TranscriptsDisabled, NoTranscriptFound, VideoUnavailable
import torch
import gradio as gr

# Use a pipeline as a high-level helper
from transformers import pipeline

text_summary = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16)
# Load model directly
# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
# model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-cnn-12-6")

def summarize_text(input):
    output = text_summary(input)
    return output[0]['summary_text']


# def summarize_text(text, max_length=130, min_length=30):
#     """
#     Summarizes the input text using the DistilBART model.

#     Parameters:
#         text (str): The input text to summarize.
#         max_length (int): Maximum length of the summary.
#         min_length (int): Minimum length of the summary.

#     Returns:
#         str: The summarized text.
#     """
#     # Load the tokenizer and model
#     tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
#     model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-cnn-12-6")
    
#     # Tokenize the input text with truncation
#     inputs = tokenizer(
#         text,
#         max_length=1024,  # DistilBART's max input length
#         truncation=True,
#         return_tensors="pt"
#     )
    
#     # Generate the summary
#     summary_ids = model.generate(
#         inputs["input_ids"],
#         max_length=max_length,
#         min_length=min_length,
#         length_penalty=2.0,
#         num_beams=4,
#         early_stopping=True
#     )
    
#     # Decode the generated summary
#     summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
#     return summary

def get_youtube_video_id(url):
    """
    Extracts the video ID from a YouTube URL.
    """
    import re
    video_id_pattern = r'(?:https?:\/\/)?(?:www\.)?youtu(?:be\.com\/(?:[^\/\n\s]+\/\S+\/|(?:v|e(?:mbed)?)\/|\S*?[?&]v=)|\.be\/)([a-zA-Z0-9_-]{11})'
    match = re.match(video_id_pattern, url)
    if match:
        return match.group(1)
    else:
        raise ValueError("Invalid YouTube URL provided.")

def get_youtube_transcript(url):
    """
    Fetches the transcript for a given YouTube video URL.
    """
    try:
        video_id = get_youtube_video_id(url)
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        # Combine transcript text into a single string
        full_transcript = "\n".join([entry['text'] for entry in transcript])
        summary_text = summarize_text(full_transcript)
        return summary_text
        return full_transcript
    except TranscriptsDisabled:
        return "Transcripts are disabled for this video."
    except NoTranscriptFound:
        return "No transcript found for this video."
    except VideoUnavailable:
        return "The video is unavailable or the URL is incorrect."
    except Exception as e:
        return f"An error occurred: {str(e)}"

if __name__ == "__main__":
    # Input YouTube video URL

    # youtube_url = input("Enter the YouTube video URL: ").strip()
    # transcript = get_youtube_transcript(youtube_url)
    # print("\n=== Transcript ===\n")
    # print(transcript)
    # =============================================================

    gr.close_all()

    demo = gr.Interface(fn=get_youtube_transcript,
                        inputs=[gr.Textbox(label="Input Youtube URL", lines=2)],
                        outputs=[gr.Textbox(label="Youtube Summarized Content", lines=10)],
                        title = "Gen AI Youtube Summarization",
                        description= "Enter youtube URL the tool will give you summarized text"
                        )
    demo.launch(share=True)