"""
ComfyUI-PromptFormat

Cleans up prompts: removes duplicate tags, fixes bracket/comma spacing,
swaps underscores to spaces (with exclusions), and optionally appends
a trailing comma per line.

"""

from __future__ import annotations

import re
from pathlib import Path

import folder_paths  # type: ignore  # provided by ComfyUI at runtime


# --------------------------------------------------------------------------
# Core formatter (Python port of LeFormatter from prompt_format.js)
# --------------------------------------------------------------------------


class PromptFormatter:
    """Stateless formatter. All methods are classmethods; instance-free."""

    # Placeholder tokens used during pipeline so regex passes don't mangle
    # protected content.
    _SHY = "$SHY$"
    _CAT = "$CAT$"
    _TI_TEMPLATE = "@TEXTUAL{}INVERSION@"

    # --- Network (lora/lyco) stripping -----------------------------------

    # Matches angle-bracket networks like <lora:name:1.0> and lbw block weights.
    _RE_BLOCK_WEIGHT = re.compile(r"(lbw=)?\s*(\d+(\.\d+)?)(\s*,\s*(\d+(\.\d+)?))+")
    _RE_NETWORK = re.compile(r"\s*<.+?>\s*")

    @classmethod
    def _strip_networks(cls, text: str) -> str:
        """Remove <lora:...> / <lyco:...> / block-weight syntax entirely."""
        text = cls._RE_BLOCK_WEIGHT.sub("", text)
        text = cls._RE_NETWORK.sub(" ", text)
        return text

    # --- Expression protection (emoticons like "> <" and ":3") -----------

    _RE_SHY_IN = re.compile(r"(?:,|\n|^)\s*> <\s*(?:,|\n|$)")
    _RE_CAT_IN = re.compile(r"(?:,|\n|^)\s*:3\s*(?:,|\n|$)")

    @classmethod
    def _protect_expressions(cls, text: str) -> str:
        text = cls._RE_SHY_IN.sub(f", {cls._SHY},", text)
        text = cls._RE_CAT_IN.sub(f", {cls._CAT},", text)
        return text

    @classmethod
    def _restore_expressions(cls, text: str) -> str:
        return text.replace(cls._SHY, "> <").replace(cls._CAT, ":3")

    # --- Underscore handling ---------------------------------------------

    @classmethod
    def _rm_underscore(cls, text: str, protected: list[str]) -> str:
        """Replace underscores with spaces, but not inside protected names
        (embedding / textual-inversion filenames or user-excluded tags)."""
        if not text.strip():
            return ""

        # Swap protected names out for placeholders
        for i, name in enumerate(protected):
            if name and name in text:
                text = text.replace(name, cls._TI_TEMPLATE.format(i))

        text = text.replace("_", " ")

        for i, name in enumerate(protected):
            text = text.replace(cls._TI_TEMPLATE.format(i), name)

        return text

    # --- Deduplication ---------------------------------------------------

    _RE_KEYWORD = re.compile(r"^(AND|BREAK)$")
    _RE_BRACKETS = re.compile(r"[\[\]\(\)]")
    _RE_MULTISPACE = re.compile(r"\s+")

    @classmethod
    def _dedupe(cls, tags: list[str], aliases: list[tuple[re.Pattern, str]]) -> list[str]:
        unique: set[str] = set()
        out: list[str] = []

        for tag in tags:
            cleaned = cls._RE_BRACKETS.sub("", tag)
            cleaned = cls._RE_MULTISPACE.sub(" ", cleaned).strip()

            # Keep prompt-control keywords (AND, BREAK) as-is, always.
            if cls._RE_KEYWORD.match(cleaned):
                out.append(tag)
                continue

            # Pure numbers (e.g. weights) pass through.
            try:
                float(cleaned)
                out.append(tag)
                continue
            except ValueError:
                pass

            # Alias lookup: if tag matches a pattern, substitute with main tag
            substitute: str | None = None
            for pattern, main_tag in aliases:
                if pattern.match(cleaned):
                    substitute = main_tag
                    break

            if substitute is None and cleaned not in unique:
                unique.add(cleaned)
                out.append(tag)
                continue

            if substitute is not None and substitute not in unique:
                unique.add(substitute)
                out.append(tag.replace(cleaned, substitute))
                continue

            # Duplicate: blank out the content but keep the slot so brackets
            # upstream still balance. The later "prune empty chunks" pass
            # removes it.
            out.append(tag.replace(cleaned, ""))

        return out

    # --- Main per-line formatter -----------------------------------------

    # Comma-after-bracket cleanups
    _RE_FIX_COMMAS = [
        (re.compile(r",+\s*\)"), "),"),
        (re.compile(r",+\s*\]"), "],"),
        (re.compile(r",+\s*\}"), "},"),
        (re.compile(r"\(\s*,+"), ",("),
        (re.compile(r"\[\s*,+"), ",["),
        (re.compile(r"\{\s*,+"), ",{"),
    ]
    _RE_FIX_SPACES = [
        (re.compile(r"\s+\)"), ")"),
        (re.compile(r"\s+\]"), "]"),
        (re.compile(r"\s+\}"), "}"),
        (re.compile(r"\(\s+"), "("),
        (re.compile(r"\[\s+"), "["),
        (re.compile(r"\{\s+"), "{"),
    ]
    _RE_PIPE = re.compile(r"\s*\|\s*")
    _RE_COLON = re.compile(r"\s*\:\s*")
    _RE_EMPTY_BRACKET = re.compile(r"\(\s*\)|\[\s*\]")
    _RE_FRANCHISE = re.compile(r"\\\(([^\\\)]+?):([^\\\)]+?)\\\)")
    _RE_COLON_CLEAN = re.compile(r",\s*:(\d)")

    @classmethod
    def format_line(
        cls,
        line: str,
        dedupe: bool,
        rm_underscore: bool,
        protected: list[str],
        aliases: list[tuple[re.Pattern, str]],
    ) -> str:
        # 1) Strip LoRA/network syntax outright (ComfyUI handles those elsewhere)
        line = cls._strip_networks(line)

        # 2) Underscore replacement (with embedding name protection)
        if rm_underscore:
            line = cls._rm_underscore(line, protected)

        # 3) Protect emoticons
        line = cls._protect_expressions(line)

        # 4) Normalize commas around brackets
        for pattern, repl in cls._RE_FIX_COMMAS:
            line = pattern.sub(repl, line)
        for pattern, repl in cls._RE_FIX_SPACES:
            line = pattern.sub(repl, line)

        # 5) Tighten pipe and colon syntax
        line = cls._RE_PIPE.sub("|", line)
        line = cls._RE_COLON.sub(":", line)

        # 6) Split into tags, dedupe, rejoin
        tags = [t.strip() for t in line.split(",")]
        if dedupe:
            tags = cls._dedupe(tags, aliases)

        line = ", ".join(tags)
        line = cls._RE_MULTISPACE.sub(" ", line)

        # 7) Drop empty brackets iteratively
        while cls._RE_EMPTY_BRACKET.search(line):
            line = cls._RE_EMPTY_BRACKET.sub("", line)

        # 8) Franchise escape: \(series: name\) -> \(series: name\)  (space after colon)
        line = cls._RE_FRANCHISE.sub(r"\\(\1: \2\\)", line)

        # 9) Prune empty chunks
        line = ", ".join(t.strip() for t in line.split(",") if t.strip())

        # 10) Clean up stray empty-before-colon like ",:2"
        line = cls._RE_COLON_CLEAN.sub(r":\1", line)

        # 11) Restore protected emoticons
        line = cls._restore_expressions(line)

        return line

    @classmethod
    def format_pipeline(
        cls,
        text: str,
        dedupe: bool,
        rm_underscore: bool,
        append_comma: bool,
        protected: list[str],
        aliases: list[tuple[re.Pattern, str]],
    ) -> str:
        lines = text.split("\n")
        formatted = [
            cls.format_line(line, dedupe, rm_underscore, protected, aliases)
            for line in lines
        ]

        if not append_comma:
            return "\n".join(formatted)

        joined = ",\n".join(formatted)
        joined = re.sub(r"\n,\n", "\n\n", joined)
        joined = re.sub(r"\s*,\s*$", "", joined)
        return joined


# --------------------------------------------------------------------------
# Helpers for gathering embedding filenames + parsing user aliases
# --------------------------------------------------------------------------


def get_embedding_names() -> list[str]:
    """Scan ComfyUI's embedding folders for filenames to protect during
    underscore replacement."""
    extensions = (".pt", ".pth", ".ckpt", ".safetensors", ".sft")
    names: set[str] = set()

    try:
        folders = folder_paths.get_folder_paths("embeddings")
    except Exception:
        folders = []

    for folder in folders:
        base = Path(folder)
        if not base.is_dir():
            continue
        for ext in extensions:
            for file in base.rglob(f"*{ext}"):
                names.add(file.stem)

    return sorted(names)


def parse_aliases(raw: str) -> list[tuple[re.Pattern, str]]:
    """Parse alias config (same syntax as Forge extension):

        main_tag: alt1, alt2, regex_pattern
        1girl: girl, woman, lady
        adult: \\d+\\s*(y\\.?o\\.?|[Yy]ear[s]? [Oo]ld)
    """
    aliases: list[tuple[re.Pattern, str]] = []
    if ":" not in raw:
        return aliases

    for line in raw.splitlines():
        if ":" not in line:
            continue
        tag, words = line.split(":", 1)
        main = tag.strip()
        if not main:
            continue
        for word in (w.strip() for w in words.split(",")):
            if not word or word == main:
                continue
            # Strip anchors; we always anchor full-match in the formatter.
            cleaned = word.lstrip("^").rstrip("$")
            try:
                pattern = re.compile(f"^{cleaned}$")
            except re.error:
                continue
            aliases.append((pattern, main))

    return aliases


def parse_exclusions(raw: str) -> list[str]:
    """User-provided tags that should NOT have their underscores stripped
    (e.g. score_9, score_8_up)."""
    return [t.strip() for t in raw.split(",") if t.strip()]


# --------------------------------------------------------------------------
# ComfyUI Node
# --------------------------------------------------------------------------


class PromptFormatNode:
    """Clean up a prompt string.

    Inputs: raw prompt text + toggles.
    Output: cleaned prompt, ready to feed into a text encoder.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": (
                    "STRING",
                    {
                        "multiline": True,
                        "dynamicPrompts": True,
                        "placeholder": "your prompt here...",
                        "default": "",
                    },
                ),
                "dedupe": ("BOOLEAN", {"default": True}),
                "remove_underscores": ("BOOLEAN", {"default": True}),
                "append_comma": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "exclusions": (
                    "STRING",
                    {
                        "multiline": False,
                        "default": "",
                        "placeholder": "score_9, score_8_up, score_7_up",
                    },
                ),
                "aliases": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "",
                        "placeholder": "1girl: girl, woman, lady\nadult: \\d+\\s*(y\\.?o\\.?|[Yy]ear[s]? [Oo]ld)",
                    },
                ),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)
    FUNCTION = "format"
    CATEGORY = "Cyberdelia/text"
    OUTPUT_NODE = False

    # Cache embedding list across calls within a session; invalidate cheaply.
    _embedding_cache: list[str] | None = None

    @classmethod
    def _get_protected(cls, exclusions: str) -> list[str]:
        if cls._embedding_cache is None:
            cls._embedding_cache = get_embedding_names()
        extras = parse_exclusions(exclusions)
        # Longest names first so "score_8_up" matches before "score".
        return sorted(set(cls._embedding_cache + extras), key=len, reverse=True)

    def format(
        self,
        text: str,
        dedupe: bool,
        remove_underscores: bool,
        append_comma: bool,
        exclusions: str = "",
        aliases: str = "",
    ):
        protected = self._get_protected(exclusions)
        alias_list = parse_aliases(aliases)

        cleaned = PromptFormatter.format_pipeline(
            text,
            dedupe=dedupe,
            rm_underscore=remove_underscores,
            append_comma=append_comma,
            protected=protected,
            aliases=alias_list,
        )
        return (cleaned,)


class PromptFormatEncodeNode:
    """Format a prompt AND encode it with a CLIP model in one step.

    Equivalent to chaining PromptFormat -> CLIPTextEncode.
    Output is CONDITIONING, ready for a sampler.
    """

    @classmethod
    def INPUT_TYPES(cls):
        base = PromptFormatNode.INPUT_TYPES()
        base["required"] = {"clip": ("CLIP",), **base["required"]}
        return base

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION = "encode"
    CATEGORY = "Cyberdelia/text"

    def encode(
        self,
        clip,
        text: str,
        dedupe: bool,
        remove_underscores: bool,
        append_comma: bool,
        exclusions: str = "",
        aliases: str = "",
    ):
        formatter = PromptFormatNode()
        (cleaned,) = formatter.format(
            text, dedupe, remove_underscores, append_comma, exclusions, aliases
        )

        tokens = clip.tokenize(cleaned)
        output = clip.encode_from_tokens_scheduled(tokens)
        return (output,)


# --------------------------------------------------------------------------
# API route for the frontend "Format" button
# --------------------------------------------------------------------------

try:
    from aiohttp import web
    from server import PromptServer  # type: ignore

    @PromptServer.instance.routes.post("/cyberdelia/prompt_format")
    async def _format_route(request):
        """Frontend calls this when the user clicks the ✨ Format button.
        Returns the cleaned prompt so the JS can write it back into the
        widget's textarea."""
        try:
            data = await request.json()
        except Exception:
            return web.json_response({"error": "invalid json"}, status=400)

        text = data.get("text", "") or ""
        dedupe = bool(data.get("dedupe", True))
        rm_underscore = bool(data.get("remove_underscores", True))
        append_comma = bool(data.get("append_comma", True))
        exclusions = data.get("exclusions", "") or ""
        aliases = data.get("aliases", "") or ""

        protected = PromptFormatNode._get_protected(exclusions)
        alias_list = parse_aliases(aliases)

        cleaned = PromptFormatter.format_pipeline(
            text,
            dedupe=dedupe,
            rm_underscore=rm_underscore,
            append_comma=append_comma,
            protected=protected,
            aliases=alias_list,
        )
        return web.json_response({"text": cleaned})

except ImportError:
    # If we're imported outside ComfyUI (tests, tooling), skip route setup.
    pass