BAN_Expression_To_Default.py

from __future__ import annotations

import re


DEFAULT_EXPRESSION_1 = "calm expression"
DEFAULT_EXPRESSION_2 = "neutral expression"
DEFAULT_EXPRESSION_3 = "slight smug smile"


# Matches:
# expression.1=...
# expression.2: ...
# expression_tag1=...
# expression1=...
_EXPR_RE = re.compile(
    r"^expression(?:\.|_tag)?([1-5])\s*[:=]\s*(.*)$",
    flags=re.IGNORECASE,
)

_SKIN_RE = re.compile(
    r"^skin(?:\.|_tag)?([1-5])\s*[:=]",
    flags=re.IGNORECASE,
)

_AESTHETIC_RE = re.compile(
    r"^aesthetic(?:\.|_tag)?([1-5])\s*[:=]",
    flags=re.IGNORECASE,
)

_EQUIP_RE = re.compile(
    r"^(?:equip|equipment)\.",
    flags=re.IGNORECASE,
)

_GPT_BAM_BLOCK_RE = re.compile(
    r"GPT_BAM_START###(.*?)###GPT_BAM_END",
    flags=re.IGNORECASE | re.DOTALL,
)


def _split_segments(payload: str) -> list[str]:
    """Split a GPT_BAM payload into ### segments, removing empty/newline-only parts."""
    payload = (payload or "").replace("\r", "\n")
    return [seg.strip() for seg in payload.split("###") if seg.strip()]


def _is_expression_segment(seg: str) -> bool:
    return _EXPR_RE.match(seg) is not None


def _find_insertion_index(segments: list[str]) -> int:
    """
    Where to insert expression.1/2/3 if no expression exists yet:

    1) first existing expression position
    2) after last skin tag
    3) after last aesthetic tag
    4) after last equip.* tag
    5) otherwise at end
    """
    first_expr_idx = None
    last_skin_idx = None
    last_aesthetic_idx = None
    last_equip_idx = None

    for i, seg in enumerate(segments):
        if _is_expression_segment(seg) and first_expr_idx is None:
            first_expr_idx = i
        if _SKIN_RE.match(seg):
            last_skin_idx = i
        if _AESTHETIC_RE.match(seg):
            last_aesthetic_idx = i
        if _EQUIP_RE.match(seg):
            last_equip_idx = i

    if first_expr_idx is not None:
        return first_expr_idx
    if last_skin_idx is not None:
        return last_skin_idx + 1
    if last_aesthetic_idx is not None:
        return last_aesthetic_idx + 1
    if last_equip_idx is not None:
        return last_equip_idx + 1
    return len(segments)


def _rewrite_payload(payload: str) -> str:
    segments = _split_segments(payload)
    insertion_index = _find_insertion_index(segments)

    new_expression_segments = [
        f"expression.1={DEFAULT_EXPRESSION_1}",
        f"expression.2={DEFAULT_EXPRESSION_2}",
        f"expression.3={DEFAULT_EXPRESSION_3}",
    ]

    out_segments: list[str] = []
    inserted = False

    for i, seg in enumerate(segments):
        if not inserted and i == insertion_index:
            out_segments.extend(new_expression_segments)
            inserted = True

        # Remove ALL existing expression.1..5 / expression_tag1..5
        if _is_expression_segment(seg):
            continue

        out_segments.append(seg)

    if not inserted:
        out_segments.extend(new_expression_segments)

    return "###".join(out_segments)


def _rewrite_gpt_bam_text(text: str) -> str:
    """
    If a GPT_BAM block is found, only rewrite that block and preserve any text outside it.
    If no GPT_BAM block is found, treat the whole input as a raw payload and rewrite it.
    """
    text = text or ""
    match = _GPT_BAM_BLOCK_RE.search(text)

    if not match:
        # Fallback: treat entire text as payload and wrap it back into GPT_BAM
        rewritten_payload = _rewrite_payload(text)
        return f"GPT_BAM_START###{rewritten_payload}###GPT_BAM_END"

    original_payload = match.group(1)
    rewritten_payload = _rewrite_payload(original_payload)
    new_block = f"GPT_BAM_START###{rewritten_payload}###GPT_BAM_END"

    return text[:match.start()] + new_block + text[match.end():]


class BAM_expression_default:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "BAM-format_In": ("STRING", {"multiline": True, "default": ""}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("BAM-format_OUT",)
    FUNCTION = "apply"
    CATEGORY = "BAM"

    def apply(self, **kwargs):
        bam_in = kwargs.get("BAM-format_In", "")
        bam_out = _rewrite_gpt_bam_text(bam_in)
        return (bam_out,)


NODE_CLASS_MAPPINGS = {
    "BAM_expression_default": BAM_expression_default,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BAM_expression_default": "BAM_expression_default",
}