Bam_Prompt_Parser.py

from __future__ import annotations

import re
from typing import Dict, List, Tuple


# ============================================================================
# Output definitions (order matters!)
# ============================================================================
_OUTPUT_DEFS: List[Tuple[str, str]] = [
    # Core identity
    ("gender_str", "STRING"),
    ("gender_int", "INT"),
    ("age_str", "STRING"),
    ("age_int", "INT"),
    ("identity_str", "STRING"),
    ("eyecolor_str", "STRING"),
    ("hairstyle_str", "STRING"),

    # Equipment (strings)
    ("topwear_str", "STRING"),
    ("bellywear_str", "STRING"),
    ("breastwear_str", "STRING"),
    ("handwear_left_str", "STRING"),
    ("handwear_right_str", "STRING"),
    ("wristwear_left_str", "STRING"),
    ("wristwear_right_str", "STRING"),
    ("forearm_left_str", "STRING"),
    ("forearm_right_str", "STRING"),
    ("elbow_left_str", "STRING"),
    ("elbow_right_str", "STRING"),
    ("upperarm_left_str", "STRING"),
    ("upperarm_right_str", "STRING"),
    ("shoulder_left_str", "STRING"),
    ("shoulder_right_str", "STRING"),
    ("shank_left_str", "STRING"),
    ("shank_right_str", "STRING"),
    ("knee_left_str", "STRING"),
    ("knee_right_str", "STRING"),
    ("foot_left_str", "STRING"),
    ("foot_right_str", "STRING"),
    ("necklace_str", "STRING"),
    ("earring_left_str", "STRING"),
    ("earring_right_str", "STRING"),
    ("kneewear_str", "STRING"),
    ("headwear_str", "STRING"),
    ("facemask_str", "STRING"),
    ("sunglasses_str", "STRING"),
    ("glasses_str", "STRING"),
    ("crotch_str", "STRING"),
    ("one_piece_str", "STRING"),

    # Tags (strings)
    ("aesthetic_tag1", "STRING"),
    ("aesthetic_tag2", "STRING"),
    ("aesthetic_tag3", "STRING"),
    ("aesthetic_tag4", "STRING"),
    ("aesthetic_tag5", "STRING"),

    ("skin_tag1", "STRING"),
    ("skin_tag2", "STRING"),
    ("skin_tag3", "STRING"),
    ("skin_tag4", "STRING"),
    ("skin_tag5", "STRING"),

    ("expression_tag1", "STRING"),
    ("expression_tag2", "STRING"),
    ("expression_tag3", "STRING"),
    ("expression_tag4", "STRING"),
    ("expression_tag5", "STRING"),

    # Unique headwear slot AFTER expression tags (your step 23)
    ("headwear_str_2", "STRING"),

    # Equipment section simplified (values only, comma-separated)
    ("old_bam", "STRING"),
]

RETURN_NAMES_TUPLE = tuple(name for name, _t in _OUTPUT_DEFS)
RETURN_TYPES_TUPLE = tuple(_t for _name, _t in _OUTPUT_DEFS)


# ============================================================================
# Equipment parsing helpers
# ============================================================================
def _extract_parenthesized_items(s: str) -> List[str]:
    """
    Extract top-level (...) items from a string, handling nested parentheses.
    Returns inside text of each (...) (outer parens removed).
    """
    items: List[str] = []
    depth = 0
    buf: List[str] = []

    for ch in s or "":
        if ch == "(":
            if depth == 0:
                buf = []
            else:
                buf.append(ch)
            depth += 1
        elif ch == ")":
            if depth > 0:
                depth -= 1
            if depth == 0:
                item = "".join(buf).strip()
                if item:
                    items.append(item)
            else:
                buf.append(ch)
        else:
            if depth > 0:
                buf.append(ch)

    return items


def _normalize_key(k: str) -> str:
    k = (k or "").strip().lower()
    k = k.replace(" ", "_").replace("-", "_")
    k = re.sub(r"_+", "_", k)
    return k


# Map key spellings to a canonical internal key.
_KEY_CANONICAL: Dict[str, str] = {
    "belly": "bellywear",
    "bellywear": "bellywear",

    "topwear": "topwear",

    "breast": "breastwear",
    "breastwear": "breastwear",

    "hand": "handwear",
    "handwear": "handwear",

    "wrist": "wristwear",
    "wristwear": "wristwear",

    "forearm": "forearm",
    "elbow": "elbow",

    "upperarm": "upperarm",
    "upper_arm": "upperarm",

    "shoulder": "shoulder",
    "shank": "shank",

    "knee": "knee",

    "foot": "foot",
    "footwear": "foot",   # your example uses Footwear:
    "shoe": "foot",
    "shoes": "foot",

    "necklace": "necklace",

    "earring": "earring",
    "earrings": "earring",

    "kneewear": "kneewear",
    "headwear": "headwear",

    "facemask": "facemask",
    "face_mask": "facemask",
    "mask": "facemask",

    "sunglasses": "sunglasses",
    "glasses": "glasses",

    "crotch": "crotch",

    "onepiece": "one_piece",
    "one_piece": "one_piece",
    "one_piecewear": "one_piece",
}

# Canonical keys that are side-aware (Left/Right)
_SIDE_FIELDS: Dict[str, Tuple[str, str]] = {
    "handwear": ("handwear_left_str", "handwear_right_str"),
    "wristwear": ("wristwear_left_str", "wristwear_right_str"),
    "forearm": ("forearm_left_str", "forearm_right_str"),
    "elbow": ("elbow_left_str", "elbow_right_str"),
    "upperarm": ("upperarm_left_str", "upperarm_right_str"),
    "shoulder": ("shoulder_left_str", "shoulder_right_str"),
    "shank": ("shank_left_str", "shank_right_str"),
    "knee": ("knee_left_str", "knee_right_str"),
    "foot": ("foot_left_str", "foot_right_str"),
    "earring": ("earring_left_str", "earring_right_str"),
}

# Canonical keys that map to a single output
_SINGLE_FIELDS: Dict[str, str] = {
    "topwear": "topwear_str",
    "bellywear": "bellywear_str",
    "breastwear": "breastwear_str",
    "necklace": "necklace_str",
    "kneewear": "kneewear_str",
    "headwear": "headwear_str",
    "facemask": "facemask_str",
    "sunglasses": "sunglasses_str",
    "glasses": "glasses_str",
    "crotch": "crotch_str",
    "one_piece": "one_piece_str",
}

_ALL_EQUIP_OUTPUTS = set(_SINGLE_FIELDS.values())
for lf, rf in _SIDE_FIELDS.values():
    _ALL_EQUIP_OUTPUTS.add(lf)
    _ALL_EQUIP_OUTPUTS.add(rf)


def _parse_equipment(equip: str) -> Tuple[Dict[str, str], str]:
    """
    Parse equipment section like:
      (Knee_Left: Blue Skater Kneepad), (Knee:Green Skater Kneepads), (Belly:), ...

    Rules implemented:
    - Missing entries => output stays ""
    - Empty "(Belly:)" => output ""
    - Side-less "(Knee:...)" => assigns to BOTH left and right outputs
    - old_bam => values-only, comma-separated, in original order
    """
    out: Dict[str, str] = {name: "" for name in _ALL_EQUIP_OUTPUTS}
    old_values: List[str] = []

    for item in _extract_parenthesized_items(equip or ""):
        if ":" not in item:
            continue

        raw_key, raw_val = item.split(":", 1)
        key = _normalize_key(raw_key)
        val = (raw_val or "").strip()

        # Trim a trailing comma, if any
        if val.endswith(","):
            val = val[:-1].rstrip()

        # old_bam: collect only non-empty values
        if val:
            old_values.append(val)

        # Detect side suffix
        side = None
        base_key = key
        if base_key.endswith("_left"):
            side = "left"
            base_key = base_key[:-5]
        elif base_key.endswith("_right"):
            side = "right"
            base_key = base_key[:-6]
        base_key = base_key.rstrip("_")

        canonical = _KEY_CANONICAL.get(base_key, base_key)

        if canonical in _SIDE_FIELDS:
            left_name, right_name = _SIDE_FIELDS[canonical]
            if side == "left":
                out[left_name] = val
            elif side == "right":
                out[right_name] = val
            else:
                # "(Knee:...)" style => both sides
                out[left_name] = val
                out[right_name] = val
        elif canonical in _SINGLE_FIELDS:
            out[_SINGLE_FIELDS[canonical]] = val
        else:
            # Unknown equipment key => ignored for structured outputs,
            # but still included in old_bam via old_values.
            pass

    old_bam = ", ".join(v for v in old_values if v)
    return out, old_bam


# ============================================================================
# BAM parsing
# ============================================================================
def _get_part(parts: List[str], idx: int) -> str:
    return parts[idx] if idx < len(parts) else ""


def _safe_int(s: str, default: int = 0) -> int:
    try:
        return int((s or "").strip())
    except Exception:
        return default


def _parse_bam(bam: str) -> Dict[str, object]:
    """
    Expected structure after the first START### marker:

      0  gender_num
      1  age
      2  identity
      3  eyecolor
      4  hairstyle
      5  equipment
      6..10   aesthetic_tag1..5
      11..15  skin_tag1..5
      16..20  expression_tag1..5
      21 headwear_str_2
      22+ irrelevant
    """
    bam = bam or ""
    marker = "START###"
    idx = bam.find(marker)
    payload = bam[idx + len(marker):] if idx != -1 else bam

    parts = payload.split("###")

    gender_token = _get_part(parts, 0).strip()
    gender_int = 1 if gender_token == "1" else 2
    gender_str = "boy" if gender_int == 1 else "girl"

    age_str = _get_part(parts, 1).strip()
    age_int = _safe_int(age_str, default=0)

    identity_str = _get_part(parts, 2).strip()
    eyecolor_str = _get_part(parts, 3).strip()
    hairstyle_str = _get_part(parts, 4).strip()

    equipment_raw = _get_part(parts, 5).strip()
    equip_out, old_bam = _parse_equipment(equipment_raw)

    out: Dict[str, object] = {
        "gender_str": gender_str,
        "gender_int": gender_int,
        "age_str": age_str,
        "age_int": age_int,
        "identity_str": identity_str,
        "eyecolor_str": eyecolor_str,
        "hairstyle_str": hairstyle_str,
        "old_bam": old_bam,
    }

    # Equipment
    out.update(equip_out)

    # Aesthetic tags
    for i in range(5):
        out[f"aesthetic_tag{i+1}"] = _get_part(parts, 6 + i).strip()

    # Skin tags
    for i in range(5):
        out[f"skin_tag{i+1}"] = _get_part(parts, 11 + i).strip()

    # Expression tags
    for i in range(5):
        out[f"expression_tag{i+1}"] = _get_part(parts, 16 + i).strip()

    # Outer-layer headwear slot
    out["headwear_str_2"] = _get_part(parts, 21).strip()

    return out


# ============================================================================
# ComfyUI node
# ============================================================================
class BAM_Prompt_Parser:
    """
    ComfyUI custom node: parses your BAM string and exposes each field as outputs.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "bam_string": ("STRING", {"multiline": True, "default": ""}),
            }
        }

    RETURN_TYPES = RETURN_TYPES_TUPLE
    RETURN_NAMES = RETURN_NAMES_TUPLE
    FUNCTION = "parse"
    CATEGORY = "BAM"

    def parse(self, bam_string: str):
        parsed = _parse_bam(bam_string)

        # Guarantee every declared output exists (malformed BAM => defaults)
        for name, t in _OUTPUT_DEFS:
            if name not in parsed:
                parsed[name] = 0 if t == "INT" else ""

        return tuple(parsed[name] for name in RETURN_NAMES_TUPLE)


NODE_CLASS_MAPPINGS = {
    "BAM_Prompt_Parser": BAM_Prompt_Parser,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BAM_Prompt_Parser": "BAM Prompt Parser",
}