import ollama
import pymupdf
import os
import shutil
from minio import Minio
from minio.error import S3Error
import chunker
import lib
import json
import psycopg2
import pymongo

import mongo_conn
import pgvec_conn
import redis_conn
from tqdm import tqdm

def retrieve_file_contents(path, pPullFromS3):
    pulled_bucket_data = os.path.join(".", "MINIO_TEMP_DATA")

    #Minio Connection
    minioClient = Minio(
        "localhost:9000",
        access_key= "PVFOeJbx87rQyi0WXF1X",
        secret_key = "Am8Cd9auYGbEGuEXfJtnWEPsMwJCx9N58NCNHCgs",
        secure=False,
    )

    file_extraction_functions = {
        "pdf": lambda path: lib.extract_text_and_pictures(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),
    }

    if pPullFromS3:
        path = pull_files_from_s3(pulled_bucket_data, minioClient, )

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

    contents = []

    print("parsing files...")
    for file in tqdm(files):
        content = file_extraction_functions[file[0]](file[1])
        if not pPullFromS3:
            push_files_into_s3(file, content, minioClient)

        contents.append({
            "type": file[0],
            "path": file[1] if not pPullFromS3 else "Diese Datei kommt ursprünglich nicht von dir, sondern aus einem S3-Bucket. ",
            "filename": file[2],
            "content": content
        })

    shutil.rmtree(pulled_bucket_data)

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

    print("generating embeddings...")
    for content in tqdm(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,
    }
    result = collection.insert_one(dokument)

    if result.acknowledged:
        collection.create_index("doc_id")


def push_files_into_s3(file, content, client):
    with open(f"./{file[2]}.txt", "w", encoding="UTF-8") as future_s3_file:
        future_s3_file.writelines(content)
        future_s3_file.flush()

    client.fput_object(
        bucket_name="datafiles",
        object_name=f"{file[2]}.txt",
        file_path=f"./{file[2]}.txt",
    )
    os.remove(f"./{file[2]}.txt")

def pull_files_from_s3(pPath, client):
    if not os.path.exists(pPath):
        os.makedirs(pPath)
    file_objects = client.list_objects("datafiles")
    for file in file_objects:
        local_path = os.path.join(pPath, file.object_name)
        client.fget_object("datafiles", file.object_name, local_path)
    return pPath


def reload_files(pPullFromS3):
    pgvec_conn.flush_pg()
    redis_conn.flush_redis()
    mongo_conn.flush_mongo()
    if not (pPullFromS3):
        provided_path = input("Please Provide Full Qualified Path: ")
        contents = retrieve_file_contents(provided_path, pPullFromS3)
    else:
        provided_path = ""
        contents = retrieve_file_contents(provided_path, pPullFromS3)

    create_embeddings(contents)


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


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

        response = redis_conn.load_response_from_redis(prompt.lower().strip())

        lib.prompt_embedding(prompt) if response is None else print("Cached Response:", response)


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

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

get_user_action()

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