import logging
import json
import hmac
import streamlit as st
from openai import OpenAI

def check_password():
    """Returns `True` if the user had a correct password."""

    def login_form():
        """Form with widgets to collect user information"""
        with st.form("Credentials"):
            st.text_input("Username", key="username")
            st.text_input("Password", type="password", key="password")
            st.form_submit_button("Log in", on_click=password_entered)

    def password_entered():
        """Checks whether a password entered by the user is correct."""
        if hmac.compare_digest(st.session_state["password"], st.secrets["PASSWORD"]) \
            and hmac.compare_digest(st.session_state["username"], st.secrets["USERNAME"]):

            st.session_state["password_correct"] = True
            #del st.session_state["password"]  # Don't store the password.
            #del st.session_state["username"]
        else:
            st.session_state["password_correct"] = False

    if st.session_state.get("password_correct", False):
        return True

    login_form()

    if "password_correct" in st.session_state:
        st.error("😕 User not known or password incorrect")
    return False

if not check_password():
    st.stop()

def on_suggested_reply_click_handler():
    response = st.session_state.reply_suggestion
    st.session_state.messages.append({"role": "Emily", "content": response})

logging.basicConfig(filename="demo.log", filemode="a+", level=logging.INFO)

openai_api_key = st.secrets["OPEN_AI_API_KEY"]
client = OpenAI(api_key=openai_api_key)

st.title("Anti debuffing")
col1, col2 = st.columns(2)

instruction = st.secrets["INSTRUCTION"]

if "messages" not in st.session_state:
    st.session_state.messages = []

with col1:
    st.header("Alison's screen")

    if prompt := st.chat_input("Alison's message", key="elison_prompt"):
        st.session_state.messages.append({"role": "Alison", "content": prompt})

if prompt:
    try:
        dialog = "".join([f"{message['role']} : {message['content']}\n" for message in st.session_state.messages])

        completion = client.chat.completions.create(
            model="gpt-4-turbo-preview",
            messages=[
                {"role": "system", "content": "You are a good psychologist"},
                {"role": "user", "content": instruction.format(DIALOG=dialog)}
            ]
        )
        reply = completion.choices[0].message.content
        print(f"Reply: {reply}")
        reply_suggestion = json.loads(reply)["possible_reply"]
        manipulations = json.loads(reply)["manipulations"]
    except Exception as e:
        logging.error(e, exc_info=True)
        reply_suggestion = "Hey! I don't understand:)"
        manipulations = None

with col2:
    st.header("Emily's screen")
    autoreply = st.checkbox('Autoreply', value=True, key='autoreply')

    if not autoreply:
        if response := st.chat_input("Emily's message",  key="emily_prompt"):
            st.session_state.messages.append({"role": "Emily", "content": response})
    elif prompt:
        response = reply_suggestion
        st.session_state.messages.append({"role": "Emily", "content": response})

    if prompt and manipulations:
        try:
            for m in manipulations:
                m_title = m["manipulation_short_name"]
                m_desc = m["explanation"]
                st.warning(f"Detected manipulation: {m_title}\n\n{m_desc}")
        except Exception as e:
            logging.error(e, exc_info=True)

        if not autoreply:
            with st.form(key="use_suggested_reply"):
                response = reply_suggestion
                st.session_state.reply_suggestion = response
                st.info(f"Suggested reply:\n\n{reply_suggestion}")
                submitted = st.form_submit_button("Use suggested reply", on_click=on_suggested_reply_click_handler)
        else:
            st.info(f"Suggested reply:\n\n{reply_suggestion}")

with st.sidebar:
    if st.button("New dialog"):
        st.session_state.messages = []
        st.session_state.reply_suggestion = None

    st.header("Active dialog")
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])