import json
import os
import re
import shutil
import uuid
from dataclasses import dataclass, asdict
from pathlib import Path

import torch

import folder_paths

STORE_PATH = Path("prompt_bench") / "store"


@dataclass
class PBStoreItem:
    prompt: str
    rating: str
    model: str
    text: str


def store_path(store_id: str):
    path_to_store = folder_paths.get_user_directory() / STORE_PATH / store_id
    os.makedirs(path_to_store, exist_ok=True)
    return path_to_store


def get_store_path_items(store_id: str):
    return len(list(filter(lambda i: ".json" in i, os.listdir(store_path(store_id)))))


def save_item(store_id: str, data: PBStoreItem):
    items = get_store_path_items(store_id)
    path = os.path.join(store_path(store_id), f"{items + 1}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump(asdict(data), f, indent=2)


class PBStore:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING",),
                "rating": ("RATING",),
                "model": ("STRING",),
                "text": ("STRING",),
            },
            "optional": {
                "images": ("IMAGE", {"tooltip": "The images to save."})
            },
        }

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return str(uuid.uuid4())

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    FUNCTION = "process"
    CATEGORY = "Prompt Bench"
    OUTPUT_NODE = True

    def process(self, prompt: str, rating: dict, model: str, text: str, images: torch.Tensor | None = None):
        item = PBStoreItem(
            prompt=prompt,
            rating=json.dumps(rating, ensure_ascii=True, sort_keys=True),
            model=model,
            text=text,
        )

        # Keep store id at max 128 if longer.
        store_id = re.sub("\W+", "_", text)[:128]

        if not store_id or len(store_id) >= 0:
            store_id = "unknown"

        save_item(store_id, item)

        if images is not None:
            for image in images:
                pass
                # shutil.copy2(store_path(store_id), None)

        return ()