import chromadb
from chromadb.utils import embedding_functions


def get_client():
    client = chromadb.PersistentClient(path="./chromadb_linux/")
    MODEL_NAME: str = "mixedbread-ai/mxbai-embed-large-v1"  # ~ 0.5 gb
    COLLECTION_NAME: str = "scheme"
    EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name=MODEL_NAME
    )
    schemer = client.get_collection(
        name=COLLECTION_NAME,
        embedding_function=EMBEDDING_FUNC,
    )
    return schemer


def update_collection(iter: int, text: object, client: chromadb.Collection):
    client.add(documents=[text["text"]], metadatas=[{"source": "pdf"}], ids=[text["content"] + str(iter)])


def encode_image(img_path: str):
    import base64

    with open(img_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


async def image_to_text(image) -> object:
    from openai import OpenAI
    import json

    client = OpenAI()

    response = client.chat.completions.create(
        model="gpt-4-turbo",
        response_format={"type": "json_object"},
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Transcribe the contents of this image and return a JSON object that contains the text. It must be structured in the following manner: two entries with the following keys: 'content' and 'text'. Content will be a line describing what the content of text will be, and text will be a simple transcription of the image"},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64;,{image}",
                            "detail": "high",
                        },
                    },
                ],
            }
        ],
    )
    return json.loads(response.choices[0].message.content)


async def start_troggin_off(dir: str):
    import os
    from pdf2image import convert_from_path

    client = get_client()

    for folder in os.listdir(dir):
        folder_path = os.path.join(dir, folder)
        if os.path.isdir(folder_path):
            for file in os.listdir(folder_path):
                if file.endswith(".pdf"):
                    print("Processing", file)
                    pdf_path = os.path.join(folder_path, file)
                    images = convert_from_path(pdf_path)

                    for i, image in enumerate(images):
                        image.save(f"out{i}.jpg", "JPEG")
                        encoded_image = encode_image(f"out{i}.jpg")
                        text = await image_to_text(encoded_image)
                        update_collection(i, text, client)


if __name__ == "__main__":
    import asyncio
    asyncio.run(start_troggin_off("data/Class Notes/"))