"""BooruGacha Ban List — expand concepts into a full ban set."""

from ..data import flatten, parse_weighted_tag, tokenize


class BooruGachaBanList:
    """Takes user-specified concepts and expands them into a ``BANLIST``
    of related tags using token overlap and co-occurrence data.

    Expansion level controls how aggressively related tags are included:
    0 = exact matches only, higher = broader conceptual banning.
    """

    DESCRIPTION = (
        "Builds a ban list from concepts. Tags are normalised, then "
        "expanded via token overlap and co-occurrence statistics. "
        "Connect the output to the ban_list input of Search or Roll nodes."
    )

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "tags": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "",
                        "tooltip": (
                            "Comma-separated concepts to ban, "
                            "e.g. 'cum, feet, guro'"
                        ),
                    },
                ),
                "expansion": (
                    "FLOAT",
                    {
                        "default": 0.3,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.05,
                        "tooltip": (
                            "How aggressively to expand. "
                            "0 = exact only, 0.3 = token overlap + close "
                            "co-occurrences, 1.0 = broad conceptual ban."
                        ),
                    },
                ),
                "same_category": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": (
                            "Only expand to tags in the same category as "
                            "the seed tag. Prevents e.g. banning '3d' "
                            "(art_style) from also banning copyrights "
                            "like 'super_mario_3d_world'."
                        ),
                    },
                ),
                "content_filter": (
                    ["any", "sfw", "nudity", "sexual", "gore"],
                    {
                        "default": "any",
                        "tooltip": (
                            "Only expand into tags with this content "
                            "rating. 'any' = no restriction."
                        ),
                    },
                ),
            },
        }

    RETURN_TYPES = ("BANLIST", "STRING")
    RETURN_NAMES = ("ban_list", "debug_info")
    OUTPUT_TOOLTIPS = (
        "Set of banned tag names",
        "Expansion details per input concept",
    )
    FUNCTION = "build"
    CATEGORY = "BooruGacha"

    def build(
        self,
        tags: str,
        expansion: float,
        same_category: bool = True,
        content_filter: str = "any",
    ):
        from ..search import get_cooccurrence, get_engine

        engine = get_engine()
        registry = engine.registry
        cooccurrence = get_cooccurrence()

        sub_queries = [q.strip() for q in tags.split(",") if q.strip()]
        ban_set: set[str] = set()
        debug_lines: list[str] = []

        for sq in sub_queries:
            raw, _ = parse_weighted_tag(sq)

            # Normalise via search engine
            match = engine.find_one(raw, threshold=0.3)
            seed_tag = None
            if match:
                seed_db = match[0]
                idx = registry.danbooru_to_idx.get(seed_db)
                if idx is not None:
                    seed_tag = registry.tags[idx]
                seed_name = seed_tag["name"] if seed_tag else seed_db
            else:
                seed_name = raw.strip().lower().replace(" ", "_")

            # Determine category filter from seed tag
            seed_group = seed_tag["group"] if seed_tag else None
            seed_category = seed_tag.get("category") if seed_tag else None

            expanded = _expand_tag(
                seed_name,
                expansion,
                registry,
                cooccurrence,
                category_filter=(seed_group, seed_category)
                if same_category
                else None,
                filtering_filter={content_filter}
                if content_filter != "any"
                else None,
            )

            # Include both normalized and danbooru forms
            full = set()
            for name in expanded:
                full.add(name)
                ni = registry.name_to_idx.get(name)
                if ni is not None:
                    full.add(registry.tags[ni]["danbooru_name"])

            ban_set.update(full)

            added = sorted(expanded - {seed_name})
            debug_lines.append(
                f"{sq} -> {seed_name} (+{len(added)} expanded)"
            )
            for t in added:
                debug_lines.append(f"  + {t}")

        debug_lines.insert(0, f"Total banned: {len(ban_set)} tags")
        return (ban_set, "\n".join(debug_lines))


def _expand_tag(
    tag_name: str,
    expansion: float,
    registry,
    cooccurrence,
    category_filter: tuple[str | None, str | None] | None = None,
    filtering_filter: set[str] | None = None,
) -> set[str]:
    """Expand a single tag into related tags to ban.

    When *category_filter* is ``(group, category)``, only tags sharing
    the same fine-grained category (or group, for flat groups like
    artist/character) are included in the expansion.

    When *filtering_filter* is a set, only tags whose ``filtering``
    value is in the set are included.
    """
    banned: set[str] = {tag_name}

    if expansion <= 0:
        return banned

    def _allowed(idx: int) -> bool:
        cand = registry.tags[idx]
        if category_filter is not None:
            seed_group, seed_cat = category_filter
            if seed_cat and cand.get("category"):
                if cand["category"] != seed_cat:
                    return False
            elif cand["group"] != seed_group:
                return False
        if filtering_filter is not None:
            if cand.get("filtering", "sfw") not in filtering_filter:
                return False
        return True

    # --- Token containment: tags where all of the banned tag's tokens appear
    tag_tokens = set(tokenize(tag_name))
    if tag_tokens:
        candidate_sets = [
            registry.token_index.get(t, set()) for t in tag_tokens
        ]
        shared = candidate_sets[0].copy()
        for s in candidate_sets[1:]:
            shared &= s
        for idx in shared:
            if _allowed(idx):
                banned.add(registry.tags[idx]["name"])

    # --- Co-occurrence expansion (kicks in at expansion > 0.2)
    if expansion > 0.2:
        max_entries = int(expansion * 30)
        min_jaccard = max(0.005, 0.15 * (1.0 - expansion))

        related = cooccurrence.get_related(tag_name)
        count = 0
        for cooccur_tag, jaccard, _pmi, _lmi in related:
            if count >= max_entries:
                break
            if jaccard < min_jaccard:
                continue
            ni = registry.name_to_idx.get(cooccur_tag)
            if ni is not None and _allowed(ni):
                banned.add(cooccur_tag)
                count += 1

    return banned
