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
from datetime import datetime
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 = []
    ctime = ""
    mtime = ""

    print("parsing files...")
    for file in tqdm(files):
        if not pPullFromS3:
            ctime = datetime.fromtimestamp(os.path.getctime(file[1])).strftime('%Y-%m-%d %H:%M:%S')
            mtime = datetime.fromtimestamp(os.path.getmtime(file[1])).strftime('%Y-%m-%d %H:%M:%S')
            print(ctime, mtime)
        content = file_extraction_functions[file[0]](file[1])
        if not pPullFromS3:
            push_files_into_s3(file, content, minioClient)

        contents.append({
            "type": file[0],
            "ctime": ctime if not pPullFromS3 else 'Diese Datei kommt ursprünglich nicht von dir, sondern aus einem S3-Bucket. ',
            "mtime": mtime if not pPullFromS3 else 'Diese Datei kommt ursprünglich nicht von dir, sondern aus einem S3-Bucket. ',
            "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
        })

    if os.path.exists(pulled_bucket_data):
        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, content['ctime'], content["mtime"])
    conn.commit()
    conn.close()

def insert_data_mongo(id, filepath,pChunk, pCtime, pMtime):
    client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/')
    mongodb = client['document_table']
    collection = mongodb['documents']
    dokument = {
        'doc_id': id,
        'ctime' : pCtime,
        'mtime' : pMtime,
        '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"))