import ollama
import pymupdf
import os

import chunker
import lib
import json
import psycopg2
import pymongo

count = 0

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):
    global count
    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:
        for chunk in chunker.generate_chunks(content["content"]):
            merged_info = "Dateiname: " + content["filename"] + " Dateiinhalt: " + chunk
            # 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) Returning id;"
            cur.execute(insert_data, (response["embedding"],))
            doc_id = cur.fetchone()[0]
            insert_data_mongo(doc_id, content['path'], chunk)
    conn.commit()
    conn.close()

def insert_data_mongo(id, filepath,pChunk):
    client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/')
    mongodb = client['document_table']
    collection = mongodb['documents']
    dokument = {
        'doc_id': id,
        'filepath': filepath,
        'chunk_content': pChunk,
    }
    collection.insert_one(dokument)

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"))