from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
import os
import openai
import pandas as pd
import numpy as np

# Set up OpenAI API key
openai.api_key = os.getenv("OPENAI_API_KEY")

# Load data
persist_directory = 'trading_psychology_db'
embedding = OpenAIEmbeddings()
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)

pre_text_trading_psychology = """
You are a helpful AI assistant that helps people with their trading psychology.
You can use the provided context as knowledge source.
Context is generated from a trading psychology podcast called 'Chat with Traders' Hosted by 'Aaron Fifield'.
You answer the question by using the context and your own knowledge. Try to give definitive answers.
Do not mention that the answer is based on the context.
"""

def search(query):
    similar_docs = vectordb.similarity_search(query=query, k=10)
    similar_texts = [doc.page_content for doc in similar_docs]
    context = "\n\n".join(similar_texts)
    return context

def get_context(prompt):
    # Get the context based on the prompt
    context = search(prompt) 

    # Concatenate the prompt and context
    formatted_prompt = f"""
    {pre_text_trading_psychology}

    User Question: {prompt}

    Context: ```{context}```

    Your answer:
    """
    return formatted_prompt

def get_reply(message, messages_archived, messages_current):
    
    if message:
        messages_current = messages_archived.copy()
        context = get_context(message)
        messages_current.append(
            {"role": "user", "content": context}
        )
        chat = openai.ChatCompletion.create(
            model="gpt-3.5-turbo", messages=messages_current, temperature=0
        )
    
        reply = chat.choices[0].message.content
        messages_archived.append({"role": "user", "content": message})
        messages_archived.append({"role": "assistant", "content": reply})
    # If no message is provided, return a string that says "No Message Received"
    else:
        reply = "No Message Received"

    return reply, messages_archived, messages_current