from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import DirectoryLoader, TextLoader
import bibtexparser
import langchain
import os
import glob
from dotenv import load_dotenv
import openai
import constants
import time

# Set OpenAI API Key
load_dotenv()
os.environ["OPENAI_API_KEY"] = constants.APIKEY
openai.api_key = constants.APIKEY 

# Set paths
source_path = './data/src/'
destination_file = './data/ingested.txt'
store_path = './vectorstore/'
bibtex_file_path = '/home/wouter/Tools/Zotero/bibtex/library.bib'

# Load documents
print("===Loading documents===")
text_loader_kwargs={'autodetect_encoding': True}
loader = DirectoryLoader(source_path,
                         show_progress=True,
                         use_multithreading=True,
                         loader_cls=TextLoader,
                         loader_kwargs=text_loader_kwargs)
documents = loader.load()

if len(documents) == 0:
    print("No new documents found")
    quit()
# Add metadata based in bibliographic information
print("===Adding metadata===")

# Read the BibTeX file
with open(bibtex_file_path) as bibtex_file:
    bib_database = bibtexparser.load(bibtex_file)

# Get a list of all text file names in the directory
text_file_names = os.listdir(source_path)
metadata_store = []

# Go through each entry in the BibTeX file
for entry in bib_database.entries:
    # Check if the 'file' key exists in the entry
    if 'file' in entry:
        # Extract the file name from the 'file' field and remove the extension
        pdf_file_name = os.path.basename(entry['file']).replace('.pdf', '')

         # Check if there is a text file with the same name
        if f'{pdf_file_name}.txt' in text_file_names:
            # If a match is found, append the metadata to the list
            metadata_store.append(entry)

for document in documents:
    for entry in metadata_store:
        doc_name = os.path.basename(document.metadata['source']).replace('.txt', '')
        ent_name = os.path.basename(entry['file']).replace('.pdf', '')
        if doc_name == ent_name:
            document.metadata.update(entry)

# Initialize text splitter
print("===Splitting documents into chunks===")
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 1500,
    chunk_overlap  = 150,
    length_function = len,
    add_start_index = True,
)

split_documents = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings(
    show_progress_bar=True,
    request_timeout=60,
)

print("===Embedding text and creating database===")
new_db = FAISS.from_documents(split_documents, embeddings)

print("===Merging new and old database===")
old_db = FAISS.load_local(store_path, embeddings)
old_db.merge_from(new_db)
old_db.save_local(store_path, "index")

# Record the files that we have added
print("===Recording ingested files===")
with open(destination_file, 'a') as f:
    for document in documents:
        f.write(os.path.basename(document.metadata['source']))
        f.write('\n')

langchain.document_loaders.DirectoryLoader langchain.vectorstores.FAISS.load_local langchain.text_splitter.RecursiveCharacterTextSplitter langchain.vectorstores.FAISS.from_documents langchain.embeddings.OpenAIEmbeddings
[((380, 393), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (391, 393), False, 'from dotenv import load_dotenv\n'), ((764, 898), 'langchain.document_loaders.DirectoryLoader', 'DirectoryLoader', (['source_path'], {'show_progress': '(True)', 'use_multithreading': '(True)', 'loader_cls': 'TextLoader', 'loader_kwargs': 'text_loader_kwargs'}), '(source_path, show_progress=True, use_multithreading=True,\n    loader_cls=TextLoader, loader_kwargs=text_loader_kwargs)\n', (779, 898), False, 'from langchain.document_loaders import DirectoryLoader, TextLoader\n'), ((1364, 1387), 'os.listdir', 'os.listdir', (['source_path'], {}), '(source_path)\n', (1374, 1387), False, 'import os\n'), ((2324, 2437), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(1500)', 'chunk_overlap': '(150)', 'length_function': 'len', 'add_start_index': '(True)'}), '(chunk_size=1500, chunk_overlap=150,\n    length_function=len, add_start_index=True)\n', (2354, 2437), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((2536, 2596), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'show_progress_bar': '(True)', 'request_timeout': '(60)'}), '(show_progress_bar=True, request_timeout=60)\n', (2552, 2596), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((2670, 2719), 'langchain.vectorstores.FAISS.from_documents', 'FAISS.from_documents', (['split_documents', 'embeddings'], {}), '(split_documents, embeddings)\n', (2690, 2719), False, 'from langchain.vectorstores import FAISS\n'), ((2774, 2814), 'langchain.vectorstores.FAISS.load_local', 'FAISS.load_local', (['store_path', 'embeddings'], {}), '(store_path, embeddings)\n', (2790, 2814), False, 'from langchain.vectorstores import FAISS\n'), ((1261, 1291), 'bibtexparser.load', 'bibtexparser.load', (['bibtex_file'], {}), '(bibtex_file)\n', (1278, 1291), False, 'import bibtexparser\n'), ((3045, 3090), 'os.path.basename', 'os.path.basename', (["document.metadata['source']"], {}), "(document.metadata['source'])\n", (3061, 3090), False, 'import os\n'), ((1664, 1695), 'os.path.basename', 'os.path.basename', (["entry['file']"], {}), "(entry['file'])\n", (1680, 1695), False, 'import os\n'), ((2019, 2064), 'os.path.basename', 'os.path.basename', (["document.metadata['source']"], {}), "(document.metadata['source'])\n", (2035, 2064), False, 'import os\n'), ((2104, 2135), 'os.path.basename', 'os.path.basename', (["entry['file']"], {}), "(entry['file'])\n", (2120, 2135), False, 'import os\n')]