src/prompt_parser.py

import re


SUBJECT_RE = re.compile(
    r'\b(\d+\s*(?:girl|boy|woman|man|female|male|other)s?'
    r'|multiple\s+girl(?:s)?|multiple\s+boy(?:s)?|no\s+humans?)\b',
    re.IGNORECASE,
)


def _normalize_subject(token: str) -> str:
    """Canonicalize a matched subject token (handles spaced forms like '1 woman')."""
    t = token.lower().replace(" ", "")
    m = re.fullmatch(r"(\d+)(girl|boy|woman|man|female|male|other)s?", t)
    if m:
        n, kind = m.group(1), m.group(2)
        if kind in ("girl", "woman", "female"):
            base = "girl"
        elif kind in ("boy", "man", "male"):
            base = "boy"
        else:
            base = "other"
        if n == "1":
            return "1" + base
        return n + ("girls" if base == "girl" else "boys" if base == "boy" else "others")
    if t.startswith("multiple girl"):
        return "multiple girls"
    if t.startswith("multiple boy"):
        return "multiple boys"
    if t.startswith("no human"):
        return "no humans"
    return token.lower()

QUALITY_TAGS = frozenset({
    "masterpiece", "best quality", "good quality", "normal quality",
    "low quality", "worst quality", "high quality",
    "score_9", "score_8", "score_7", "score_6", "score_5",
    "score_4", "score_3", "score_2", "score_1",
})

META_TAGS = frozenset({
    "highres", "absurdres", "incredibly absurdres", "ultra highres",
    "anime screencap", "official art", "game_cg", "visual novel",
    "light novel", "manga", "comic", "4koma", "doujinshi",
    "jpeg artifacts", "lineart", "no lineart", "sketch",
    "rough sketch", "clean sketch", "monochrome", "greyscale",
    "full color", "colored", "monochrome lineart",
})

SAFETY_TAGS = frozenset({"safe", "sensitive", "questionable", "nsfw", "explicit"})

YEAR_RE = re.compile(r'\byear\s+(\d{4})\b', re.IGNORECASE)
PERIOD_TAGS = frozenset({"newest", "recent", "mid", "early", "old"})

ARTIST_RE = re.compile(r'(?:@|artist:)([^,]+?)(?:\s*,|\s*$)', re.IGNORECASE)

COPYRIGHT_PAREN_RE = re.compile(r'\(([^)]+)\)')

# SD emphasis / LoRA tokens are extracted BEFORE parsing so that e.g.
# "(blue eyes:1.2)" is not misread as a booru "(series)" parenthetical,
# and "<lora:name:0.8>" survives the comma-splitter intact.
WEIGHT_TOKEN_RE = re.compile(
    r"\(\(([^()]+)\)\)"                           # ((tag)) strong emphasis
    r"|\(([^()]+):\s*(-?\+?\d+(?:\.\d+)?)\s*\)"   # (tag:1.2) explicit weight
)
LORA_RE = re.compile(r"<(?:lora|lyco):[^>]+>", re.IGNORECASE)

SERIES_IDENTIFIERS = frozenset({
    "vocaloid", "touhou", "genshin impact", "honkai star rail",
    "fate", "fate/grand order", "fgo", "azur lane", "kancolle",
    "kantai collection", "blue archive", "arknights", "uma musume",
    "love live", "idolmaster", "the idolmaster", "hololive",
    "nijisanji", "碧蓝航线", "原神", "starrail",
})


class ParsedPrompt:
    __slots__ = (
        "subject", "character", "series", "artists",
        "quality_tags", "meta_tags", "year_tag", "period_tag",
        "safety_tag", "general_tags", "nl_text", "weighted_tokens",
    )

    def __init__(self):
        self.subject: str = ""
        self.character: str = ""
        self.series: str = ""
        self.artists: list[str] = []
        self.quality_tags: list[str] = []
        self.meta_tags: list[str] = []
        self.year_tag: str = ""
        self.period_tag: str = ""
        self.safety_tag: str = ""
        self.general_tags: list[str] = []
        self.nl_text: str = ""
        # Verbatim SD tokens extracted before parsing: "(tag:1.2)", "((tag))",
        # "<lora:name:0.8>". They bypass tag processing and are appended to output.
        self.weighted_tokens: list[str] = []

    @property
    def has_booru_structure(self) -> bool:
        return bool(self.subject or self.quality_tags or self.general_tags)

    def __repr__(self):
        return (
            f"ParsedPrompt(subject={self.subject!r}, character={self.character!r}, "
            f"series={self.series!r}, artists={self.artists!r}, "
            f"year={self.year_tag!r}, safety={self.safety_tag!r}, "
            f"quality={self.quality_tags}, meta={self.meta_tags}, "
            f"general={len(self.general_tags)} tags, nl={self.nl_text!r})"
        )


def parse_prompt(raw: str) -> ParsedPrompt | None:
    if not raw or not raw.strip():
        return None

    result = ParsedPrompt()
    text = raw.strip()

    # Extract explicit weight / emphasis / LoRA tokens FIRST so later stages
    # (comma split, regex passes) cannot mangle or misclassify them.
    # They ride along as verbatim strings and are appended at render time.
    loras = LORA_RE.findall(text)
    text = LORA_RE.sub(" ", text)

    weighted: list[str] = []
    def _stash_weight(m: re.Match) -> str:
        if m.group(1) is not None:          # ((tag)) strong emphasis
            inner = m.group(1).strip()
            if inner:
                weighted.append(f"(({inner}))")
        else:                               # (tag:1.2) explicit weight
            inner, w = m.group(2).strip(), m.group(3)
            if inner:
                weighted.append(f"({inner}:{w})")
        return " "

    text = WEIGHT_TOKEN_RE.sub(_stash_weight, text)
    result.weighted_tokens = weighted + loras

    subject_m = SUBJECT_RE.search(text)
    if subject_m:
        result.subject = _normalize_subject(subject_m.group(1))
        text = text[:subject_m.start()] + text[subject_m.end():]

    for m in ARTIST_RE.finditer(text):
        name = m.group(1).strip()
        if name:
            result.artists.append(name)
    text = ARTIST_RE.sub("", text)

    paren_m = COPYRIGHT_PAREN_RE.search(text)
    if paren_m:
        inner = paren_m.group(1).strip()
        if inner and inner.lower() not in ("style", "medium", "artist", "parody"):
            before = text[:paren_m.start()].rstrip()
            if before:
                last_comma = before.rfind(",")
                char_candidate = before[last_comma + 1:].strip() if last_comma != -1 else before.strip()
                if char_candidate:
                    result.character = char_candidate
            result.series = inner
            text = text[:paren_m.start()] + text[paren_m.end():]

    year_m = YEAR_RE.search(text)
    if year_m:
        result.year_tag = f"year {year_m.group(1)}"
        text = text[:year_m.start()] + text[year_m.end():]

    parts = [p.strip() for p in text.split(",")]
    general = []
    for part in parts:
        if not part:
            continue
        low = part.lower().replace("_", " ").strip()
        low = re.sub(r"\s+", " ", low)
        # Normalize spaced score tags ("score 9" / "score 9 up") into quality.
        if low.startswith("score ") and low[6:].split()[0].isdigit():
            rest = low[6:].split()
            low = "score_" + rest[0] + ("_up" if len(rest) > 1 and rest[1] == "up" else "")
        if low in QUALITY_TAGS:
            result.quality_tags.append(low)
        elif low in META_TAGS:
            result.meta_tags.append(low)
        elif low in SAFETY_TAGS:
            result.safety_tag = low
        elif low == part.lower() and low in PERIOD_TAGS:
            # Guard against stealing common English words ("old church",
            # "early morning") that are split at comma but mean something else.
            # A single-word comma segment is far more likely to be a real
            # period tag than a word embedded mid-phrase.
            result.period_tag = low
        elif low in SERIES_IDENTIFIERS and not result.series:
            result.series = low
        else:
            general.append(part)

    result.general_tags = general

    # "solo" with no explicit count implies a single subject (usually 1girl).
    if not result.subject and "solo" in [g.lower().strip() for g in general]:
        result.subject = "1girl"

    if not result.subject and not result.general_tags and not result.quality_tags:
        result.nl_text = raw.strip()

    return result