import gradio as gr
from langchain import PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.chains import LLMChain
from langchain_community.retrievers import WikipediaRetriever
import os


def song_meaning(song, artist):
    artist_input = artist.title()
    song_input = song.title()
    query_input = f"{song_input} by {artist_input}"

    retriever = WikipediaRetriever()
    docs = retriever.get_relevant_documents(query=query_input)

    template_song_meaning = """
      {artist} has released a song called {song}.

      {content}

      based on the the content above what does the song {song} by {artist} tell us about? give me a long explanations

      """

    prompt_template_song_meaning = PromptTemplate(input_variables=["artist", "song", "content"],
                                                  template=template_song_meaning)

    llm = ChatOpenAI(openai_api_key=os.environ['OPENAI_API_KEY'], model_name="gpt-3.5-turbo", temperature=0)
    chain = LLMChain(llm=llm, prompt=prompt_template_song_meaning)
    results = chain.run(artist=artist_input, song=song_input, content=docs[0].page_content)

    return results


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    song = gr.Textbox(label="Song")
    artist = gr.Textbox(label="Artist")
    output = gr.Textbox(label="Meaning")
    gr.Interface(fn=song_meaning, inputs=[song, artist], outputs=output)
    example = gr.Examples([['Maroon', 'Taylor Swift'], ['Devil In Her Heart', 'The Beatles'],
                           ['Time Machine', 'Jay Chou'], ['Last Farewell', 'BIGBANG']], [song, artist])

demo.launch()