src/tag_searcher.py

import functools
import json
import os

_COOCCURRENCE_PATH = os.path.join(
    os.path.dirname(os.path.dirname(__file__)),
    "data", "tag_cooccurrence.json",
)
_COOCURRENCE_DATA = None

DANBOORU_TAG_API = "https://danbooru.donmai.us/tags.json"
_REQUEST_TIMEOUT = 8

# relation weights: curated/semantic connections beat co-category noise.
_RELATION_BONUS = {
    "character_tag": 3.0,
    "common_pair": 2.0,
    "synonym": 1.5,
    "modifier": 1.2,
    "co_category": 1.0,
}


def _load_cooccurrence() -> dict:
    global _COOCURRENCE_DATA
    if _COOCURRENCE_DATA is not None:
        return _COOCURRENCE_DATA
    try:
        with open(_COOCCURRENCE_PATH, "r", encoding="utf-8") as f:
            _COOCURRENCE_DATA = json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        _COOCURRENCE_DATA = {"cooccurrence": {}}
    return _COOCURRENCE_DATA


def reload_cooccurrence():
    global _COOCURRENCE_DATA
    _COOCURRENCE_DATA = None
    _load_cooccurrence()


def search_local_tags(query: str, warehouse, limit: int = 30) -> list[dict]:
    q = query.lower().strip()
    if not q:
        return []
    results = []
    seen = set()
    for cat_name in warehouse.pools:
        pool = warehouse.get_pool(cat_name)
        if pool is None:
            continue
        for tag_dict in pool.get_all_tag_dicts():
            tag = tag_dict.get("tag", "")
            tl = tag.lower()
            if q in tl and tl not in seen:
                # Rank: exact match > starts-with > contains; shorter tags win ties.
                if tl == q:
                    rank = 0
                elif tl.startswith(q):
                    rank = 1
                else:
                    rank = 2
                results.append({
                    "tag": tag,
                    "category": cat_name,
                    "rating": tag_dict.get("rating", "sfw"),
                    "source": "local",
                    "_rank": (rank, len(tl)),
                })
                seen.add(tl)
    results.sort(key=lambda x: x["_rank"])
    out = []
    for r in results[:limit]:
        r.pop("_rank", None)
        out.append(r)
    return out


@functools.lru_cache(maxsize=512)
def search_danbooru_tags(query: str, limit: int = 20) -> list[dict]:
    import urllib.request
    import urllib.parse
    import urllib.error
    q = query.lower().strip()
    if not q:
        return []
    url = f"{DANBOORU_TAG_API}?search[name_or_name_matches]={urllib.parse.quote(q)}&limit={limit}"
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Whyx-PROmpTea/1.0"})
        with urllib.request.urlopen(req, timeout=_REQUEST_TIMEOUT) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        results = []
        for item in data:
            results.append({
                "tag": item.get("name", ""),
                "category": item.get("category", 0),
                "post_count": item.get("post_count", 0),
                "source": "danbooru",
            })
        results.sort(key=lambda x: -x.get("post_count", 0))
        return results[:limit]
    except (urllib.error.URLError, urllib.error.HTTPError, OSError, json.JSONDecodeError):
        return []


def search_tags(query: str, warehouse, limit: int = 30) -> dict:
    local = search_local_tags(query, warehouse, limit)
    web = search_danbooru_tags(query, max(5, limit // 2))
    return {
        "local": local,
        "web": web,
    }


def get_cooccurrence_tags(tag: str, limit: int = 10) -> list[dict]:
    """Co-occurring tags for ``tag``, strongest first.

    Related entries are scored by base weight × relation bonus so that
    curated/semantic links (synonym, common_pair, character_tag) outrank raw
    co-category noise.
    """
    co = _load_cooccurrence().get("cooccurrence", {})
    tl = tag.lower().strip()
    related = co.get(tl, [])
    scored = sorted(
        related,
        key=lambda r: -(
            float(r.get("weight", 1.0)) * _RELATION_BONUS.get(r.get("relation", "co_category"), 1.0)
        ),
    )
    return scored[:limit]


HUMAN_SUBJECT_TAGS = frozenset({
    "1girl", "1boy", "solo", "female", "male", "human", "person",
})

SUBJECT_COUNT_CONFLICTS = {
    "1girl": ["2girls", "3girls", "multiple_girls", "1boy", "2boys", "3boys",
              "multiple_boys", "male_focus", "no_humans", "group"],
    "2girls": ["1girl", "solo", "no_humans"],
    "3girls": ["1girl", "solo", "no_humans"],
    "1boy": ["2girls", "3girls", "multiple_girls", "female_focus", "1girl",
             "no_humans"],
    "2boys": ["1boy", "solo", "no_humans"],
    "3boys": ["1boy", "solo", "no_humans"],
    "solo": ["2girls", "3girls", "multiple_girls", "2boys", "3boys", "group",
             "multiple_boys"],
    "multiple_girls": ["1girl", "solo"],
    "multiple_boys": ["1boy", "solo"],
    "no_humans": ["1girl", "2girls", "3girls", "1boy", "2boys", "3boys",
                  "solo", "group", "male_focus", "female_focus"],
    "male_focus": ["1girl", "female_focus"],
    "female_focus": ["1boy", "male_focus"],
    "group": ["1girl", "1boy", "solo"],
}

# Canonicalize to spaced form: the pipeline works in spaced canonical tags
# (e.g. "multiple girls", "male focus"), so the conflict keys/values must
# match. Build a normalized copy that replaces underscores with spaces.
SUBJECT_COUNT_CONFLICTS = {
    k.replace("_", " "): [v.replace("_", " ") for v in vals]
    for k, vals in SUBJECT_COUNT_CONFLICTS.items()
}


def has_human_subject(user_tags: list[str]) -> bool:
    if not user_tags:
        return False
    return bool({t.lower().strip() for t in user_tags} & HUMAN_SUBJECT_TAGS)


def filter_subject_conflicts(tags: list[str], user_tags: list[str] | None = None) -> list[str]:
    protect_src = user_tags if user_tags is not None else tags
    protect_lower = {t.lower().strip().replace("_", " ") for t in protect_src}
    blocked = set()
    for tag in protect_lower:
        if tag in SUBJECT_COUNT_CONFLICTS:
            blocked.update(SUBJECT_COUNT_CONFLICTS[tag])
    if not blocked:
        return tags
    return [t for t in tags if t.lower().strip().replace("_", " ") not in blocked]


def enrich_prompt_tags(
    parsed,
    warehouse,
    max_tags: int = 5,
    user_tags: list[str] | None = None,
) -> list[str]:
    enrich_tags = []
    context_sources = []
    if parsed.subject:
        context_sources.append(parsed.subject)
    if parsed.character:
        context_sources.append(parsed.character)
    if parsed.series:
        context_sources.append(parsed.series)
    if parsed.general_tags:
        context_sources.extend(parsed.general_tags[:5])
    seen = set()
    for src in context_sources:
        src_lower = src.lower().strip()
        if src_lower in seen:
            continue
        seen.add(src_lower)
        related = get_cooccurrence_tags(src, limit=5)
        for item in related:
            tag = item.get("tag", "")
            if tag.lower() not in seen:
                enrich_tags.append(tag)
                seen.add(tag.lower())
    enrich_tags = enrich_tags[:max_tags]
    if len(enrich_tags) < max_tags:
        web_results = search_danbooru_tags(context_sources[0] if context_sources else "", limit=10)
        for item in web_results:
            tag = item.get("tag", "")
            if tag.lower() not in seen:
                enrich_tags.append(tag)
                seen.add(tag.lower())
                if len(enrich_tags) >= max_tags:
                    break
    enrich_tags = filter_subject_conflicts(enrich_tags, user_tags)
    return enrich_tags[:max_tags]


def build_COOCURRENCE_DATA(warehouse) -> dict:
    co = {}
    for cat_name in warehouse.pools:
        pool = warehouse.get_pool(cat_name)
        if pool is None:
            continue
        tags_in_cat = pool.get_all_tags(max_rating="explicit")
        for tag in tags_in_cat:
            tl = tag.lower().strip()
            if tl not in co:
                co[tl] = []
            related = [t for t in tags_in_cat if t.lower().strip() != tl]
            for rt in related[:15]:
                co[tl].append({"tag": rt, "relation": "co_category", "weight": 1.0})
    return {"cooccurrence": co}