src/model_formatter.py

import hashlib

from src.prompt_parser import ParsedPrompt
from src.safety_filter import get_safety_tag
from src.dedup_engine import smart_dedup
from src.tag_categories import order_tags_booru
from src.tag_format import normalize_tag


def _dedup(tags: list[str]) -> list[str]:
    seen = set()
    out = []
    for t in tags:
        k = t.lower().strip()
        if k and k not in seen:
            seen.add(k)
            out.append(t)
    return out


def _is_score_tag(tag: str) -> bool:
    return tag.lower().strip().startswith("score_")


def _ensure_quality(quality_tags: list[str], model: str = "anima") -> list[str]:
    q = list(quality_tags)
    has_masterpiece = any("masterpiece" in x.lower() for x in q)
    has_best = any("best quality" in x.lower() for x in q)
    has_score = any("score_" in x.lower() for x in q)
    # Anima (circlestone-labs) uses score_7, NOT score_9 (which is Pony V6 XL).
    # Anima's official prefix: "masterpiece, best quality, score_7, safe,"
    score_tag = "score_7" if model == "anima" else "score_9"
    if not has_masterpiece:
        q.insert(0, "masterpiece")
    if not has_best:
        q.insert(1, "best quality")
    if not has_score:
        q.append(score_tag)
    return _dedup(q)


def _ensure_quality_illustrious(quality_tags: list[str]) -> list[str]:
    q = [t for t in quality_tags if not _is_score_tag(t)]
    has_masterpiece = any("masterpiece" in x.lower() for x in q)
    has_best = any("best quality" in x.lower() for x in q)
    has_score_7 = any("score_7" in x.lower() for x in (quality_tags or []))
    if not has_masterpiece:
        q.insert(0, "masterpiece")
    if not has_best:
        q.insert(1, "best quality")
    if not has_score_7:
        q.insert(2, "score_7")
    return _dedup(q)


QUALITY_WEIGHTS = {"masterpiece": 1.2, "best quality": 1.1}

LIGHT_IMPACT_TAGS = frozenset({
    "bloom", "depth of field", "motion blur", "film grain", "chromatic aberration",
    "vignette", "lens flare", "glow", "soft focus", "sharp focus", "bokeh",
    "dramatic lighting", "cinematic lighting", "dramatic shadows", "rim lighting",
    "backlighting", "crepuscular rays", "god rays", "light particles", "dust motes",
    "fireflies", "glowing", "neon lighting", "studio lighting", "hard lighting",
    "soft lighting", "natural lighting", "volumetric lighting", "foggy", "misty",
    "particles", "fire", "water", "electricity", "explosion", "magic", "sparkles",
    "reflection", "refraction", "iridescent", "metallic", "glass", "translucent",
    "dynamic lighting", "ambient occlusion", "global illumination", "ray tracing",
    "subsurface scattering", "caustics", "luminance", "hdr", "tone mapping",
    "dramatic", "haunting", "ethereal", "serene", "mysterious", "tranquil",
    "whimsical", "melancholic", "romantic", "ominous", "foreboding",
    "post-apocalyptic", "dark fantasy", "cosmic horror", "cyberpunk",
    "dreamy", "vibrant", "pastel colors", "neon palette", "high contrast",
    "professional photography", "award winning", "masterpiece", "perfect anatomy",
    "incredibly absurdres", "extremely detailed", "intricate", "ornate",
    "perfect face", "detailed eyes", "beautiful detailed eyes",
})


def _hash_weight(tag: str, light: bool = False) -> float:
    # Stable across processes: hashlib is not affected by PYTHONHASHSEED,
    # unlike builtin hash() which is randomised per process for str.
    h = int.from_bytes(hashlib.blake2b(tag.lower().encode("utf-8"), digest_size=4).digest(), "big")
    if light:
        return round(1.0 + (h % 16) / 100.0, 2)
    return round(0.9 + (h % 21) / 100.0, 2)


def _wrap_tag(tag: str, weight: float) -> str:
    if weight == 1.0:
        return tag
    return f"({tag}:{weight})"


def _join_tags(tags: list[str]) -> str:
    return ", ".join(tag for tag in tags if tag)


# Adjective pairs that must never be merged (avoids nonsense like "long short hair").
_OPPOSITE_ADJ = {
    "long": "short", "short": "long",
    "big": "small", "small": "big",
    "wide": "narrow", "narrow": "wide",
    "warm": "cool", "cool": "warm",
    "hard": "soft", "soft": "hard",
    "light": "dark", "dark": "light",
    "high": "low", "low": "high",
    "open": "closed", "closed": "open",
}


# English natural adjective order (OSASCOMP-ish), as a position-rank lookup.
# When combining "blue pleated skirt"-style tags we emit adjectives in this
# order so the merged phrase sounds natural ("long white pleated skirt").
_ADJ_ORDER = {
    "opinion": 0,   # beautiful, gorgeous, ugly, nice...
    "size": 1,      # big, huge, long, short, tall... (quality tags like "beautiful" too)
    "age": 2,       # old, new, ancient, young...
    "shape": 3,     # round, square, wavy, straight, pleated...
    "color": 4,     # blue, red, white, golden...
    "origin": 5,    # japanese, gothic, medieval...
    "material": 6,  # silk, leather, metallic, glass...
    "purpose": 7,   # wedding (dress), school (uniform)...
}

_ADJ_COLOR_SET = {
    "white", "black", "red", "blue", "green", "yellow", "purple", "pink",
    "orange", "brown", "gray", "grey", "silver", "golden", "gold", "blonde",
    "aqua", "cyan", "magenta", "violet", "scarlet", "crimson", "azure",
    "multicolored", "pastel", "neon", "rainbow", "dark", "light", "pale",
}
_ADJ_SIZE_SET = {
    "big", "small", "huge", "tiny", "large", "short", "long", "tall", "wide",
    "narrow", "thick", "thin", "absurdly", "gigantic", "micro", "mini",
}
_ADJ_AGE_SET = {"old", "new", "ancient", "young", "vintage", "retro", "modern", "futuristic"}
_ADJ_SHAPE_SET = {
    "wavy", "straight", "curly", "pleated", "round", "square", "messy",
    "spiky", "twintails", "ponytail", "braided", "drill", "ruffled",
    "detailed", "intricate", "ornate",
}
_ADJ_MATERIAL_SET = {
    "silk", "leather", "cotton", "wool", "lace", "denim", "metallic", "metal",
    "glass", "crystal", "wooden", "plastic", "fur", "satin", "velvet",
}


_QUALITY_KEYWORDS = {"best", "good", "high", "great", "amazing", "beautiful",
                    "detailed", "aesthetic", "stunning", "perfect", "pretty",
                    "gorgeous", "lovely", "cute", "elegant"}

_QUAL_ORDER = _ADJ_ORDER["opinion"]
_SIZE_ORDER = _ADJ_ORDER["size"]


def _adj_position(word: str) -> tuple[int, str]:
    """Return sort key (bucket, word) so ties preserve discovery order."""
    if word in _ADJ_SIZE_SET:
        return (_ADJ_ORDER["size"], word)
    if word in _QUALITY_KEYWORDS:
        return (_ADJ_ORDER["opinion"], word)
    if word in _ADJ_AGE_SET:
        return (_ADJ_ORDER["age"], word)
    if word in _ADJ_SHAPE_SET:
        return (_ADJ_ORDER["shape"], word)
    if word in _ADJ_COLOR_SET:
        return (_ADJ_ORDER["color"], word)
    if word in _ADJ_MATERIAL_SET:
        return (_ADJ_ORDER["material"], word)
    return (_ADJ_ORDER["opinion"], word)


def _merge_combinable_tags(tags: list[str]) -> list[str]:
    """Combine adjectival+noun tags that share a noun (Booru Prompt Gallery's
    'Combine tags' feature), e.g. 'skirt', 'white skirt', 'pleated skirt'
    becomes 'white pleated skirt'. Skips groups whose adjectives are opposites."""
    groups: dict[str, list[str]] = {}
    for t in tags:
        parts = t.split()
        if parts:
            groups.setdefault(parts[-1], []).append(t)
    if not groups:
        return tags
    result = list(tags)
    for noun, group in groups.items():
        if len(group) < 2:
            continue
        # Only combine when a bare noun is present (e.g. "skirt" + "white skirt"),
        # otherwise two different adjectival forms (e.g. "blue eyes" + "red eyes")
        # would collapse into nonsense like "blue red eyes".
        if not any(len(t.split()) == 1 for t in group):
            continue
        adj_sets = [set(t.split()[:-1]) for t in group]
        conflict = False
        for i in range(len(group)):
            for j in range(i + 1, len(group)):
                for a in adj_sets[i]:
                    if _OPPOSITE_ADJ.get(a) in adj_sets[j]:
                        conflict = True
                        break
                if conflict:
                    break
            if conflict:
                break
        if conflict:
            continue
        all_adjs: list[str] = []
        for adjs in adj_sets:
            for a in adjs:
                if a not in all_adjs:
                    all_adjs.append(a)
        if not all_adjs:
            continue
        all_adjs.sort(key=_adj_position)
        merged = " ".join(all_adjs + [noun])
        for t in group:
            if t in result:
                result.remove(t)
        result.append(merged)
    return result


def _dedup_tokens(tokens: list[str]) -> list[str]:
    """Dedup rendered tags across sections by normalized key (weight wrapper ignored)."""
    seen = set()
    out = []
    for t in tokens:
        if not t:
            continue
        key = t.lower().strip()
        if key.startswith("(") and key.endswith(")") and ":" in key:
            key = key[1:-1].rsplit(":", 1)[0].strip()
        if key and key not in seen:
            seen.add(key)
            out.append(t)
    return out


def _flatten_sections(sections: list[str]) -> list[str]:
    tokens: list[str] = []
    for s in sections:
        if not s:
            continue
        tokens.extend(part.strip() for part in s.split(",") if part.strip())
    return tokens


def _apply_quality_weights(tag: str, weight_mode: str) -> str:
    if weight_mode == "off":
        return tag
    w = QUALITY_WEIGHTS.get(tag.lower())
    if w and weight_mode in ("light", "on"):
        return _wrap_tag(tag, w)
    return tag


def _tag_weight(tag: str, weight_mode: str, tag_weights: dict | None) -> float:
    if weight_mode == "off":
        return 1.0
    if tag_weights:
        w = tag_weights.get(tag.lower())
        if w is not None:
            return w
    if weight_mode == "on":
        return _hash_weight(tag, light=False)
    if weight_mode == "light":
        tl = tag.lower()
        if tl in LIGHT_IMPACT_TAGS:
            return _hash_weight(tag, light=True)
        return 1.0
    return 1.0


def _format_artist_entry(artist_name: str, tag_weights: dict | None) -> str:
    entry = f"artist:{artist_name}"
    if tag_weights:
        w = tag_weights.get(artist_name.lower(), 1.0)
        if w != 1.0:
            entry = f"(artist:{artist_name}:{w})"
    return entry


def format_anima(
    parsed: ParsedPrompt,
    rating: str = "pg",
    quality_enabled: bool = True,
    weight_mode: str = "off",
    tag_weights: dict | None = None,
    model: str = "anima",
    output_format: str = "prompt",
) -> str:
    safety = get_safety_tag(rating)
    sections = []

    if quality_enabled:
        quality = _ensure_quality(parsed.quality_tags, model=model)
        if weight_mode in ("light", "on"):
            quality = [_apply_quality_weights(t, weight_mode) for t in quality]
        sections.append(_join_tags(quality))
    else:
        quality = list(parsed.quality_tags) if parsed.quality_tags else []

    if parsed.meta_tags:
        meta = _dedup(parsed.meta_tags)
        meta = [_wrap_tag(t, _tag_weight(t, weight_mode, tag_weights)) for t in meta]
        sections.append(_join_tags(meta))

    if parsed.year_tag:
        sections.append(parsed.year_tag)
    if parsed.period_tag:
        sections.append(parsed.period_tag)

    sections.append(safety)

    if parsed.subject:
        subj = parsed.subject
        w = _tag_weight(subj, weight_mode, tag_weights)
        if w != 1.0:
            subj = _wrap_tag(subj, w)
        sections.append(subj)
    if parsed.character:
        char = parsed.character
        w = _tag_weight(char, weight_mode, tag_weights)
        if w != 1.0:
            char = _wrap_tag(char, w)
        sections.append(char)
    if parsed.series:
        sec = parsed.series
        w = _tag_weight(sec, weight_mode, tag_weights)
        if w != 1.0:
            sec = _wrap_tag(sec, w)
        sections.append(sec)
    if parsed.artists:
        artists_out = [_format_artist_entry(a, tag_weights) for a in parsed.artists]
        sections.append(_join_tags(artists_out))

    if parsed.general_tags:
        sections.append("BREAK")

    all_general = _merge_combinable_tags(parsed.general_tags)
    all_general = smart_dedup(all_general, model="anima")
    quality_lower = {x.lower() for x in quality}
    safety_lower = {safety}
    meta_lower = {x.lower() for x in (parsed.meta_tags or [])}
    subject_lower = {(parsed.subject or "").lower()}
    char_lower = {(parsed.character or "").lower()}
    series_lower = {(parsed.series or "").lower()}
    skip = quality_lower | safety_lower | meta_lower | subject_lower | char_lower | series_lower
    all_general = [g for g in all_general if g.lower() not in skip and not _is_score_tag(g)]

    if all_general:
        all_general = order_tags_booru(all_general)
        all_general = [_wrap_tag(t, _tag_weight(t, weight_mode, tag_weights)) for t in all_general]
        sections.append(_join_tags(all_general))

    if parsed.nl_text:
        sections.append(parsed.nl_text)

    if getattr(parsed, "weighted_tokens", None):
        sections.extend(parsed.weighted_tokens)

    tokens = _dedup_tokens(_flatten_sections([s for s in sections if s]))
    tokens = [normalize_tag(tok, output_format) for tok in tokens]
    return ", ".join(tokens)


def format_illustrious(
    parsed: ParsedPrompt,
    rating: str = "pg",
    quality_enabled: bool = True,
    weight_mode: str = "off",
    tag_weights: dict | None = None,
    output_format: str = "prompt",
) -> str:
    rating_map = {
        "pg": "general", "pg13": "general", "pg16": "sensitive",
        "r": "nsfw", "r+": "explicit",
    }
    illustrious_rating = rating_map.get(rating, "general")

    if quality_enabled:
        quality = _ensure_quality_illustrious(parsed.quality_tags)
    else:
        quality = [t for t in (parsed.quality_tags or []) if not _is_score_tag(t)]
    quality = _dedup(quality)
    if weight_mode in ("light", "on"):
        quality = [_apply_quality_weights(t, weight_mode) for t in quality]

    ordered = []
    if quality:
        ordered.extend(quality)
    if parsed.subject:
        subj = parsed.subject
        w = _tag_weight(subj, weight_mode, tag_weights)
        if w != 1.0:
            subj = _wrap_tag(subj, w)
        ordered.append(subj)
    if parsed.character:
        char = parsed.character
        w = _tag_weight(char, weight_mode, tag_weights)
        if w != 1.0:
            char = _wrap_tag(char, w)
        ordered.append(char)
    if parsed.series:
        sec = parsed.series
        w = _tag_weight(sec, weight_mode, tag_weights)
        if w != 1.0:
            sec = _wrap_tag(sec, w)
        ordered.append(sec)
    if parsed.artists:
        artists_out = [_format_artist_entry(a, tag_weights) for a in parsed.artists]
        ordered.extend(artists_out)

    if parsed.general_tags:
        ordered.append("BREAK")

    all_general = _merge_combinable_tags(parsed.general_tags)
    all_general = smart_dedup(all_general, model="illustrious")
    quality_lower = {x.lower() for x in quality}
    subject_lower = {(parsed.subject or "").lower()}
    char_lower = {(parsed.character or "").lower()}
    series_lower = {(parsed.series or "").lower()}
    skip = quality_lower | subject_lower | char_lower | series_lower
    all_general = [g for g in all_general if g.lower() not in skip and not _is_score_tag(g)]
    all_general = order_tags_booru(all_general)
    all_general = [_wrap_tag(t, _tag_weight(t, weight_mode, tag_weights)) for t in all_general]
    ordered.extend(all_general)

    if parsed.nl_text:
        ordered.append(parsed.nl_text)

    if getattr(parsed, "weighted_tokens", None):
        ordered.extend(parsed.weighted_tokens)

    tokens = [normalize_tag(t, output_format) for t in _dedup_tokens(ordered)]
    tag_str = ", ".join(tokens)
    return f"rating:{illustrious_rating}, {tag_str}" if tag_str else f"rating:{illustrious_rating}"


_NOOBAI_RATING_MAP = {
    "pg": "general", "pg13": "general", "pg16": "sensitive",
    "r": "nsfw", "r+": "explicit",
}


def format_noobai(
    parsed: ParsedPrompt,
    rating: str = "pg",
    quality_enabled: bool = True,
    weight_mode: str = "off",
    tag_weights: dict | None = None,
    output_format: str = "prompt",
) -> str:
    """NoobAI-XL puts the aesthetic score tags around the quality block:
        masterpiece, best quality, very aesthetic, absurdres
    followed by subject and remaining tags. No explicit score_N; it uses
    ``very aesthetic``-style boosters instead of the Pony score scale and keeps
    the rating tag separate from quality.
    """
    noobai_rating = _NOOBAI_RATING_MAP.get(rating, "general")

    if quality_enabled:
        quality = [t for t in (parsed.quality_tags or []) if not _is_score_tag(t)]
        # Desired block order: masterpiece, best quality, very aesthetic, absurdres
        wanted = [
            ("masterpiece", None),
            ("best quality", None),
            ("very aesthetic", None),
            ("absurdres", "meta"),
        ]
        block: list[str] = []
        for phrase, _ in wanted:
            if not any(phrase in x.lower() for x in block):
                block.append(phrase)
        # user-supplied extra quality tags keep tail position
        for t in quality:
            if not any(t.lower() in b.lower() or b.lower() in t.lower() for b in block):
                block.append(t)
        quality = _dedup(block)
    else:
        quality = [t for t in (parsed.quality_tags or []) if not _is_score_tag(t)]

    if weight_mode in ("light", "on"):
        quality = [_apply_quality_weights(t, weight_mode) for t in quality]

    ordered: list[str] = []
    if quality:
        ordered.extend(quality)
    if parsed.subject:
        subj = parsed.subject
        w = _tag_weight(subj, weight_mode, tag_weights)
        if w != 1.0:
            subj = _wrap_tag(subj, w)
        ordered.append(subj)
    if parsed.character:
        char = parsed.character
        w = _tag_weight(char, weight_mode, tag_weights)
        if w != 1.0:
            char = _wrap_tag(char, w)
        ordered.append(char)
    if parsed.series:
        sec = parsed.series
        w = _tag_weight(sec, weight_mode, tag_weights)
        if w != 1.0:
            sec = _wrap_tag(sec, w)
        ordered.append(sec)
    if parsed.artists:
        ordered.extend(_format_artist_entry(a, tag_weights) for a in parsed.artists)

    if parsed.general_tags:
        ordered.append("BREAK")

    all_general = _merge_combinable_tags(parsed.general_tags)
    all_general = smart_dedup(all_general, model="illustrious")
    quality_lower = {x.lower() for x in quality}
    skip = quality_lower | {
        (parsed.subject or "").lower(),
        (parsed.character or "").lower(),
        (parsed.series or "").lower(),
    }
    all_general = [g for g in all_general if g.lower() not in skip and not _is_score_tag(g)]
    all_general = order_tags_booru(all_general)
    all_general = [_wrap_tag(t, _tag_weight(t, weight_mode, tag_weights)) for t in all_general]
    ordered.extend(all_general)

    if parsed.nl_text:
        ordered.append(parsed.nl_text)
    if getattr(parsed, "weighted_tokens", None):
        ordered.extend(parsed.weighted_tokens)

    tokens = [normalize_tag(t, output_format) for t in _dedup_tokens(ordered)]
    tag_str = ", ".join(tokens)
    return f"rating:{noobai_rating}, {tag_str}" if tag_str else f"rating:{noobai_rating}"


def format_prompt(
    parsed: ParsedPrompt,
    model: str = "anima",
    rating: str = "pg",
    quality_enabled: bool = True,
    weight_mode: str = "off",
    tag_weights: dict | None = None,
    output_format: str = "prompt",
) -> str:
    if model == "noobai":
        return format_noobai(parsed, rating, quality_enabled=quality_enabled, weight_mode=weight_mode, tag_weights=tag_weights, output_format=output_format)
    if model == "illustrious":
        return format_illustrious(parsed, rating, quality_enabled=quality_enabled, weight_mode=weight_mode, tag_weights=tag_weights, output_format=output_format)
    return format_anima(parsed, rating, quality_enabled=quality_enabled, weight_mode=weight_mode, tag_weights=tag_weights, model=model, output_format=output_format)