import os
import streamlit as st
from langchain_openai import OpenAI
from langchain.chains import RetrievalQA
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, AIMessage

def run_query(query, chat_history, k, temperature,openai_api_key):
    load_dotenv()
    persist_directory = 'db'

    db = Chroma(persist_directory=persist_directory, embedding_function=OpenAIEmbeddings(openai_api_key=openai_api_key))  # access db
    retriever = db.as_retriever(search_kwargs={"k": k})  # kwargs determines how many docs it uses

    llm = OpenAI(api_key=openai_api_key, max_tokens=1500, temperature=temperature)  # api key self explanatory. max_tokens provides how long of a response 
    # we get from the llm (do note the llm has a cap of 4097.) and temperature provides a scale form 0.0 to 1.0 of how much freedom 
    # the llm should take in its response (how closely it should adhere to docs vs how freely)

    # Perform similarity search
    search_results = retriever.get_relevant_documents(query)

    # Extract texts from the retrieved documents
    context = "\n".join([doc.page_content for doc in search_results])
    
    # Combine the context with the query
    full_query = f"{context}\n\n{query}"
    
    # Get response from LLM
    llm_response = llm(full_query)

    # Store the interaction in chat history
    chat_history.append((HumanMessage(content=query), AIMessage(content=llm_response)))
    
    return llm_response

def main():
    st.title("EmeraldEnergy™️ AI Copilot")
    api_key = st.text_input("Enter your OpenAI API key:")
    chat_history = []

    role = 'Pretend I am an HVAC technician installing a Mitsubishi heat pump. I do not have access to the manual. Help me install this using your knowledge of the mitsubishi heatpump. Please consider the section under instruction of troubleshooting. If you can not help me, then provide generalized guidance. Please ask for which model of the product we are using at the end. \n'

    query = st.text_area("How can I help you today?:", height=100)

    k = st.slider("Number of documents to retrieve (k):", min_value=1, max_value=4, value=4)

    temperature = st.slider("Temperature:", min_value=0.0, max_value=1.0, value=0.2)

    if st.button("Submit"):
        full_query = role + query
        result = run_query(full_query, chat_history, k, temperature,api_key)
        st.write(f"AI: {result}")

    # Display chat history
    st.write("### Chat History")
    for human_msg, ai_msg in chat_history:
        st.write(f"**Human:** {query}")
        st.write(f"**AI:** {ai_msg.content}")

if __name__ == "__main__":
    main()