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 Custom Tag Extractor - 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
        
    ##############################
    def abort_process(self, status=True):
        self.abort = status
    
    def stay_alive():
        if self.abort == False:
            return True
    ###########################
    
    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)
                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)
                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)
            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)
            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!")
                        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)
            if response.status_code == 200:
                self.log_append("Document Embedded!")
            else:
                self.log_append("ERROR: Failed to Embed Document!")
    
    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)
            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):
        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})
                        if response.status_code == 200:
                            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
                                self.log_append("Document Uploaded Successfully!")
                                return True
                        else:
                            self.log_append("ERROR: Failed to Upload Document!")
                            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)
            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)
        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)
        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)
        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)
        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)
            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('\\\\', ' ')
                #print(returned_payload)
                if returned_payload == "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)
                    if pljson:
                        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,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.content) )
                return
        else:
            return
        
