import os 
import time
import gradio as gr
import logging
from youtube_transcript_api import YouTubeTranscriptApi
from langchain.docstore.document import Document
from langchain_groq import ChatGroq
from langchain.chains.summarize.chain import load_summarize_chain
import chatops 

logger = logging.getLogger(__name__)

DEVICE = 'cpu'

DEFAULT_CHAR_LENGTH = 1000


def youtube_link_dataloader(video_link,max_video_length=1000):
    video_text = ""
    meta_data = {"source": f"{video_link}"} 
    video_id = video_link.split("watch?v=")[1].split("&")[0]
    srt = YouTubeTranscriptApi.get_transcript(video_id)
    for text_data in srt:
        video_text = video_text + " " + text_data.get("text")
    if len(video_text) > max_video_length:
        video_text = video_text[0:max_video_length]
    document = [Document(page_content= video_text, metadata= meta_data)]
    return document


def youtube_chat(API_key=None,llm_service='mistralai/Mistral-7B-v0.1',youtube_link=None,char_length=2000):
    
    video_document  = youtube_link_dataloader(video_link=youtube_link,max_video_length=char_length)
    print("docuemt:",video_document)
    
    if  llm_service== 'mistralai/Mistral-7B-v0.1':
        llm = chatops.get_hugging_face_model(
                            model_id="mistralai/Mistral-7B-v0.1",
                            API_key=API_key,
                            temperature=0.1,
                            max_tokens=2048
                            )
    elif llm_service == 'OpenAI':
        llm = chatops.get_openai_chat_model(API_key=API_key)
    elif llm_service == 'llama3-8b-8192':
        os.environ["GROQ_API_KEY"] = API_key
        llm = ChatGroq(model="llama3-8b-8192")

    summarize_chain = load_summarize_chain(llm=llm, chain_type='stuff', verbose = True )
    results = summarize_chain.invoke({'input_documents':video_document})
    return results['output_text']

iface = gr.Interface(
    fn = youtube_chat,
    inputs = [
        gr.Textbox(label="Add API key", type="password"),
        gr.Dropdown(['mistralai/Mistral-7B-v0.1','llama3-8b-8192'],label='Large Language Model',info='LLM Service'),
        gr.Textbox(label='You tube link'),
        gr.Slider(DEFAULT_CHAR_LENGTH,5000,label="Video link Length in seconds",info="Length of video in seconds")
    ],
    outputs="text",
    description ="""Summarize your You tube link using Large Language Models  
        The Objective of the space is to use the Large Language models to generate a small Summary of the You tube Link provided.
        It Facilitates to generate notes if you are using you tube for Educational purposes.
    """,
)

iface.launch()