Maxed-Out-99-ComfyUI-MaxedOut/CharacterPrompts/storage.py

import json
import os
from datetime import datetime, timezone

import folder_paths

SCHEMA_VERSION = 1
EXTENSION_DIRNAME = "ComfyUI-CharacterPrompts-MXD"
FILENAME = "characters.json"


def normalize_key(value: str) -> str:
    return " ".join(value.strip().lower().split())


def parse_aliases_csv(aliases_csv: str) -> list[str]:
    if not aliases_csv:
        return []
    aliases = []
    seen = set()
    for raw in aliases_csv.split(","):
        alias = normalize_key(raw)
        if alias and alias not in seen:
            aliases.append(alias)
            seen.add(alias)
    return aliases


def get_library_dir() -> str:
    return os.path.join(folder_paths.get_user_directory(), "default", EXTENSION_DIRNAME)


def get_library_path() -> str:
    return os.path.join(get_library_dir(), FILENAME)


def _empty_schema() -> dict:
    return {"version": SCHEMA_VERSION, "characters": {}}


def _write_json_atomic(path: str, data: dict) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp_path = f"{path}.tmp"
    with open(tmp_path, "w", encoding="utf-8") as handle:
        json.dump(data, handle, indent=2, ensure_ascii=False)
    os.replace(tmp_path, path)


def _normalize_entry(canonical: str, entry: dict) -> dict:
    name = str(entry.get("name", canonical)).strip() or canonical
    prompt = str(entry.get("prompt", "")).strip()
    aliases_raw = entry.get("aliases", [])
    aliases = []
    seen = {canonical}
    if isinstance(aliases_raw, list):
        for alias_raw in aliases_raw:
            alias = normalize_key(str(alias_raw))
            if alias and alias not in seen:
                aliases.append(alias)
                seen.add(alias)
    updated_at = str(entry.get("updated_at", _utc_now()))
    return {
        "name": name,
        "prompt": prompt,
        "aliases": aliases,
        "updated_at": updated_at,
    }


def _normalize_schema(raw: dict) -> dict:
    if not isinstance(raw, dict):
        return _empty_schema()
    characters_raw = raw.get("characters", {})
    if not isinstance(characters_raw, dict):
        characters_raw = {}
    normalized = {}
    for key, entry in characters_raw.items():
        canonical = normalize_key(str(key))
        if not canonical or not isinstance(entry, dict):
            continue
        normalized[canonical] = _normalize_entry(canonical, entry)
    return {"version": SCHEMA_VERSION, "characters": normalized}


def _backup_corrupt_file(path: str) -> None:
    if not os.path.exists(path):
        return
    backup_path = f"{path}.bak"
    try:
        if os.path.exists(backup_path):
            os.remove(backup_path)
        os.replace(path, backup_path)
    except OSError:
        pass


def _utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def load_schema() -> dict:
    path = get_library_path()
    if not os.path.exists(path):
        schema = _empty_schema()
        _write_json_atomic(path, schema)
        return schema

    try:
        with open(path, "r", encoding="utf-8") as handle:
            raw = json.load(handle)
    except (OSError, json.JSONDecodeError):
        _backup_corrupt_file(path)
        schema = _empty_schema()
        _write_json_atomic(path, schema)
        return schema

    schema = _normalize_schema(raw)
    if schema != raw:
        _write_json_atomic(path, schema)
    return schema


def save_schema(schema: dict) -> None:
    normalized = _normalize_schema(schema)
    _write_json_atomic(get_library_path(), normalized)


def list_character_names(schema: dict | None = None) -> list[str]:
    schema = schema or load_schema()
    names = [entry["name"] for entry in schema["characters"].values()]
    return sorted(names, key=lambda value: value.lower())


def save_or_update_character(
    character_name: str,
    character_prompt: str,
    aliases_csv: str | None = "",
) -> tuple[dict, dict]:
    canonical = normalize_key(character_name)
    if not canonical:
        raise ValueError("Character name cannot be empty.")

    prompt = str(character_prompt or "").strip()
    if not prompt:
        raise ValueError("Character prompt cannot be empty.")

    schema = load_schema()
    existing_entry = schema["characters"].get(canonical)

    if aliases_csv is None:
        aliases_raw = existing_entry.get("aliases", []) if existing_entry else []
        aliases = []
        seen = {canonical}
        for alias_raw in aliases_raw:
            alias = normalize_key(str(alias_raw))
            if alias and alias not in seen:
                aliases.append(alias)
                seen.add(alias)
    else:
        aliases = [alias for alias in parse_aliases_csv(aliases_csv) if alias != canonical]

    schema["characters"][canonical] = {
        "name": character_name.strip() or canonical,
        "prompt": prompt,
        "aliases": aliases,
        "updated_at": _utc_now(),
    }
    save_schema(schema)
    return schema, schema["characters"][canonical]


def load_character(character_name: str) -> tuple[dict, str, dict]:
    canonical = normalize_key(character_name)
    if not canonical:
        raise ValueError("Character name cannot be empty.")
    schema = load_schema()
    entry = schema["characters"].get(canonical)
    if entry is None:
        raise KeyError(f"Character '{character_name}' was not found.")
    return schema, canonical, entry


def delete_character(character_name: str) -> tuple[dict, str]:
    canonical = normalize_key(character_name)
    if not canonical:
        raise ValueError("Character name cannot be empty.")
    schema = load_schema()
    if canonical not in schema["characters"]:
        raise KeyError(f"Character '{character_name}' was not found.")
    removed_name = schema["characters"][canonical]["name"]
    del schema["characters"][canonical]
    save_schema(schema)
    return schema, removed_name