"""De-duplication (spec §3.3.5): exact dedup on normalized function text, then
near-duplicate detection via 5-gram Jaccard >= 0.9 using MinHash/LSH.

Operates at the pair level (a pair is one dedup unit — dropping only one half
of a pair would violate pair integrity), keyed by each pair's func_vuln text.
"""
from __future__ import annotations

from collections import defaultdict

from datasketch import MinHash, MinHashLSH

from src.data.normalize_text import shingles, strip_comments_and_whitespace

MINHASH_PERM = 128
JACCARD_THRESHOLD = 0.9


def exact_dedup_pairs(rows: list[dict]) -> tuple[list[dict], dict]:
    """Drops whole pairs whose func_vuln normalizes to a text seen before
    (first pair with that text wins). Returns (kept_rows, stats)."""
    by_pair: dict[str, list[dict]] = defaultdict(list)
    for r in rows:
        by_pair[r["pair_id"]].append(r)

    seen_norm_text: set[str] = set()
    kept_pair_ids: set[str] = set()
    n_dropped_pairs = 0

    for pair_id, pair_rows in by_pair.items():
        vuln_row = next((r for r in pair_rows if r["target"] == 1), pair_rows[0])
        norm = strip_comments_and_whitespace(vuln_row["func_vuln"])
        if norm in seen_norm_text:
            n_dropped_pairs += 1
            continue
        seen_norm_text.add(norm)
        kept_pair_ids.add(pair_id)

    kept_rows = [r for r in rows if r["pair_id"] in kept_pair_ids]
    stats = {
        "n_pairs_before": len(by_pair),
        "n_pairs_dropped_exact": n_dropped_pairs,
        "n_pairs_after": len(by_pair) - n_dropped_pairs,
    }
    return kept_rows, stats


def near_dup_pair_groups(rows: list[dict], threshold: float = JACCARD_THRESHOLD) -> list[set[str]]:
    """Returns groups (as sets of pair_id) of pairs whose func_vuln text is a
    near-duplicate of another's (Jaccard over 5-gram shingles >= threshold),
    via MinHash/LSH. Does not drop anything itself — spec §3.3.5 only asks to
    "log how many samples were removed" and spec §3.3's near-dup check is
    primarily used as a cross-partition leakage check in the split builder
    (`tests/test_splits.py`: "no near-dup across partitions (spot-check)"),
    not necessarily a blanket corpus-wide drop.
    """
    by_pair: dict[str, list[dict]] = defaultdict(list)
    for r in rows:
        by_pair[r["pair_id"]].append(r)

    lsh = MinHashLSH(threshold=threshold, num_perm=MINHASH_PERM)
    minhashes: dict[str, MinHash] = {}

    for pair_id, pair_rows in by_pair.items():
        vuln_row = next((r for r in pair_rows if r["target"] == 1), pair_rows[0])
        norm = strip_comments_and_whitespace(vuln_row["func_vuln"])
        sh = shingles(norm, n=5)
        mh = MinHash(num_perm=MINHASH_PERM)
        for s in sh:
            mh.update(s.encode("utf-8"))
        minhashes[pair_id] = mh
        lsh.insert(pair_id, mh)

    visited: set[str] = set()
    groups: list[set[str]] = []
    for pair_id, mh in minhashes.items():
        if pair_id in visited:
            continue
        matches = set(lsh.query(mh))
        if len(matches) > 1:
            groups.append(matches)
            visited.update(matches)
    return groups