import os

import openai
import streamlit as st
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.llms import OpenAI as l_OpenAI
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM

from helpers.foundation_models import *
import requests


OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
SERPAPI_API_KEY = os.environ["SERPAPI_API_KEY"]
openai_client = openai.OpenAI(api_key=OPENAI_API_KEY)


API_URL = "https://sks7h7h5qkhoxwxo.us-east-1.aws.endpoints.huggingface.cloud"
headers = {
	"Accept" : "application/json",
	"Content-Type": "application/json" 
}


def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.json()


def llama2_7b_ysa(prompt: str) -> str:
    output = query({
        "inputs": prompt,
        "parameters": {}
    })

    response = output[0]['generated_text']

    return response

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []


# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])


with st.expander("Instructions"):
    st.sidebar.markdown(
        r"""
        # 🌟 Streamlit + Hugging Face Demo 🤖

        ## Introduction 📖

        This demo showcases how to interact with Large Language Models (LLMs) on Hugging Face using Streamlit. 
        """
    )


option = st.sidebar.selectbox(
    "Which task do you want to do?",
    ("Sentiment Analysis", "Medical Summarization", "Llama2 on YSA", "ChatGPT", "ChatGPT (with Google)"),
)


clear_button = st.sidebar.button("Clear Conversation", key="clear")

st.sidebar.write("---")

st.sidebar.markdown("Yiqiao Yin: [Site](https://www.y-yin.io/) | [LinkedIn](https://www.linkedin.com/in/yiqiaoyin/)")


# Reset everything
if clear_button:
    st.session_state.messages = []


# React to user input
if prompt := st.chat_input("What is up?"):
    # Display user message in chat message container
    st.chat_message("user").markdown(prompt)
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.spinner("Wait for it..."):
        if option == "Sentiment Analysis":
            pipe_sentiment_analysis = pipeline("sentiment-analysis")
            if prompt:
                out = pipe_sentiment_analysis(prompt)
                final_response = f"""
                    Prompt: {prompt}
                    Sentiment: {out[0]["label"]}
                    Score: {out[0]["score"]}
                """
        elif option == "Medical Summarization":
            pipe_summarization = pipeline(
                "summarization", model="Falconsai/medical_summarization"
            )
            if prompt:
                out = pipe_summarization(prompt)
                final_response = out[0]["summary_text"]
        elif option == "Llama2 on YSA":
            if prompt:
                out = llama2_7b_ysa(query=prompt)
                engineered_prompt = f"""
                    The user asked the question: {prompt}

                    We have found relevant content: {out}

                    Answer the user question based on the above content in paragraphs.
                """
                final_response = call_chatgpt(query=engineered_prompt)
        elif option == "ChatGPT":
            if prompt:
                out = call_chatgpt(query=prompt)
                final_response = out
        elif option == "ChatGPT (with Google)":
            if prompt:
                ans_langchain = call_langchain(prompt)
                prompt = f"""
                    Based on the internet search results: {ans_langchain};

                    Answer the user question: {prompt}
                """
                out = call_chatgpt(query=prompt)
                final_response = out
        else:
            final_response = ""

    response = f"{final_response}"
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st.markdown(response)
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})