import os
import fitz  # PyMuPDF zum Extrahieren von PDF-Text
import marqo

# Verbindung zur Marqo-Instanz herstellen
client = marqo.Client(url="http://192.168.11.16:8882")  # Marqo-URL

# Marqo-Indexname
index_name = "pdf-index"

# Ordner mit PDFs
pdf_folder = "C:\Firma\KI\Local_RAG\dateien"

# Index-Einstellungen
index_settings = {
    "treatUrlsAndPointersAsImages": False,
    "model": "hf/e5-base-v2",  # Du kannst das Model nach Bedarf anpassen
}

# Funktion zum Extrahieren von Text aus einem PDF
def extract_text_from_pdf(pdf_path):
    text = ""
    try:
        doc = fitz.open(pdf_path)
        for page_num in range(doc.page_count):
            page = doc.load_page(page_num)
            text += page.get_text("text")
        doc.close()
    except Exception as e:
        print(f"Fehler beim Lesen der Datei {pdf_path}: {e}")
    return text

# Alle PDFs aus einem Ordner laden und in Marqo importieren
def import_pdfs_to_marqo(folder_path, index_name):
    for filename in os.listdir(folder_path):
        if filename.endswith(".pdf"):
            pdf_path = os.path.join(folder_path, filename)
            print(f"Verarbeite {filename}...")

            # Text aus PDF extrahieren
            pdf_text = extract_text_from_pdf(pdf_path)

            if pdf_text:
                # Dokument in Marqo importieren
                document = {
                    "pdf_filename": filename,
                    "pdf_text": pdf_text
                }
                client.index(index_name).add_documents(
                    [document],
                    tensor_fields=["pdf_text"]  # Hier wird das Textfeld für die Vektorisierung angegeben
                )
                print(f"{filename} erfolgreich in Marqo importiert.")

# Funktion zur Überprüfung, ob ein Index existiert
def index_exists(client, index_name):
    try:
        client.index(index_name).get_settings()  # Abfrage, ob der Index existiert
        return True
    except marqo.errors.IndexNotFoundError:
        return False

# Hauptprogramm
if __name__ == "__main__":
    # Prüfen, ob der Index existiert
    if not index_exists(client, index_name):
        print(f"Index {index_name} existiert nicht. Erstelle den Index...")
        client.create_index(index_name, settings_dict=index_settings)
        print(f"Index {index_name} wurde neu erstellt.")
    else:
        print(f"Index {index_name} existiert bereits. Importiere neue Dokumente...")

    # PDFs importieren
    import_pdfs_to_marqo(pdf_folder, index_name)
    print("Fertig!")