import streamlit as st
from openai import OpenAI
from params import params
from Chroma import Chroma

# from add_data import start_troggin_off, create_client

CLIENT = Chroma()
APP_NAME: str = "Groove-GPT"
history = []
st.set_page_config(layout="wide")

# INFO
st.title(APP_NAME)

l_col, r_col = st.columns((3, 1))

if "trigger" not in st.session_state:
    st.session_state["trigger"] = False


def on_enter():
    st.session_state["trigger"] = True


# param column
with r_col:
    (
        submit_button,
        remember_chat_history,
        temperature,
        num_samples,
        access_key,
        gpt_type,
    ) = params()

# input & response
with l_col:
    user_question: str = st.text_area(
        "Enter your groovy questions here",
        on_change=on_enter,
    )

    # ON BUTTON CLICK
    if (
        (submit_button | st.session_state["trigger"])
        & (access_key != "")
        & (user_question != "")
    ):
        openai_client = OpenAI(api_key=access_key)

        with st.spinner("Loading..."):
            collection = CLIENT.get_collection()
            results = collection.query(
                query_texts=[user_question],
                n_results=num_samples,
                include=["documents"],
            )
            documents = results["documents"]
            response = openai_client.chat.completions.create(
                model=gpt_type,
                messages=[
                    {
                        "role": "system",
                        "content": "You are an expert in functional programming in R5RS, with great knowledge on programming paradigms. You wish to teach the user everything you know about programming paradigms in R5RS - so you explain everything thoroughly. Surround Latex equations in dollar signs as such Inline equation: $equation$ & Display equation: $$equation$$.",
                    },
                    {"role": "user", "content": user_question},
                    {"role": "assistant", "content": str(documents)},
                    {"role": "user", "content": f"Conversation History: {history}"},
                ],
                temperature=temperature,
                stream=True,
            )

        st.header("The Super Duper Schemer Says ...")
        text_placeholder = st.empty()

        content = ""
        for i, chunk in enumerate(response):
            if chunk.choices[0].delta.content is not None:
                content += chunk.choices[0].delta.content

                text_placeholder.markdown(content)

        history.append({user_question: content} if remember_chat_history else {})
    else:
        st.write("Please provide an input and (valid) API key")