"""
Thread-safe JSON database for storing prompts with categories.
Prompts can optionally have LoRA attachments.
"""

import json
import os
import uuid
import threading

DB_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "prompts.json")
_lock = threading.Lock()

DEFAULT_DB = {
    "categories": ["Location", "Character", "Style", "Mood"],
    "prompts": []
}


def load_db():
    """Load the database from disk. Returns default if file missing or corrupted."""
    with _lock:
        if not os.path.exists(DB_FILE):
            save_db_unsafe(DEFAULT_DB)
            return DEFAULT_DB.copy()
        try:
            with open(DB_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
            if "categories" not in data or "prompts" not in data:
                data = DEFAULT_DB.copy()
                save_db_unsafe(data)
            # Ensure all prompts have a loras field
            for p in data.get("prompts", []):
                if "loras" not in p:
                    p["loras"] = []
            return data
        except (json.JSONDecodeError, IOError):
            save_db_unsafe(DEFAULT_DB)
            return DEFAULT_DB.copy()


def save_db_unsafe(data):
    """Save database without acquiring lock (caller must hold lock)."""
    with open(DB_FILE, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


def save_db(data):
    """Save database to disk (thread-safe)."""
    with _lock:
        save_db_unsafe(data)


def get_categories():
    """Return list of category names."""
    data = load_db()
    return data.get("categories", [])


def get_prompts(category=None):
    """Return prompts, optionally filtered by category."""
    data = load_db()
    prompts = data.get("prompts", [])
    if category:
        prompts = [p for p in prompts if p.get("category") == category]
    return prompts


def get_prompt_by_id(prompt_id):
    """Return a single prompt by its ID, or None."""
    data = load_db()
    for p in data.get("prompts", []):
        if p.get("id") == prompt_id:
            return p
    return None


def get_prompt_by_name(name, category=None):
    """Return a single prompt by its name and optionally category, or None."""
    data = load_db()
    for p in data.get("prompts", []):
        if p.get("name") == name:
            if category is None or p.get("category") == category:
                return p
    return None


def add_prompt(category, name, text, loras=None):
    """Add a new prompt. Returns the new prompt dict."""
    with _lock:
        data = _load_unsafe()
        prompt = {
            "id": str(uuid.uuid4()),
            "category": category,
            "name": name,
            "text": text,
            "loras": loras or [],
        }
        data["prompts"].append(prompt)
        save_db_unsafe(data)
        return prompt


def update_prompt(prompt_id, category=None, name=None, text=None, loras=None):
    """Update an existing prompt by ID. Returns updated prompt or None."""
    with _lock:
        data = _load_unsafe()
        for p in data["prompts"]:
            if p["id"] == prompt_id:
                if category is not None:
                    p["category"] = category
                if name is not None:
                    p["name"] = name
                if text is not None:
                    p["text"] = text
                if loras is not None:
                    p["loras"] = loras
                save_db_unsafe(data)
                return p
        return None


def delete_prompt(prompt_id):
    """Delete a prompt by ID. Returns True if found and deleted."""
    with _lock:
        data = _load_unsafe()
        original_len = len(data["prompts"])
        data["prompts"] = [p for p in data["prompts"] if p["id"] != prompt_id]
        if len(data["prompts"]) < original_len:
            save_db_unsafe(data)
            return True
        return False


def add_category(name):
    """Add a new category. Returns True if added, False if already exists."""
    with _lock:
        data = _load_unsafe()
        if name in data["categories"]:
            return False
        data["categories"].append(name)
        save_db_unsafe(data)
        return True


def delete_category(name):
    """Delete a category and all its prompts. Returns True if found and deleted."""
    with _lock:
        data = _load_unsafe()
        if name not in data["categories"]:
            return False
        data["categories"].remove(name)
        data["prompts"] = [p for p in data["prompts"] if p["category"] != name]
        save_db_unsafe(data)
        return True


def reorder_categories(ordered_names):
    """Reorder categories to match the given list. Returns True if successful."""
    with _lock:
        data = _load_unsafe()
        existing = set(data["categories"])
        given = set(ordered_names)
        if existing != given:
            return False
        data["categories"] = list(ordered_names)
        save_db_unsafe(data)
        return True


def reorder_prompts(ordered_ids):
    """Reorder prompts to match the given ID list. Unmentioned prompts keep their relative order at the end."""
    with _lock:
        data = _load_unsafe()
        by_id = {p["id"]: p for p in data["prompts"]}
        reordered = []
        for pid in ordered_ids:
            if pid in by_id:
                reordered.append(by_id.pop(pid))
        # Append any remaining prompts not in the ordered list
        for p in data["prompts"]:
            if p["id"] in by_id:
                reordered.append(p)
        data["prompts"] = reordered
        save_db_unsafe(data)
        return True


def get_available_loras():
    """Return list of available LoRA filenames from ComfyUI's models folder."""
    try:
        import folder_paths
        return folder_paths.get_filename_list("loras")
    except Exception:
        return []


def export_db():
    """Return full database dict for export."""
    return load_db()


def import_db(incoming, mode="merge"):
    """Import prompts from an incoming dict.
    mode='overwrite': replace entire DB.
    mode='merge': add new categories and prompts, skip duplicates (by name+category).
    Returns dict with stats: {categories_added, prompts_added, prompts_skipped}.
    """
    with _lock:
        if mode == "overwrite":
            # Ensure all prompts have IDs and loras
            for p in incoming.get("prompts", []):
                if "id" not in p:
                    p["id"] = str(uuid.uuid4())
                if "loras" not in p:
                    p["loras"] = []
            save_db_unsafe(incoming)
            return {
                "categories_added": len(incoming.get("categories", [])),
                "prompts_added": len(incoming.get("prompts", [])),
                "prompts_skipped": 0,
            }

        # Merge mode
        data = _load_unsafe()
        cats_added = 0
        prompts_added = 0
        prompts_skipped = 0

        # Add missing categories
        existing_cats = set(data["categories"])
        for cat in incoming.get("categories", []):
            if cat not in existing_cats:
                data["categories"].append(cat)
                existing_cats.add(cat)
                cats_added += 1

        # Add prompts that don't exist yet (by name + category)
        existing_prompts = {(p["name"], p["category"]) for p in data["prompts"]}
        for p in incoming.get("prompts", []):
            key = (p.get("name", ""), p.get("category", ""))
            if key in existing_prompts:
                prompts_skipped += 1
                continue
            prompt = {
                "id": p.get("id") or str(uuid.uuid4()),
                "name": p.get("name", ""),
                "category": p.get("category", ""),
                "text": p.get("text", ""),
                "loras": p.get("loras", []),
            }
            data["prompts"].append(prompt)
            existing_prompts.add(key)
            prompts_added += 1

        save_db_unsafe(data)
        return {
            "categories_added": cats_added,
            "prompts_added": prompts_added,
            "prompts_skipped": prompts_skipped,
        }


def _load_unsafe():
    """Load database without lock (internal use only)."""
    if not os.path.exists(DB_FILE):
        save_db_unsafe(DEFAULT_DB)
        return DEFAULT_DB.copy()
    try:
        with open(DB_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        # Ensure all prompts have loras field
        for p in data.get("prompts", []):
            if "loras" not in p:
                p["loras"] = []
        return data
    except (json.JSONDecodeError, IOError):
        save_db_unsafe(DEFAULT_DB)
        return DEFAULT_DB.copy()