import gradio as gr
from youtube_transcript_api import YouTubeTranscriptApi
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import re
import torch

# ✅ FORCE lightweight model
model_name = "sshleifer/distilbart-cnn-12-6"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# move to CPU explicitly (safe)
device = "cpu"
model = model.to(device)


def extract_video_id(url):
    regex = r"(?:v=|\/)([0-9A-Za-z_-]{11})"
    match = re.search(regex, url)
    return match.group(1) if match else None


def get_transcript(video_id):
    try:
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        return " ".join([t['text'] for t in transcript])
    except:
        return None


def summarize_text(text):
    max_chunk = 500   # 🔥 smaller chunks = safer
    chunks = [text[i:i+max_chunk] for i in range(0, len(text), max_chunk)]

    final_summary = ""

    for chunk in chunks:
        inputs = tokenizer(
            "summarize: " + chunk,
            return_tensors="pt",
            max_length=512,
            truncation=True
        ).to(device)

        summary_ids = model.generate(
            **inputs,
            max_length=100,
            min_length=25,
            num_beams=2,   # 🔥 reduced for speed
        )

        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        final_summary += summary + " "

    return final_summary


def process_video(url):
    video_id = extract_video_id(url)

    if not video_id:
        return "Invalid URL", "", ""

    transcript = get_transcript(video_id)

    if not transcript:
        return "Transcript not available", "", ""

    summary = summarize_text(transcript)

    video_embed = f"""
    <iframe width="100%" height="315"
    src="https://www.youtube.com/embed/{video_id}"
    frameborder="0" allowfullscreen></iframe>
    """

    return summary, transcript[:2000], video_embed  # 🔥 limit transcript


with gr.Blocks() as demo:
    gr.Markdown("# 🎥 YouTube Video Summarizer")

    url_input = gr.Textbox(label="Enter YouTube URL")
    btn = gr.Button("Generate")

    summary_output = gr.Textbox(label="Summary")
    transcript_output = gr.Textbox(label="Transcript (trimmed)")
    video_output = gr.HTML()

    btn.click(
        process_video,
        inputs=url_input,
        outputs=[summary_output, transcript_output, video_output]
    )

demo.launch(ssr_mode=False)