"""BooruGacha Tag Expand — expand tags into related co-occurring tags."""

from ..data import format_weighted_tag, parse_weighted_tag, tokenize

SCORING_MODES = {
    "Similarity (Jaccard)": 1,
    "Association (PMI)": 2,
    "Weighted (LMI)": 3,
}

SCOPE_OPTIONS = ["any", "same group", "same subgroup", "same category"]


class BooruGachaTagExpand:
    """Expands input tags into related tags via token overlap and
    co-occurrence data.

    Token overlap provides the conceptual link (``catgirl`` →
    ``cat_ears``), while co-occurrence ranks and supplements with
    contextually relevant tags.  Weight decays by rank.
    """

    DESCRIPTION = (
        "Expands each input tag into related tags. "
        "Useful for reinforcing a concept — e.g. 'catgirl' expands "
        "into cat_ears, cat_tail, etc. Weight decays by rank."
    )

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "tags": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "",
                        "tooltip": (
                            "Comma-separated tags to expand, e.g. "
                            "'catgirl, beach'"
                        ),
                    },
                ),
                "max_per_tag": (
                    "INT",
                    {
                        "default": 5,
                        "min": 1,
                        "max": 30,
                        "tooltip": "Maximum number of expansion tags per input tag.",
                    },
                ),
                "weight": (
                    "FLOAT",
                    {
                        "default": 0.8,
                        "min": 0.0,
                        "max": 2.0,
                        "step": 0.05,
                        "tooltip": (
                            "Weight for the top-ranked expansion tag. "
                            "Lower-ranked tags decay to ~30% of this value."
                        ),
                    },
                ),
                "scoring": (
                    list(SCORING_MODES.keys()),
                    {
                        "default": "Similarity (Jaccard)",
                        "tooltip": (
                            "Similarity: tags that commonly co-occur. "
                            "Association: strongly linked even if rare. "
                            "Weighted: balances association with prevalence."
                        ),
                    },
                ),
                "scope": (
                    SCOPE_OPTIONS,
                    {
                        "default": "any",
                        "tooltip": (
                            "Restrict expansions to the same group, "
                            "subgroup, or category as the input tag."
                        ),
                    },
                ),
                "content_filter": (
                    ["all", "sfw", "sfw+nudity", "sfw+nudity+sexual"],
                    {
                        "default": "all",
                        "tooltip": "Filter expanded tags by content rating.",
                    },
                ),
            },
            "optional": {
                "ban_list": (
                    "BANLIST",
                    {"tooltip": "Ban list to exclude from expansion results."},
                ),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("tags", "debug_info")
    OUTPUT_TOOLTIPS = (
        "Original + expanded tags, comma-separated",
        "Per-tag expansion details",
    )
    FUNCTION = "expand"
    CATEGORY = "BooruGacha"

    def expand(
        self,
        tags: str,
        max_per_tag: int,
        weight: float,
        scoring: str,
        scope: str,
        content_filter: str = "all",
        ban_list: set | None = None,
    ):
        from ..search import CONTENT_FILTER_SETS, get_cooccurrence, get_engine

        engine = get_engine()
        registry = engine.registry
        cooccurrence = get_cooccurrence()
        ban = ban_list or set()
        allowed = CONTENT_FILTER_SETS.get(content_filter)
        score_idx = SCORING_MODES[scoring]

        sub_queries = [q.strip() for q in tags.split(",") if q.strip()]

        # --- Pass 0: resolve all input tags ------------------------------
        resolved: list[tuple[str, str, float | None, str, dict | None, str, str]] = []
        for sq in sub_queries:
            raw, orig_weight = parse_weighted_tag(sq)
            match = engine.find_one(raw, threshold=0.3)
            if match:
                db_name, label, _score = match
                idx = registry.danbooru_to_idx.get(db_name)
                seed_tag = registry.tags[idx] if idx is not None else None
            else:
                db_name = raw.strip().lower().replace(" ", "_")
                seed_tag = None
                label = "unresolved"
            lookup_name = seed_tag["name"] if seed_tag else db_name
            resolved.append((sq, raw, orig_weight, db_name, seed_tag, lookup_name, label))

        # --- Build cross-input conflict detector -------------------------
        # For each original with a category, cache its co-occurrence
        # neighbourhood.  An expansion that shares a category with another
        # original but doesn't appear in its co-occurrence set is likely
        # conflicting (e.g. light_skinned expanded from smooth_skin
        # conflicts with the user's dark_skin input).
        conflict_detector = _build_conflict_detector(
            resolved, cooccurrence, registry,
        )

        # --- Pass 1-2: expand each tag -----------------------------------
        out_tags: list[str] = []
        debug_lines: list[str] = []
        seen: set[str] = set()

        for sq, raw, orig_weight, db_name, seed_tag, lookup_name, label in resolved:
            # Always include the original tag (user-typed)
            orig_out = format_weighted_tag(db_name, orig_weight)
            if db_name not in seen:
                out_tags.append(orig_out)
                seen.add(db_name)

            # Skip expansion for unresolved tags — we don't know what
            # the user meant, so co-occurrence would just be noise.
            if seed_tag is None:
                debug_lines.append(f"{sq} -> {db_name} [{label}] (skipped)")
                continue

            # Build co-occurrence lookup dict for ranking
            related = cooccurrence.get_related(lookup_name, limit=500)
            cooccur_map: dict[str, tuple[float, float, float]] = {}
            for ct, jac, pmi, lmi in related:
                cooccur_map[ct] = (jac, pmi, lmi)

            # --- Phase 1: token overlap candidates -----------------------
            seed_tokens = set(tokenize(lookup_name))
            token_candidates = _token_overlap_candidates(
                seed_tokens,
                lookup_name,
                db_name,
                registry,
                cooccur_map,
                score_idx,
                seed_tag,
                scope,
                allowed,
                ban,
                seen,
            )

            # --- Phase 2: co-occurrence fill -----------------------------
            cooccur_only = _cooccurrence_candidates(
                related,
                score_idx,
                registry,
                seed_tag,
                scope,
                allowed,
                ban,
                seen,
                {c[0] for c in token_candidates},
            )

            # Merge: co-occurrence-backed token candidates first, then
            # pure co-occurrence fill.  Zero-score token-only candidates
            # are dropped (coincidental overlap, not conceptual).
            # Cross-input conflict detection filters both pools — better
            # to under-expand (user can add tags) than to produce
            # contradictory expansions the user has to clean up.
            backed = [(d, m, "token") for d, m in token_candidates if m > 0]
            cooccur = [(d, m, "cooccur") for d, m in cooccur_only]

            expansions: list[tuple[str, float, str]] = []
            for pool in (backed, cooccur):
                for cand_db, metric_val, src in pool:
                    if len(expansions) >= max_per_tag:
                        break
                    if cand_db in seen:
                        continue
                    if conflict_detector(cand_db, lookup_name, registry):
                        continue
                    expansions.append((cand_db, metric_val, src))
                    seen.add(cand_db)

            # Apply decaying weights and add to output
            n = len(expansions)
            for rank, (exp_db, _m, _src) in enumerate(expansions):
                decay = 1.0 - 0.7 * rank / max(n - 1, 1) if n > 1 else 1.0
                eff_weight = weight * decay
                out_tags.append(format_weighted_tag(exp_db, eff_weight))

            # Debug
            debug_lines.append(
                f"{sq} -> {db_name} [{label}] (+{n} expanded)"
            )
            for rank, (exp_db, metric_val, src) in enumerate(expansions):
                decay = 1.0 - 0.7 * rank / max(n - 1, 1) if n > 1 else 1.0
                eff_weight = weight * decay
                debug_lines.append(
                    f"  + {format_weighted_tag(exp_db, eff_weight)}"
                    f" ({src}, score: {metric_val:.4f})"
                )

        return (", ".join(out_tags), "\n".join(debug_lines))


def _build_conflict_detector(resolved, cooccurrence, registry):
    """Return a callable that checks cross-input conflicts.

    An expansion conflicts when it shares a category with another
    original input tag but doesn't appear in that input's co-occurrence
    neighbourhood — meaning they rarely appear together in practice
    (e.g. ``light_skinned`` conflicts with user-supplied ``dark_skin``).
    """
    # Map: category → list of (lookup_name, co-occurrence name set)
    cat_to_originals: dict[str, list[tuple[str, set[str]]]] = {}
    for _sq, _raw, _ow, _db, seed_tag, lookup_name, _label in resolved:
        if seed_tag is None:
            continue
        cat = seed_tag.get("category")
        if not cat:
            continue
        related = cooccurrence.get_related(lookup_name, limit=500)
        cooccur_names = {ct for ct, *_ in related}
        cat_to_originals.setdefault(cat, []).append((lookup_name, cooccur_names))

    def check(cand_db: str, from_seed: str, reg) -> bool:
        """Return True if *cand_db* conflicts with another original."""
        idx = reg.danbooru_to_idx.get(cand_db)
        if idx is None:
            return False
        cand_cat = reg.tags[idx].get("category")
        if not cand_cat or cand_cat not in cat_to_originals:
            return False
        cand_name = reg.tags[idx]["name"]
        for orig_name, orig_cooccur in cat_to_originals[cand_cat]:
            if orig_name == from_seed:
                continue
            if cand_name not in orig_cooccur:
                return True
        return False

    return check


def _token_overlap_candidates(
    seed_tokens: set[str],
    lookup_name: str,
    db_name: str,
    registry,
    cooccur_map: dict,
    score_idx: int,
    seed_tag: dict | None,
    scope: str,
    allowed: set[str] | None,
    ban: set[str],
    seen: set[str],
) -> list[tuple[str, float]]:
    """Find tags sharing tokens with the seed, ranked by co-occurrence."""
    if not seed_tokens:
        return []

    # Gather all tag indices that share at least one token
    candidate_idxs: set[int] = set()
    for token in seed_tokens:
        hits = registry.token_index.get(token)
        if hits:
            candidate_idxs.update(hits)

    results: list[tuple[str, float]] = []
    for ci in candidate_idxs:
        cand = registry.tags[ci]
        cand_db = cand["danbooru_name"]
        cand_name = cand["name"]

        # Skip self
        if cand_name == lookup_name or cand_db == db_name:
            continue
        if cand_db in seen or cand_name in seen:
            continue
        if cand_name in ban or cand_db in ban:
            continue
        if allowed and cand.get("filtering", "sfw") not in allowed:
            continue
        if not _scope_ok(scope, seed_tag, cand):
            continue

        # Co-occurrence score for ranking (0 if no data)
        cooccur = cooccur_map.get(cand_name)
        if cooccur:
            metric_val = cooccur[score_idx - 1]
        else:
            metric_val = 0.0

        # Token overlap ratio as a quality floor — candidates with no
        # co-occurrence data are ranked by overlap alone
        cand_tokens = set(tokenize(cand_name))
        shared = len(seed_tokens & cand_tokens)
        overlap = shared / max(len(seed_tokens), len(cand_tokens))

        # Sort key: co-occurrence is primary when positive, overlap
        # is the fallback for unscored candidates
        sort_key = (metric_val if metric_val > 0 else 0.0, overlap)
        results.append((cand_db, metric_val, sort_key))

    # Sort by (co-occurrence desc, overlap desc)
    results.sort(key=lambda r: (-r[2][0], -r[2][1]))

    return [(db, metric) for db, metric, _key in results]


def _cooccurrence_candidates(
    related: list[tuple[str, float, float, float]],
    score_idx: int,
    registry,
    seed_tag: dict | None,
    scope: str,
    allowed: set[str] | None,
    ban: set[str],
    seen: set[str],
    already_selected: set[str],
) -> list[tuple[str, float]]:
    """Co-occurrence-only candidates — positive scores only."""
    # Re-sort by chosen metric
    sorted_related = sorted(related, key=lambda r: -r[score_idx])

    results: list[tuple[str, float]] = []
    for ct, jac, pmi, lmi in sorted_related:
        metric_val = (jac, pmi, lmi)[score_idx - 1]
        if metric_val <= 0:
            continue

        if ct in already_selected:
            continue

        ni = registry.name_to_idx.get(ct)
        if ni is None:
            continue

        cand = registry.tags[ni]
        cand_db = cand["danbooru_name"]

        if cand_db in seen or cand["name"] in seen:
            continue
        if cand["name"] in ban or cand_db in ban:
            continue
        if allowed and cand.get("filtering", "sfw") not in allowed:
            continue
        if not _scope_ok(scope, seed_tag, cand):
            continue

        results.append((cand_db, metric_val))

    return results


def _scope_ok(scope: str, seed_tag: dict | None, cand: dict) -> bool:
    """Check whether *cand* falls within the requested scope of *seed_tag*.

    When scope is "same category" or "same subgroup", candidates missing
    that field (null) are rejected — they must not fall back to a looser
    group-level match.
    """
    if scope == "any" or seed_tag is None:
        return True
    if scope == "same category":
        seed_cat = seed_tag.get("category")
        cand_cat = cand.get("category")
        if not seed_cat or not cand_cat:
            return False
        return cand_cat == seed_cat
    if scope == "same subgroup":
        seed_sub = seed_tag.get("subgroup")
        cand_sub = cand.get("subgroup")
        if not seed_sub or not cand_sub:
            return False
        return cand_sub == seed_sub
    if scope == "same group":
        return cand["group"] == seed_tag["group"]
    return True
