from langchain.document_loaders.unstructured import UnstructuredFileLoader 
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.schema import AIMessage, HumanMessage, SystemMessage, Document
import os


API_KEY = os.getenv("API_TOKEN")
print("API:", API_KEY)

class Agent:
    def __init__(self, args=None) -> None:
        self.embeddings = OpenAIEmbeddings(openai_api_key=API_KEY)
        self.llm = ChatOpenAI(temperature=0.5, openai_api_key=API_KEY)
        self.context_value = ""
        self.use_context = False

    def load_context(self, doc_path):
        loader = UnstructuredFileLoader(doc_path.name)
        print('Loading file:', doc_path.name)
        self.documents = loader.load()
        self.split()

        return f"Using file from {doc_path.name}"
    
    def split(self):
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
        texts = text_splitter.split_documents(self.documents)
        docsearch  = Chroma.from_documents(texts, self.embeddings)
        self.qa = RetrievalQA.from_chain_type(llm=self.llm, chain_type="stuff", retriever=docsearch.as_retriever())

        print("Context updated")
        self.use_context = True

    def asking(self, prompt):
        if self.use_context:
            print("Answering with your context")
            return self.qa.run(prompt)
        else:
            print("Answering without your context")
            return self.llm([HumanMessage(content=prompt)]).content

    
    def get_context(self, context):
        self.context_value = context

        self.documents = [Document(page_content=context, metadata={'source': ''})]
        self.split()

        self.use_context = True
    
    def load_context_file(self, file):
        print('Loading file:', file.name)
        text = ''
        for line in open(file.name, 'r', encoding='utf8'):
            text += line

        self.context_value = text
        return text
    
    def clear_context(self):
        self.context_value = ""
        self.use_context = False
        return ""