from util import load_model
from util import pipeline

import gradio as gr
cp_aug = 'minnehwg/finetune-newwiki-summarization-ver-augmented2'


def get_model(cp):
    checkpoint = cp
    tokenizer, model = load_model(checkpoint)
    return tokenizer, model

tokenizer, model = get_model(cp_aug)


def generate_summary(url):
    results = pipeline(url, model, tokenizer)
    summary = "\n".join(results)
    return summary

def generate_summary_and_video(url):
    summary = generate_summary(url)
    try:
        video_id = url.split("v=")[1].split("&")[0]
        iframe = f'<iframe width="560" height="315" src="https://www.youtube.com/embed/{video_id}" frameborder="0" allow="accelerometer; autoplay; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>'
        return f"{iframe}\n\n**Summary:**\n{summary}"
    except IndexError:
        return f"**Summary:**\n{summary}\n\nInvalid YouTube URL for video display."

demo = gr.Interface(
    fn=generate_summary_and_video,
    inputs=gr.Textbox(lines=2, placeholder="Enter URL..."),
    outputs=gr.HTML(label="YouTube Video and Summary"),
    title="YouTube Video and Summary Display",
    description="Enter the URL '.'"
)


demo.launch(share=True)