import gradio as gr
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import os

api_key = os.getenv('GROQ_API_KEY')

def get_transcript(video_url):
    from langchain_community.document_loaders import YoutubeLoader
    loader = YoutubeLoader.from_youtube_url(
        video_url, add_video_info=False
    )
    document = loader.load()
    transcript = document[0].page_content
    return transcript


def summarize_video(video_url, language):
    transcript = get_transcript(video_url)

    model = ChatGroq(
        model="llama-3.1-70b-versatile",
        temperature=0,
        max_tokens=None,
        timeout=None,
        max_retries=2,
        api_key=api_key
    )

    system_template = """Below you will see a text. Read the text. First provide the major points that they discuss and then provide summaries for each major point. Your response should be in the language specified below. 
    Transcript: {transcript}
    Language: {language}"""

    prompt_template = ChatPromptTemplate.from_messages(
        [("system", system_template)]
    )

    parser = StrOutputParser()

    chain = prompt_template | model | parser

    response = chain.invoke({"transcript": transcript, "language": language})
    return response


iface = gr.Interface(
    fn=summarize_video,
    inputs=[
        gr.Textbox(label="YouTube Video URL"),
        gr.Textbox(label="Language for Summary")
    ],
    outputs="text",
    title="YouTube Video Summarizer",
    description="Enter a YouTube video URL and the desired language for the summary."
)

iface.launch()