import streamlit as st

try:
    import dotenv
    dotenv.load_dotenv()
except ImportError:
    pass

import openai
import os
import streamlit.components.v1 as components
import requests


openai.api_key = os.getenv("OPENAI_API_KEY")
embedbase_api_key = os.getenv("EMBEDBASE_API_KEY")

URL = "https://api.embedbase.xyz"
local_history = []


def add_to_dataset(dataset_id: str, data: str):
    response = requests.post(
        f"{URL}/v1/{dataset_id}",
        headers={
            "Content-Type": "application/json",
            "Authorization": "Bearer " + embedbase_api_key,
        },
        json={
            "documents": [
                {
                    "data": data,
                },
            ],
        },
    )
    response.raise_for_status()
    return response.json()


def search_dataset(dataset_id: str, query: str, limit: int = 3):
    response = requests.post(
        f"{URL}/v1/{dataset_id}/search",
        headers={
            "Content-Type": "application/json",
            "Authorization": "Bearer " + embedbase_api_key,
        },
        json={
            "query": query,
            "top_k": limit,
        },
    )
    response.raise_for_status()
    return response.json()


def chat(user_input: str, conversation_name: str) -> str:
    local_history.append(user_input)

    history = search_dataset(
        f"infinite-pt-{conversation_name}",
        # searching using last 4 messages from local history
        "\n\n---\n\n".join(local_history[-4:]),
        limit=3,
    )
    print("history", history)
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant.",
            },
            *[
                {
                    "role": "assistant",
                    "content": h["data"],
                }
                for h in history["similarities"][-5:]
            ],
            {"role": "user", "content": user_input},
        ],
    )
    message = response.choices[0]["message"]
    add_to_dataset(f"infinite-pt-{conversation_name}", message["content"])

    local_history.append(message)

    return message["content"]


from datetime import datetime

# conversation name is date like ddmmyy_hhmmss
# conversation_name = datetime.now().strftime("%d%m%y_%H%M%S")
conversation_name = st.text_input("Conversation name", "purpose")

# eg not local dev
if not os.getenv("OPENAI_API_KEY"):
    embedbase_api_key = st.text_input(
        "Your Embedbase key", "get it here <https://app.embedbase.xyz/signup>"
    )
    openai_key = st.text_input(
        "Your OpenAI key", "get it here <https://platform.openai.com/account/api-keys>"
    )
    openai.api_key = openai_key
user_input = st.text_input("You", "How can I reach maximum happiness this year?")
if st.button("Send"):
    infinite_pt_response = chat(user_input, conversation_name)
    st.markdown(
        f"""
        Infinite-PT
        """
    )
    st.write(infinite_pt_response)

components.html(
    """
<script>
const doc = window.parent.document;
buttons = Array.from(doc.querySelectorAll('button[kind=primary]'));
const send = buttons.find(el => el.innerText === 'Send');
doc.addEventListener('keydown', function(e) {
    switch (e.keyCode) {
        case 13:
            send.click();
            break;
    }
});
</script>
""",
    height=0,
    width=0,
)


st.markdown(
    """
    [Source code](https://huggingface.co./spaces/louis030195/infinite-memory-chatgpt)
    """
)

st.markdown(
    """
    Built with ❤️ by [louis030195](https://louis030195.com).
    """
)

st.markdown(
    """
    Powered by [Embedbase](https://embedbase.xyz).
    """
)