import streamlit as st
import chromadb
from chromadb.utils import embedding_functions
from sentence_transformers import SentenceTransformer
from openai import OpenAI

# CONSTANTS
client = chromadb.PersistentClient(path="./chromadb_linux/")
MODEL_NAME: str = "mixedbread-ai/mxbai-embed-large-v1"  # ~ 0.5 gb
COLLECTION_NAME: str = "scheme"
EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name=MODEL_NAME
)
schemer = client.get_collection(
    name=COLLECTION_NAME,
    embedding_function=EMBEDDING_FUNC,
)
DATA_AVAL: bool = schemer.count() > 0
APP_NAME: str = "Groove-GPT"
history = []

# INFO
st.title(APP_NAME)
st.header("What is Groovy-GPT?")
st.write(
    "Groovy-GPT is a RAG (Retrieval-Augmented Generation) model that uses ChromaDB to retrieve relevant documents and then uses OpenAI's models to generate a response."
)
st.write(
    "The model is trained on the MIT Scheme textbook and a handful of Discrete Math and Paradigms related content that Professor Troeger posted"
)
st.write("Data Avaliable: ", DATA_AVAL)

# INPUTS
user_question: str = st.text_area("Enter your groovy questions here")

remember_chat_history = st.toggle("Remember This Chat's History")

temperature = st.slider(
    label="Creativity of Model", min_value=0.0, max_value=2.0, value=0.8
)
st.markdown("*High creativity will make it go crazy - keep it low*")

num_samples = st.slider(
    label="Amount of References to Give to Model", min_value=10, max_value=100, value=10
)
st.markdown(
    "*High amount will make it slow and expensive (and may not be relevant) - keep it low*"
)

access_key: str = st.text_input("Enter your gpt key here", type="password")
st.markdown(
    "*For more information about how to get an access key, read [this article](https://platform.openai.com/api-keys). Make sure it has money in it ☠️*",
    unsafe_allow_html=True,
)

gpt_type: str = st.selectbox(
    label="Choose GPT Type",
    options=[
        "gpt-3.5-turbo",
        "gpt-3.5-turbo-1106",
        "gpt-3.5-turbo-0125",
        "gpt-4-32k-0613",
        "gpt-4-0613",
        "gpt-4-0125-preview",
    ],
    index=0,
)
st.markdown(
    "*For more information about GPT types, read [this article](https://platform.openai.com/docs/models).*",
    unsafe_allow_html=True,
)

st.divider()

# ON BUTTON CLICK
if st.button("Start Scheming") & (access_key != "") & (user_question != ""):
    openai_client = OpenAI(api_key=access_key)

    with st.spinner("Loading..."):
        # Perform the Chromadb query.
        results = schemer.query(
            query_texts=[user_question], n_results=num_samples, include=["documents"]
        )
        documents = results["documents"]
        response = openai_client.chat.completions.create(
            model="gpt-3.5-turbo",
            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$$. You will focus your examples to work exclusively in interative and recursive apporaches",
                },
                {"role": "user", "content": user_question},
                {"role": "assistant", "content": str(documents)},
                {"role": "user", "content": f"Conversation History: {history}"},
            ],
            temperature=temperature,
            stream=True,
        )

    # history.append({user_question : response.choices[0].message.content} if remember_chat_history else {})

    st.header("The Mega Schemer Says ...")

    text_placeholder = st.empty()

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

            text_placeholder.markdown(content)
else:
    st.write("Please provide an input and (valid) API key")