import os 
save_dir= os.path.join(os.getcwd(),'docs')
if not os.path.exists(save_dir):
    os.mkdir(save_dir)


llm_model_id = "tiiuae/falcon-7b-instruct"
HF_TOKEN = os.environ.get("HF_TOKEN", None)

from youtube_transcript_api import YouTubeTranscriptApi
import pytube

# get the transcript from YouTube
def get_yt_transcript(url):
    text = ''
    vid_id = pytube.extract.video_id(url)
    temp = YouTubeTranscriptApi.get_transcript(vid_id)
    for t in temp:
        text+=t['text']+' '
    return text

from pytube import YouTube
import transformers
import torch


from pytube import YouTube
from huggingface_hub import InferenceClient
from gradio_client import Client
# transcribes the video using the Hugging Face Hub API
def transcribe_file(url):
    client = Client("https://sanchit-gandhi-whisper-jax.hf.space/")
    response = client.predict(
        url,
        "transcribe",
        False,
        api_name="/predict_2"
    )
    return response[1]

def transcribe_youtube_video(url, force_transcribe=False,use_api=False,api_token=None):
    
    yt = YouTube(str(url))
    text = ''
    # get the transcript from YouTube if available
    try:
        text = get_yt_transcript(url)
    except:
        pass
    
    # transcribes the video if YouTube did not provide a transcription 
    # or if you want to force_transcribe anyway
    if text == '' or force_transcribe:
        text = transcribe_file(url)
        transcript_source = 'The transcript was generated using Whisper Jax.'
    else:
        transcript_source = 'The transcript was downloaded from YouTube.'
    
    return yt.title, text, transcript_source

def summarize_text(title,text,temperature,words,use_api=False,api_token=None,do_sample=False):
    
    from langchain_google_genai import ChatGoogleGenerativeAI
    from langchain.prompts import PromptTemplate
    from langchain.chains import LLMChain
    GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
    genai.configure()

    llm = ChatGoogleGenerativeAI(model="gemini-pro", google_api_key=GOOGLE_API_KEY)
    llm_model_id = 'Gemini-Pro'
    summary_source = 'The summary was generated using {} via Hugging Face API.'.format(llm_model_id)

    # Map templates
    prompt_template = """
    As an AI tasked with summarizing a video, your objective is to distill the key insights without introducing new information. This prompt aims to provide a concise summary.\n
    ----------------------- \n
    TITLE: `{title}`\n
    TEXT:\n
    `{docs}`\n 
    ----------------------- \n
    Summarize the provided content, emphasizing main points, key arguments, and relevant details. Keep the summary clear and succinct.\n
    SUMMARY:\n
    """


    map_prompt = PromptTemplate(
        template = map_template, 
        input_variables = ['title','docs']
        )
    map_chain = LLMChain(llm=llm, prompt=map_prompt)

    # Reduce - Collapse
    collapse_template = """
    As an AI tasked with combining partial summaries, your goal is to create a cohesive, comprehensive summary without duplications.\n
    ----------------------- \n
    TITLE: `{title}`\n
    PARTIAL SUMMARIES:\n
    `{doc_summaries}`\n 
    ----------------------- \n
    Synthesize the information from the partial summaries into a consolidated, coherent summary. Ensure that the final summary covers all essential points without repeating redundant information.\n
    CONSOLIDATED SUMMARY:\n
    """


    collapse_prompt = PromptTemplate(
        template = collapse_template, 
        input_variables = ['title','doc_summaries']
        )
    collapse_chain = LLMChain(llm=llm, prompt=collapse_prompt)

    # Takes a list of documents, combines them into a single string, and passes this to an LLMChain
    collapse_documents_chain = StuffDocumentsChain(
        llm_chain=collapse_chain, document_variable_name="doc_summaries"
        )

    # Final Reduce - Combine
    combine_template = """
    As an AI tasked with summarizing a video, your goal is to distill the main insights without introducing new information. This prompt aims to generate a concise executive summary.\n
    ----------------------- \n
    TITLE: `{title}`\n
    PARTIAL SUMMARIES:\n
    `{doc_summaries}`\n 
    ----------------------- \n
    Extract the most critical information from the partial summaries provided. Craft an executive summary in {words} words, focusing on the main arguments, key takeaways, and supporting evidence presented in the video. Aim for clarity, brevity, and avoid repeating redundant points. Ensure the summary encapsulates the essence of the content.\n
    EXECUTIVE SUMMARY:\n
    """

    combine_prompt = PromptTemplate(
        template = combine_template, 
        input_variables = ['title','doc_summaries','words']
        )
    combine_chain = LLMChain(llm=llm, prompt=combine_prompt)

    # Takes a list of documents, combines them into a single string, and passes this to an LLMChain
    combine_documents_chain = StuffDocumentsChain(
        llm_chain=combine_chain, document_variable_name="doc_summaries"
        )

    # Combines and iteratively reduces the mapped documents
    reduce_documents_chain = ReduceDocumentsChain(
        # This is final chain that is called.
        combine_documents_chain=combine_documents_chain,
        # If documents exceed context for `StuffDocumentsChain`
        collapse_documents_chain=collapse_documents_chain,
        # The maximum number of tokens to group documents into.
        token_max=800,
        )

    # Combining documents by mapping a chain over them, then combining results
    map_reduce_chain = MapReduceDocumentsChain(
        # Map chain
        llm_chain=map_chain,
        # Reduce chain
        reduce_documents_chain=reduce_documents_chain,
        # The variable name in the llm_chain to put the documents in
        document_variable_name="docs",
        # Return the results of the map steps in the output
        return_intermediate_steps=False,
        )

    from langchain.document_loaders import TextLoader
    from langchain.text_splitter import TokenTextSplitter
    
    with open(save_dir+'/transcript.txt','w') as f:
        f.write(text)
    loader = TextLoader(save_dir+"/transcript.txt")
    doc = loader.load()
    text_splitter = TokenTextSplitter(chunk_size=800, chunk_overlap=100)
    docs = text_splitter.split_documents(doc)

    summary = map_reduce_chain.run({'input_documents':docs, 'title':title, 'words':words})

    try:
        del(map_reduce_chain,reduce_documents_chain,combine_chain,collapse_documents_chain,map_chain,collapse_chain,llm)
    except:
        pass
    torch.cuda.empty_cache()

    return summary, summary_source

import gradio as gr 
import pytube
from pytube import YouTube

def get_youtube_title(url):
    yt = YouTube(str(url))
    return yt.title

def get_video(url):
    vid_id = pytube.extract.video_id(url)
    embed_html = '<iframe width="100%" height="315" src="https://www.youtube.com/embed/{}" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>'.format(vid_id)
    return embed_html

def summarize_youtube_video(url,force_transcribe,
                            temperature=0.2,words=300,do_sample=True):
    print("URL:",url)
    api_token = HF_TOKEN
    title,text,transcript_source = transcribe_youtube_video(url,force_transcribe,True,api_token)
    print("Transcript:",text)
    summary, summary_source = summarize_text(title,text,temperature,words,True,api_token,do_sample)
    print("Summary:",summary)
    return summary, text, transcript_source, summary_source

html = '<iframe width="100%" height="315" src="https://www.youtube.com/embed/" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>'

# def change_transcribe_api(vis):
#     return gr.Checkbox(value=False, visible=vis)

# def change_api_token(vis):
#     return gr.Textbox(visible=vis)

def update_source(source):
    return gr.Textbox(info=source)

def show_temp(vis):
    return gr.Slider(visible=vis)

# Defining the structure of the UI
with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# Summarize a YouTube Video")
    
    with gr.Row():
        with gr.Column(scale=4):
            url = gr.Textbox(label="Enter YouTube video URL here:",placeholder="https://www.youtube.com/watch?v=",info="The video must not be age-restricted. Otherwise, the transcription will fail. The demo supports videos in English language only.")
        with gr.Column(scale=2):
            sum_btn = gr.Button("Summarize!")
            gr.Markdown("## Please like the repo if you find this helpful.")


    with gr.Accordion("Transcription Settings",open=False):
        with gr.Row():
            force_transcribe = gr.Checkbox(label="Transcribe even if transcription is available.", info='If unchecked, the app attempts to download the transcript from YouTube first. Check this if the transcript does not seem accurate.')
            # use_transcribe_api = gr.Checkbox(label="Transcribe using the HuggingFaceHub API.",visible=False)
        
    with gr.Accordion("Summarization Settings",open=False):
        with gr.Row():
            # use_llm_api = gr.Checkbox(label="Summarize using the HuggingFaceHub API.",visible=True)
            do_sample = gr.Checkbox(label="Set the Temperature",value=False,visible=True)
            temperature = gr.Slider(minimum=0.01,maximum=1.0,value=0.2,label="Generation temperature",visible=False)
            words = gr.Slider(minimum=100,maximum=500,value=300,label="Length of the summary")
            
    gr.Markdown("# Results")

    title = gr.Textbox(label="Video Title",placeholder="title...")
    
    with gr.Row():
        video = gr.HTML(html,scale=1)
        summary_source = gr.Textbox(visible=False,scale=0)
        summary = gr.Textbox(label="Summary",placeholder="summary...",scale=1)
    
    with gr.Row():
        with gr.Group():
            transcript = gr.Textbox(label="Full Transcript",placeholder="transcript...",show_label=True)
            transcript_source = gr.Textbox(visible=False)
    with gr.Accordion("Acknoledgement",open=True):
        gr.Markdown("""
                    I sincerely appreciate the open source tools shared by [smakamali](https://huggingface.co/smakamali) (summary_method) and [Sanchit Gandhi](https://huggingface.co/sanchit-gandhi) (Whisper-Jax API) 
                    which were instrumental in developing this project. Their publicly available innovations in AI model training and speech recognition directly enabled key capabilities. Please view their exceptional repositories on HuggingFace for additional details.\n
                    
                    [summarize_youtube](https://huggingface.co/spaces/smakamali/summarize_youtube)\n
                    Detailed instructions for recreating this tool are provided [here](https://pub.towardsai.net/a-complete-guide-for-creating-an-ai-assistant-for-summarizing-youtube-videos-part-1-32fbadabc2cc?sk=34269402931178039c4c3589df4a6ec5) and [here](https://pub.towardsai.net/a-complete-guide-for-creating-an-ai-assistant-for-summarizing-youtube-videos-part-2-a008ee18f341?sk=d59046b36a52c74dfa8befa99183e5b6).\n
                    [Whisper-Jax-api](https://sanchit-gandhi-whisper-jax.hf.space/)\n 
                    """)
    with gr.Accordion("Disclaimer",open=False):
        gr.Markdown("""
                    1. This app attempts to download the transcript from Youtube first. If the transcript is not available, or the prompts require, the video will be transcribed.\n
                    2. The app performs best on videos in which the number of speakers is limited or when the YouTube transcript includes annotations of the speakers.\n 
                    3. The trascription does not annotate the speakers which may downgrade the quality of the summary if there are more than one speaker.\n
                    """)
    
    # Defining the interactivity of the UI elements
    # force_transcribe.change(fn=change_transcribe_api,inputs=force_transcribe,outputs=use_transcribe_api)
    # use_transcribe_api.change(fn=change_api_token,inputs=use_transcribe_api,outputs=api_token)    
    # use_llm_api.change(fn=change_api_token,inputs=use_llm_api,outputs=api_token)
    transcript_source.change(fn=update_source,inputs=transcript_source,outputs=transcript)
    summary_source.change(fn=update_source,inputs=summary_source,outputs=summary)
    do_sample.change(fn=show_temp,inputs=do_sample,outputs=temperature)
    
    # Defining the functions to call on clicking the button
    sum_btn.click(fn=get_youtube_title, inputs=url, outputs=title, api_name="get_youtube_title", queue=False)
    sum_btn.click(fn=summarize_youtube_video, inputs=[url,force_transcribe,temperature,words,do_sample], 
                  outputs=[summary,transcript, transcript_source, summary_source], api_name="summarize_youtube_video", queue=True)
    sum_btn.click(fn=get_video, inputs=url, outputs=video, api_name="get_youtube_video", queue=False)

demo.queue()
demo.launch(share=False)