"""
YouTube Video Summarizer – Gradio Space
"""
import os
import gradio as gr
import whisper
from pytube import YouTube
from transformers import pipeline
import torch

# ── 1. Load models once at startup -------------------------
print("⏳ Loading Whisper base model...")
whisper_model = whisper.load_model("base")          # 74 M params, fast on CPU
print("⏳ Loading BART-large-cnn...")
summarizer = pipeline("summarization",
                      model="facebook/bart-large-cnn",
                      torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)

# ── 2. Core pipeline ---------------------------------------
def summarize_video(url: str) -> str:
    try:
        # a) Download audio stream
        yt = YouTube(url)
        audio_stream = yt.streams.filter(only_audio=True).first()
        if not audio_stream:
            return "❌ No audio stream found."

        audio_file = audio_stream.download(filename="audio.mp4")

        # b) Transcribe
        result = whisper_model.transcribe(audio_file)
        transcript = result["text"]

        # c) Summarize (BART has 1024 token limit; chunk if necessary)
        max_chunk = 900  # tokens
        chunks = [transcript[i:i+max_chunk] for i in range(0, len(transcript), max_chunk)]
        summary_parts = []
        for chunk in chunks:
            summary_parts.append(
                summarizer(chunk,
                           max_length=120,
                           min_length=30,
                           do_sample=False)[0]["summary_text"]
            )
        summary = " ".join(summary_parts)

        # cleanup
        os.remove(audio_file)
        return summary
    except Exception as e:
        return f"⚠️ Error: {e}"

# ── 3. Gradio UI -------------------------------------------
iface = gr.Interface(
    fn=summarize_video,
    inputs=gr.Textbox(label="Paste YouTube URL"),
    outputs=gr.Textbox(label="Summary", lines=10),
    title="YouTube Video Summarizer 🤖",
    description="Paste any public YouTube link and get an AI-generated summary in ~30 s."
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)