import sys
import os
lib_dir = os.path.join(os.path.dirname(__file__), 'lib')
if lib_dir not in sys.path:
    sys.path.insert(0, lib_dir)
import requests
import json
import base64
import pathlib
import re
from datetime import datetime
import time

class AnythingLLM_API():
    
    def __init__(self, API_URL, API_KEY, WORKSPACE, CONTEXT_MODE='VectorDB', progress_reporter=False):
        # class specific variables
        self.progress_reporter = progress_reporter
        self.log_file = 'AI Metadata - Job.log'
        self.API_URL = API_URL
        self.API_KEY = API_KEY
        self.WORKSPACE = WORKSPACE
        self.CONTEXT_MODE = CONTEXT_MODE
        self.API_HEADERS = {"accept": "application/json", "Authorization": "Bearer "+ self.API_KEY}
        self.UPLOAD_FOLDER = 'Calibre'
        self.instanciate_logfile()      
        self.WORKSPACE_SLUG = ''
        self.CURRENT_THREAD_SLUG = ''
        self.THREAD_SLUGS_PROCESSED = []
        self.DOC_UPLOAD_ID = ''
        self.DOC_UPLOAD_URL = ''
        self.DOC_UPLOAD_LOCATION = ''
        self.abort = False
        self.DEBUG_RESPONSE = False
        
    ##############################
    def abort_process(self, status=True):
        self.abort = status    
    def stay_alive():
        if self.abort == False:
            return True            
    def debug_reponse(self, status=True):
        self.DEBUG_RESPONSE = status
    ###########################
    
    def instanciate_logfile(self):
        from calibre.constants import CONFIG_DIR_MODE, config_dir
        self.log_file = os.path.join(config_dir, 'plugins') + os.sep + self.log_file
        print("LOGFILE: " + self.log_file)
        open(self.log_file, "w")
        
    def filter_log_message(self,message):
        if isinstance(message, bytes):
            try:
                message = message.decode()
            except (UnicodeDecodeError, AttributeError):
                pass
        if isinstance(message, dict) or isinstance(message, list):
            message = str(", ".join(message))
        return message
    
    def log_append(self, message, error=False):
        message = self.filter_log_message(message)
        print(message)
        if error == False and self.abort == False:
            if self.progress_reporter != False:
                message = (message[:75] + '..') if len(message) > 75 else message
                self.progress_reporter(message)
        with open(self.log_file, 'a') as file:
            file.write(message + "\n")
                
    def strip_filename_from_path(self, file_path):
        file_stem = pathlib.Path(file_path).stem
        extension = pathlib.Path(file_path).suffix
        file_name = file_stem + extension
        return file_name
    
    def clean_slug(self, s):
        # Remove all non-word characters (everything except numbers and letters)
        s = re.sub(r"[^\w\s]", '', s)
        # Replace all runs of whitespace with a single dash
        s = re.sub(r"\s+", '-', s)
        return s.lower()
        
    def clean_name(self, s):
        # replace dots with spaces
        s = s.replace(".", " ")
        # Remove all non-word characters (everything except numbers and letters)
        s = re.sub(r"[^\w\s]", '', s)
        s = s.replace("  ", " ")
        return s
        
    def filter_slug(self, name):
        return clean_slug(name)

    def inline_encode_file(self, file_location):
        # for openwebui
        extension = pathlib.Path(file_location).suffix
        file_name = pathlib.Path(file_location).stem + extension
        
        if extension == ".pdf":
            mime_type = 'application/pdf'
            base_64_string = 'data:'+ mime_type +';base64,' + base64_encode_file(file_location)
        elif extension == ".epub":
            mime_type = 'application/epub+zip'
            base_64_string = 'data:'+ mime_type +';base64,' + base64_encode_file(file_location)
        
        attatchment = {"name": file_name, "mime": mime_type, "contentString": base_64_string}
        
        return attatchment

    def base64_encode_file(self, file_location):

        with open(file_location, "rb") as file:
            file_read = file.read()
            result = base64.b64encode(file_read).decode('ascii')
        
        return result
        
    def authenticate(self):
        self.log_append("Authenticating..")
        URL = self.API_URL + '/auth'
        PAYLOAD = {}
        try:
            response = requests.get(URL, data=PAYLOAD, headers=self.API_HEADERS)
            if response.status_code != 200:
                self.log_append("Couldn't Authenticate! - Bad Request")
                return False
            else:
                json = response.json()
                if json['authenticated'] == True:
                    self.log_append("Authenticated!")
                    return True
                else:
                    self.log_append("Couldn't Authenticate! - Are your credentials correct?")
                    return False
        except Exception as e:
            self.log_append("Couldn't Authenticate!")
            return False

    def create_workspace(self):
        if self.abort == False:
            try:
                self.log_append("Creating Workspace...")
                URL = self.API_URL + '/workspace/new'
                PAYLOAD = {}
                PAYLOAD["name"] = self.WORKSPACE
                PAYLOAD["similarityThreshold"] = 0.25
                PAYLOAD["openAiTemp"] = 0.0
                PAYLOAD["openAiHistory"] = 1
                PAYLOAD["openAiPrompt"] = "Given the following conversation, relevant context, and a follow up question, reply with an answer to the current question the user is asking. Return only your response to the question given the above information following the users instructions as needed."
                PAYLOAD["queryRefusalResponse"] = "```json{}```"
                PAYLOAD["chatMode"] = "query"
                PAYLOAD["topN"] = 5 ## number of context snippets to use from vectorDB
                PAYLOAD["vectorSearchMode"] = "rerank"
                #print(URL)
                response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS)
                
                ## debugging
                if self.DEBUG_RESPONSE == True:
                    self.log_append( str(response.status_code) )
                    self.log_append(response.text)
                
                if response.status_code == 200:
                    return self.check_workspace_exists();
                else:
                    return False
            except Exception as e:
                self.log_append(f"Error Creating Workspace: {e}")
                return False
            
    def check_workspace_exists(self):
        if self.abort == False:
            URL = self.API_URL + '/workspaces'
            PAYLOAD = {}
            #print(URL)
            self.log_append("Checking Workspace...")
            try:
                response = requests.get(URL, data=PAYLOAD, headers=self.API_HEADERS)
                
                ## debugging
                if self.DEBUG_RESPONSE == True:
                    self.log_append( str(response.status_code) )
                    self.log_append(response.text)
                
                if response.status_code == 200:
                    workspaces = json.loads(response.content)['workspaces']
                    WORKSPACE_EXISTS = False
                    for workspace in workspaces:
                        if workspace['name'] == self.WORKSPACE:
                            WORKSPACE_EXISTS = True
                            self.WORKSPACE_SLUG = workspace['slug']
                            self.log_append('Workspace: ' + workspace['name'] + ' Found')
                            return True
                    if WORKSPACE_EXISTS != True:
                        self.log_append('Workspace Not Found Creating:' + self.WORKSPACE)
                        try:
                            return self.create_workspace()
                        except Exception as e:
                            self.log_append(f"Error creating Workspace: {e}")
                            return False
                else:
                    return False                
            except Exception as e:
                self.log_append(f"Error Connecting to AI Server: {e} Have you started the server ?")
                return False
            
    def create_thread(self, thread_name):
        if self.abort == False:
            PAYLOAD = {}
            PAYLOAD['name'] = self.clean_name(thread_name)
            URL = self.API_URL + '/workspace/'+ self.WORKSPACE_SLUG +'/thread/new'        
            self.log_append("Creating New Chat: "+ PAYLOAD['name'])        
            response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS)
            
            ## debugging
            if self.DEBUG_RESPONSE == True:
                self.log_append( str(response.status_code) )
                self.log_append(response.text)
            
            if response.status_code == 200:
                threads = json.loads(response.content)['thread']
                self.CURRENT_THREAD_SLUG = threads['slug']
                self.log_append("Chat Created!")
                return True
            else:
                self.log_append("ERROR: Failed to create Chat!")
                return False
            
    def check_document_exists(self, ebook_location):
        if self.abort == False:
            filename = self.strip_filename_from_path(ebook_location)
            self.log_append("Checking if Document has already been Uploaded..")
            URL = self.API_URL + '/documents/folder/' + self.UPLOAD_FOLDER
            PAYLOAD = {}
            response = requests.get(URL, data=PAYLOAD, headers=self.API_HEADERS)
            
            ## debugging
            if self.DEBUG_RESPONSE == True:
                self.log_append( str(response.status_code) )
                self.log_append(response.text)
            
            if response.status_code == 200:
                documents = json.loads(response.content)['documents']
                for document in documents:
                    if document['title'] == filename:
                        self.log_append("Document Already Exists!")                                                
                        self.DOC_UPLOAD_ID = document['id']
                        self.DOC_UPLOAD_LOCATION = self.UPLOAD_FOLDER + "/" + document['name']
                        return True
                        
            self.log_append("Document Doesn't Exist!")
            return False
        
    def embed_document_in_workspace(self):
        if self.abort == False:
            self.log_append("Embedding Document into Workspace..")
            URL = self.API_URL + '/workspace/'+ self.WORKSPACE_SLUG +'/update-embeddings'
            PAYLOAD = {}
            PAYLOAD["adds"] = [self.DOC_UPLOAD_LOCATION]
            response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS)
            
            ## debugging
            if self.DEBUG_RESPONSE == True:
                self.log_append( str(response.status_code) )
                self.log_append(response.text)
            
            if response.status_code == 200:
                self.log_append("Document Embedded!")
                return True
            else:
                self.log_append("ERROR: Failed to Embed Document!")
                return False
    
    def pin_document_to_workspace(self):
        if self.abort == False:
            self.log_append("Pinning Document into Workspace..")
            URL = self.API_URL + '/workspace/'+ self.WORKSPACE_SLUG +'/update-pin'
            PAYLOAD = {}
            PAYLOAD["docPath"] = f"{self.DOC_UPLOAD_LOCATION}"
            PAYLOAD["pinStatus"] = True
            response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS)            
            ## debugging
            if self.DEBUG_RESPONSE == True:
                self.log_append( str(response.status_code) )
                self.log_append(response.text)
            
            if response.status_code == 200:
                self.log_append("Document Pinned!")
            else:
                self.log_append("ERROR: Failed to pin Document!")
    
    def embed_and_pin_document(self):
        if self.embed_document_in_workspace():
            if self.CONTEXT_MODE == "FullContext":
                ## pin document
                self.pin_document_to_workspace()
    
    def upload_document_to_workspace(self, ebook_location):
        if self.abort == False:
            if self.check_document_exists(ebook_location) != True:
                self.log_append("Uploading Document to Workspace..")
                PAYLOAD = {}
                PAYLOAD["addToWorkspaces"] = self.WORKSPACE_SLUG
                URL = self.API_URL + '/document/upload/'+ self.UPLOAD_FOLDER
                file_obj = ''
                try:
                    with open(ebook_location, 'rb') as file_obj:
                        response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS, files={'file': file_obj})
                        
                        ## debugging
                        if self.DEBUG_RESPONSE == True:
                            self.log_append( str(response.status_code) )
                            self.log_append(response.text)
                        
                        if response.status_code == 200:
                            self.log_append("Document Uploaded Successfully!")
                            document = json.loads(response.content)['documents'][0]                            
                            self.DOC_UPLOAD_ID = document['id']
                            self.DOC_UPLOAD_URL = document['url']
                            self.DOC_UPLOAD_LOCATION = document['location'].replace("\\","/")
                            ## update embeddings
                            if self.abort == False: 
                                self.embed_and_pin_document()
                                ## report status                                
                                return True
                        else:
                            self.log_append("ERROR: Failed to Upload Document!: " + str(response.content))
                            return False
                except:
                    self.log_append("ERROR: Try failed to Open Document!",True)
                    return False
            else:
                self.embed_and_pin_document()
                return True

    def delete_thread(self,thread_slug=None):        
        if thread_slug != None:            
            URL = self.API_URL + '/workspace/'+ self.WORKSPACE_SLUG +'/thread/' + thread_slug
            PAYLOAD = {}
            response = requests.delete(URL, data=PAYLOAD, headers=self.API_HEADERS)
            
            ## debugging
            if self.DEBUG_RESPONSE == True:
                self.log_append( str(response.status_code) )
                self.log_append(response.text)
            
            if response.status_code == 200:
                self.log_append("Thread Deleted Sucessfully!",True)
            else:
                self.log_append("ERROR: Failed to Delete Thread!",True)
    
    def delete_processed_threads(self):
        self.log_append("Deleting Threads...", True)
        for thread_slug in self.THREAD_SLUGS_PROCESSED:
            self.delete_thread(thread_slug)
        self.log_append("Threads Deleted Sucessfully!", True)
        self.THREAD_SLUGS_PROCESSED = []
    
    def delete_embedding(self):
        self.log_append("Deleting Embedding..", True)
        URL = self.API_URL + '/workspace/'+ self.WORKSPACE_SLUG +'/update-embeddings'
        PAYLOAD = {}
        PAYLOAD['deletes'] = [self.DOC_UPLOAD_LOCATION]
        response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS)
        
        ## debugging
        if self.DEBUG_RESPONSE == True:
            self.log_append( str(response.status_code) )
            self.log_append(response.text)
        
        if response.status_code == 200:
            self.log_append("Document Embedding Removed!", True)
        else:
            self.log_append("ERROR: Failed to remove Embedding!", True)
            
    def delete_pin(self):
        self.log_append("Deleting Pin..", True)
        URL = self.API_URL + '/workspace/'+ self.WORKSPACE_SLUG +'/update-pin'
        PAYLOAD = {}
        PAYLOAD["docPath"] = [self.DOC_UPLOAD_LOCATION]
        PAYLOAD["pinStatus"] = False
        response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS)
        
        ## debugging
        if self.DEBUG_RESPONSE == True:
            self.log_append( str(response.status_code) )
            self.log_append(response.text)
        
        if response.status_code == 200:
            self.log_append("Document Pin Removed!", True)
        else:
            self.log_append("ERROR: Failed to remove Pin!", True)
            
    def delete_document_folder(self):
        self.log_append("Deleting Document Folder..", True)
        URL = self.API_URL + '/document/remove-folder'
        PAYLOAD = {}
        PAYLOAD['name'] = "Calibre"
        response = requests.delete(URL, json=PAYLOAD, headers=self.API_HEADERS)
        
        ## debugging
        if self.DEBUG_RESPONSE == True:
            self.log_append( str(response.status_code) )
            self.log_append(response.text)
        
        if response.status_code == 200:
            self.log_append("Documents Deleted Successfully!", True)
        else:
            self.log_append("ERROR: Failed to Delete Document Folder!", True)
    
    def delete_document_from_system(self):
        self.log_append("Deleting Individual Document..", True)
        # this removes the documents
        PAYLOAD = {}
        PAYLOAD["names"] = [f"{self.DOC_UPLOAD_LOCATION}"]
        URL = self.API_URL + '/system/remove-documents'
        response = requests.delete(URL, json=PAYLOAD, headers=self.API_HEADERS)
        
        ## debugging
        if self.DEBUG_RESPONSE == True:
            self.log_append( str(response.status_code) )
            self.log_append(response.text)
        
        if response.status_code == 200:
            self.log_append("Document Deleted Successfully!", True)
        else:
            self.log_append("ERROR: Failed to Delete Document!", True)
        
    def cleanup(self):
        self.log_append("Starting Cleanup Process...")
        if self.CONTEXT_MODE == "FullContext":
            self.delete_pin()
        self.delete_embedding()
        self.delete_document_from_system()
        self.delete_document_folder()
        self.delete_processed_threads()
        self.log_append("Cleanup Process Complete!")
    
    def decode_bytes_object(self, bstring):
        if isinstance(bstring, bytes):
            try:
                bstring = bstring.decode()
            except (UnicodeDecodeError, AttributeError):
                pass
        return bstring
    
    def send_ai_question(self, PAYLOAD, file_name, question):
        thread_name = file_name + " " + question
        if self.create_thread(thread_name):
            self.log_append("Asking A.I. Question..")            
            if question != "":
                self.log_append(question)
            self.log_append("A.I. Thinking.....")
            # we time the thinking process
            start_time = time.time()
            URL = self.API_URL + '/workspace/'+ self.WORKSPACE_SLUG +'/thread/' + self.CURRENT_THREAD_SLUG + '/chat'
            response = requests.post(URL, json=PAYLOAD, headers=self.API_HEADERS)
            
            ## debugging
            if self.DEBUG_RESPONSE == True:
                self.log_append( str(response.status_code) )
                self.log_append(response.text)
            
            if response.status_code == 200:
                time_taken = round(time.time() - start_time, 2)
                returned_payload = response.json()
                returned_payload = json.dumps(returned_payload['textResponse'], ensure_ascii=False).replace("\\r\\n"," ").replace("\\n"," ").replace('\\"', '"').replace('\\', '/').replace('/ / ', ' ')
                
                ## debugging
                if self.DEBUG_RESPONSE == True:
                    self.log_append("Returned Payload:")
                    self.log_append( str(returned_payload) )

                if returned_payload == "There is no relevant information in this workspace to answer your query.":
                    self.log_append("AI Returned: There is no relevant information in this workspace to answer your query.")
                    return json.loads("{}")
                else:
                    thinking = re.search(r"(?<=<think>).*?(?=<\/think>)", returned_payload)
                    if thinking:
                        thinking = thinking.group(0)
                        self.log_append("------------------------------------------------------------",True)
                        self.log_append("############             Thinking:              ############",True)
                        self.log_append("------------------------------------------------------------",True)
                        self.log_append(thinking,True)
                        self.log_append("------------------------------------------------------------",True)
                    
                    
                    pljson = re.search(r"(?<=```json).*?(?=```)", returned_payload)
                    
                    ## debugging
                    if self.DEBUG_RESPONSE == True:
                        self.log_append("Payload JSON:")
                        self.log_append( str(pljson) )
                    
                    if pljson != None:
                        pljson = pljson.group(0)
                        self.log_append("------------------------------------------------------------",True)
                        self.log_append("############             Response:              ############",True)
                        self.log_append("------------------------------------------------------------",True)
                        self.log_append(pljson,True)
                        self.log_append("------------------------------------------------------------",True)
                        returned_dict = json.loads(pljson)
                        returned_dict =  {k.lower().replace(" ", "_"): v for k, v in returned_dict.items()}
                        ## got a response so update the log
                        self.log_append("Response Received Successfully!")
                        self.log_append(f"A.I thinking took " + str(time_taken) +" seconds!")
                    else:
                        
                        #NO JSON CODE WRAPPER FOUND, TRY WITH JSON PATTERN REGEX CAPTURE
                        #\s*\{[\s\S]*?\}\s*$
                        
                        ## debugging
                        if self.DEBUG_RESPONSE == True:
                            self.log_append("Thinking Variable:")
                            self.log_append( str(thinking) )
                        
                        # strip the thinking from the payload incase it contains json
                        if thinking != None:
                            returned_payload = returned_payload.replace(thinking,"").replace("<think>","").replace("</think>","").strip()                        
                            ## debugging
                            if self.DEBUG_RESPONSE == True:
                                self.log_append("Payload Minus Thinking:")
                                self.log_append( str(returned_payload) )
                        
                        pljson = re.search(r"{[\w\W]+?}", returned_payload)
                        ## debugging
                        if self.DEBUG_RESPONSE == True:
                            self.log_append("Payload Regex Found:")
                            self.log_append( str(pljson) )
                            
                        if pljson != None:                            
                            pljson = pljson.group(0)
                            self.log_append("------------------------------------------------------------",True)
                            self.log_append("############             Response:              ############",True)
                            self.log_append("------------------------------------------------------------",True)
                            self.log_append(pljson,True)
                            self.log_append("------------------------------------------------------------",True)
                            returned_dict = json.loads(pljson)
                            returned_dict =  {k.lower().replace(" ", "_"): v for k, v in returned_dict.items()}
                            ## got a response so update the log
                            self.log_append("Response Received Successfully!")
                            self.log_append(f"A.I thinking took " + str(time_taken) +" seconds!")
                        else:
                            self.log_append("ERROR: Failed to get Parse Response from the A.I.! Possibly No Context Found!")
                            self.log_append(returned_payload)
                            self.log_append("ERROR: Failed to get Parse Response from the A.I.! Possibly No Context Found!", True)
                            self.log_append(returned_payload,True)
                            return
                
                self.THREAD_SLUGS_PROCESSED.append(self.CURRENT_THREAD_SLUG)
                
                return returned_dict
            else:
                self.log_append("ERROR: Failed to get a Response from the A.I.! Response Code:" + str(response.status_code) + "Response:" + str(response.text) )
                return
        else:
            return