import openai
import streamlit as st


openai.api_key = st.secrets['OpenAI_API_Key']

def main():
    # Use markdown for the title with a custom color
    st.sidebar.markdown("<h1 style='color: #000000;'>Chat w/MidikanGPT-4</h1>", unsafe_allow_html=True)

    # Creating a text input field for user question context
    context = st.sidebar.text_area("Context (if any):")
    
    # Creating a text input field for user questions
    question = st.sidebar.text_area("Question/command:")
    
    # Creating a button for the user to send their message
    send_button = st.sidebar.button("Send")

    # Initialize chat history in the first run
    if 'chat_history' not in st.session_state:
        st.session_state['chat_history'] = ''
    
    # Handle the send button click
    if send_button:
        # Set the system prompt
        system_prompt = "ChatGPT, please act as a highly experienced, intelligent assistant who has access to vast amounts of knowledge across a wide range of disciplines. Your responses should demonstrate a deep understanding of any given topic, presenting nuanced, comprehensive information in a detailed and understandable manner. When asked a question or presented with a task, your approach should be to provide not only the answer or solution but also to offer relevant context, potential implications, comparisons and contrasts, and any other pertinent information that could deepen the user's understanding of the topic or task at hand. For complex or multifaceted inquiries, structure your responses to first offer a succinct summary, followed by a deeper dive into the matter. Consider the needs and expertise of the person asking the question and tailor your explanations accordingly, being aware that while some may require simplified explanations, others may prefer more technical details. And above all, prioritize the usefulness and accuracy of the information provided. If any additional details or follow-up actions are needed, proactively suggest them in your responses."

        # Prepare 'prompt_1'
        txt_1 = 'Question. '
        txt_2 = "Answer: Let's work this out in a step by step way to be sure we have the right answer."
        prompt_1 = context + '\n\n' + txt_1 + question + '\n' + txt_2
        prompt_1 = prompt_1.strip()
        
        # Save 'prompt_1' to the chat history
        st.session_state['chat_history'] += "<div style='color: blue;'>You: </div>"
        st.session_state['chat_history'] += f"{prompt_1}"
        st.session_state['chat_history'] += "<br><hr style='border: 1px solid black;'>"

        # Hit the OpenAI API with 'prompt_1'
        response = openai.ChatCompletion.create(
            model='gpt-4',
            messages=[
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': prompt_1}
            ],
            temperature=1, # Default value
            n=3 # Number of chat completions to generate
        )

        messages = response['choices']

        out_list = []
        for idx, message in enumerate(messages):
            idx = idx + 1
            output = 'Answer Option ' + str(idx) + '.\n' + message['message']['content']
            out_list.append(output)
        
        answers_1 = '\n\n'.join(out_list)
        
        # Save GPT's response to the chat history
        st.session_state['chat_history'] += "<div style='color: red;'>MidikanGPT-4: </div>"
        st.session_state['chat_history'] += f"{answers_1}"
        st.session_state['chat_history'] += "<br><hr style='border: 1px solid black;'>"
        
        # Prepare 'prompt_2'
        prompt_2 = context + '\n\n' + txt_1 + question + '\n\n' + answers_1 + '\n\n' + "You are a researcher tasked with investigating the 3 answer options provided. List the flaws and faulty logic of each answer option. Let's work this out in a step by step way to be sure we have all the errors:"
        prompt_2 = prompt_2.strip()
        
        # Save 'prompt_2' to the chat history
        st.session_state['chat_history'] += "<div style='color: blue;'>You: </div>"
        st.session_state['chat_history'] += "You are a researcher tasked with investigating the 3 answer options provided. List the flaws and faulty logic of each answer option. Let's work this out in a step by step way to be sure we have all the errors:"
        st.session_state['chat_history'] += "<br><hr style='border: 1px solid black;'>"
        
        # Hit the OpenAI API with 'prompt_2'
        response = openai.ChatCompletion.create(
            model='gpt-4',
            messages=[
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': prompt_2}
            ],
            temperature=1 # Default value
        )
        
        answers_2 = response['choices'][0]['message']['content']
        
        # Save GPT's response to the chat history
        st.session_state['chat_history'] += "<div style='color: red;'>MidikanGPT-4: </div>"
        st.session_state['chat_history'] += f"{answers_2}"
        st.session_state['chat_history'] += "<br><hr style='border: 1px solid black;'>"
        
        # Prepare 'prompt_3'
        prompt_3 = answers_1 + '\n\n' + answers_2 + '\n\n' + "You are a resolver tasked with 1) finding which of the 3 answer options the researcher thought was best 2) improving that answer, and 3) printing the improved answer in full. Let's work this out in a step by step way to be sure we have the right answer:"

        # Save the end of 'prompt_3' to the chat history
        st.session_state['chat_history'] += "<div style='color: blue;'>You: </div>"
        st.session_state['chat_history'] += "You are a resolver tasked with 1) finding which of the 3 answer options the researcher thought was best 2) improving that answer, and 3) printing the improved answer in full. Let's work this out in a step by step way to be sure we have the right answer:"
        st.session_state['chat_history'] += "<br><hr style='border: 1px solid black;'>"
        
        # Hit the OpenAI API with 'prompt_3'
        response = openai.ChatCompletion.create(
            model='gpt-4',
            messages=[
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': prompt_3}
            ],
            temperature=1 # Default value
        )

        answers_3 = response['choices'][0]['message']['content']
        
        # Save GPT's response to the chat history
        st.session_state['chat_history'] += "<div style='color: red;'>MidikanGPT-4: </div>"
        st.session_state['chat_history'] += f"{answers_3}"
        st.session_state['chat_history'] += "<br><hr style='border: 1px solid black;'>"
        
    # Creating an area to display the chat history
    st.markdown(st.session_state['chat_history'], unsafe_allow_html=True)


if __name__ == "__main__":
    main()