import os
import streamlit as st
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, SequentialChain
from langchain.memory import ConversationBufferMemory
from langchain.utilities import WikipediaAPIWrapper
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load the GPT-2 model and tokenizer
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
OpenAI = GPT2LMHeadModel.from_pretrained(model_name)


# app framework
st.title('Youtube Content Studio 🚀')
prompt=st.text_input('Plug in your prompt here')

#Prompt templates
title_template=PromptTemplate(
    input_variables=['topic'],
    template='write me a youtube video title about {topic}'
)
script_template=PromptTemplate(
    input_variables=['title','wikipedia_research'],
    template='write me a youtube script based on this title TITLE: {title} while leveraging this wikipedia research: {wikipedia_research}'
)

# Memory
title_memory=ConversationBufferMemory(input_key='topic',memory_key='chat_history')
script_memory=ConversationBufferMemory(input_key='title',memory_key='chat_history')
# LLMS
llm=OpenAI(temperature=0.9)
title_chain=LLMChain(llm=llm,prompt=title_template,verbose=True,output_key='title',memory=title_memory)
script_chain=LLMChain(llm=llm,prompt=script_template,verbose=True,output_key='script',memory=script_memory)

wiki=WikipediaAPIWrapper()

# Show stuff to the screen if there is a prompt
if prompt:
    title=title_chain.run(prompt)
    wiki_research=wiki.run(prompt)
    script=script_chain.run(title=title,wikipedia_research=wiki_research)
    st.write(title)
    st.write(script)
    
    
    with st.expander('Title History'):
        st.info(title_memory.buffer)
    with st.expander('Script History'):
        st.info(script_memory.buffer)    
    with st.expander('Wikipedia research History'):
        st.info(wiki_research)