import ollama
import pymupdf
import os
import lib
import json
import psycopg2



def retrieve_file_contents(path):
    file_extraction_functions = {
        "pdf": lambda path: lib.extract_pdf_content(path),
        "jpg": lambda path: lib.extract_image_content(path),
        "png": lambda path: lib.extract_image_content(path),
        "txt": lambda path: lib.extract_pdf_content(path),
        "mp3": lambda path: lib.extract_mp3_content(path),
    }

    lib.read_files(path, files := [])

    contents = []

    for file in files:
        content = file_extraction_functions[file[0]](file[1])
        contents.append({
            "type": file[0],
            "path": file[1],
            "filename": file[2],
            "content": content
        })

    return contents


def create_embeddings(pContent):
    conn = psycopg2.connect(
        dbname="embeddings",
        user="python",
        password="PasswordPassword123",
        host="localhost",
        port="5555"
    )
    cur = conn.cursor()

    create_table_query = '''
        create table if not exists dbtable (
        id SERIAL PRIMARY KEY,
        filepath TEXT NOT NULL,
        embedding VECTOR NOT NULL
        );
    '''

    cur.execute('CREATE EXTENSION IF NOT EXISTS vector;')

    cur.execute(create_table_query)

    conn.commit()

    for content in pContent:
        merged_info = "Dateiname: " + content["filename"] + " Dateiinhalt: " + content[
            "content"]  # Bessere Embeddings mit Dateiname // Information vorne dran?
        # print(merged_info)
        response = ollama.embeddings(model="mxbai-embed-large", prompt=merged_info)
        #embedding_list.append(response["embedding"])
        insert_data = f"insert into dbtable (filepath, embedding) Values ('{content['path']}', %s);"
        cur.execute(insert_data, (response["embedding"],))

    conn.commit()
    conn.close()


def reload_files():
    # path = input("Please provider path to folder: ")
    contents = retrieve_file_contents("C:\\SoftwareEng")
    create_embeddings(contents)


def add_files():
    path = input("Please provider path to folder: ")
    pass


def prompt_cycle():
    while True:
        prompt = input("Please enter prompt: ")
        lib.prompt_embedding(prompt)


def get_user_action():
    user_action = input("Reload Files (r), Add Files (a), Prompt (p): ")

    {
        "r": lambda: reload_files(),
        "a": lambda: add_files(),
        "p": lambda: prompt_cycle()
    }[user_action]()

get_user_action()

#print(json.dumps(contents, indent="\t"))