"""
Deduplication pipeline for the code corpus.

Three stages of deduplication, applied in order:

  1. EXACT DEDUP (document-level)
     - SHA-256 hash of normalized content (whitespace-collapsed, lowercased)
     - Drops byte-identical and whitespace-only-differing documents.

  2. NEAR-DEDUP (document-level, MinHash + LSH)
     - Shingles each document into character n-grams.
     - Computes MinHash signatures.
     - Uses LSH banding to cluster near-duplicates (Jaccard >= threshold).
     - Keeps one representative per cluster (the longest, highest-quality).

  3. LINE-LEVEL DEDUP (cross-document)
     - Builds a global set of exact line hashes.
     - Removes lines that appear in >threshold documents (boilerplate:
       license headers, auto-generated stubs, repeated imports, etc.).
     - Drops documents that become too short / too sparse after line removal.

Input:  data/raw_large/documents.jsonl
Output: data/dedup/documents_dedup.jsonl

Reports statistics at each stage so we can verify the pipeline is working.
"""

import hashlib
import json
import os
import re
import sys
from collections import Counter, defaultdict

from datasketch import MinHash, MinHashLSH

PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA_DIR = os.path.join(PROJECT_DIR, "data")
RAW_DOCS_PATH = os.path.join(DATA_DIR, "raw_large", "documents.jsonl")
DEDUP_DIR = os.path.join(DATA_DIR, "dedup")
DEDUP_DOCS_PATH = os.path.join(DEDUP_DIR, "documents_dedup.jsonl")
STATS_PATH = os.path.join(DEDUP_DIR, "dedup_stats.json")

# ─── MinHash / LSH parameters ────────────────────────────────────────────────
SHINGLE_SIZE = 5          # character 5-grams
NUM_PERM = 128            # MinHash permutations (higher = more accurate)
LSH_THRESHOLD = 0.5       # Jaccard similarity threshold for near-dup
LSH_BANDS = 32            # LSH bands (NUM_PERM / LSH_ROWS should be integer)
LSH_ROWS = NUM_PERM // LSH_BANDS  # 128 / 32 = 4

# ─── Line-level dedup parameters ────────────────────────────────────────────
LINE_BOILERPLATE_THRESHOLD = 50  # lines appearing in >= this many docs are boilerplate
MIN_LINES_AFTER_DEDUP = 5        # drop docs with fewer lines after line dedup
MIN_CHARS_AFTER_DEDUP = 200      # drop docs shorter than this after line dedup


def normalize_content(text: str) -> str:
    """Normalize content for exact-hash comparison: collapse whitespace, lowercase."""
    # Collapse all whitespace sequences to single space
    text = re.sub(r"\s+", " ", text)
    return text.strip().lower()


def exact_hash(text: str) -> str:
    """SHA-256 of normalized content."""
    return hashlib.sha256(normalize_content(text).encode("utf-8")).hexdigest()


def shingles(text: str, k: int = SHINGLE_SIZE) -> set[str]:
    """Character k-gram shingles of normalized text."""
    text = re.sub(r"\s+", " ", text).strip().lower()
    if len(text) < k:
        return {text}
    return {text[i : i + k] for i in range(len(text) - k + 1)}


def minhash_signature(text: str, num_perm: int = NUM_PERM) -> MinHash:
    """Compute MinHash signature for a document."""
    m = MinHash(num_perm=num_perm)
    for s in shingles(text):
        m.update(s.encode("utf-8"))
    return m


# ─── Stage 1: Exact dedup ────────────────────────────────────────────────────

def stage1_exact_dedup(docs: list[dict]) -> tuple[list[dict], dict]:
    """Remove exact (and whitespace-only) duplicates. Keep first occurrence."""
    print("\n" + "=" * 60)
    print("STAGE 1: Exact document-level dedup (SHA-256 of normalized content)")
    print("=" * 60)

    seen_hashes = {}
    kept = []
    dropped = 0

    for i, doc in enumerate(docs):
        h = exact_hash(doc["content"])
        if h in seen_hashes:
            dropped += 1
            continue
        seen_hashes[h] = i
        kept.append(doc)

        if (i + 1) % 5000 == 0:
            print(f"  Processed {i+1}/{len(docs)} | kept {len(kept)} | dropped {dropped}")

    stats = {
        "input": len(docs),
        "kept": len(kept),
        "dropped_exact": dropped,
        "drop_rate": dropped / max(len(docs), 1),
    }
    print(f"  Input: {len(docs)} | Kept: {len(kept)} | Dropped: {dropped} "
          f"({stats['drop_rate']:.1%})")
    return kept, stats


# ─── Stage 2: Near-dedup (MinHash LSH) ───────────────────────────────────────

def stage2_near_dedup(docs: list[dict]) -> tuple[list[dict], dict]:
    """Cluster near-duplicates via MinHash LSH, keep best representative per cluster."""
    print("\n" + "=" * 60)
    print(f"STAGE 2: Near-dedup (MinHash LSH, threshold={LSH_THRESHOLD}, "
          f"perm={NUM_PERM}, bands={LSH_BANDS})")
    print("=" * 60)

    # Build LSH index
    lsh = MinHashLSH(
        threshold=LSH_THRESHOLD,
        num_perm=NUM_PERM,
        params=(LSH_BANDS, LSH_ROWS),
    )

    signatures = {}
    kept_indices = set()

    for i, doc in enumerate(docs):
        content = doc["content"]
        if len(content) < 50:
            # Too short for meaningful MinHash — keep as-is
            kept_indices.add(i)
            continue

        mh = minhash_signature(content)
        signatures[i] = mh

        # Query for existing near-duplicates
        candidates = lsh.query(mh)

        if candidates:
            # This doc is a near-duplicate of an existing one.
            # Compare lengths and keep the longer one.
            existing_best = max(candidates, key=lambda j: len(docs[int(j)]["content"]))
            if len(content) > len(docs[int(existing_best)]["content"]):
                # New doc is longer — replace the existing representative
                kept_indices.discard(int(existing_best))
                kept_indices.add(i)
                # We can't easily remove from LSH, so just add the new one.
                # The old one is dropped via kept_indices tracking.
                try:
                    lsh.insert(str(i), mh)
                except Exception:
                    pass
            # else: keep existing, drop this one
        else:
            # No near-duplicate found — keep this doc
            kept_indices.add(i)
            try:
                lsh.insert(str(i), mh)
            except Exception:
                pass

        if (i + 1) % 5000 == 0:
            print(f"  Processed {i+1}/{len(docs)} | clusters so far: {len(kept_indices)}")

    kept = [docs[i] for i in sorted(kept_indices)]
    dropped = len(docs) - len(kept)

    stats = {
        "input": len(docs),
        "kept": len(kept),
        "dropped_near": dropped,
        "drop_rate": dropped / max(len(docs), 1),
        "threshold": LSH_THRESHOLD,
        "num_perm": NUM_PERM,
    }
    print(f"  Input: {len(docs)} | Kept: {len(kept)} | Dropped: {dropped} "
          f"({stats['drop_rate']:.1%})")
    return kept, stats


# ─── Stage 3: Line-level dedup ───────────────────────────────────────────────

def stage3_line_dedup(docs: list[dict]) -> tuple[list[dict], dict]:
    """Remove boilerplate lines that appear across many documents."""
    print("\n" + "=" * 60)
    print(f"STAGE 3: Line-level dedup (boilerplate threshold={LINE_BOILERPLATE_THRESHOLD})")
    print("=" * 60)

    # First pass: count how many documents each normalized line appears in
    line_doc_count = Counter()
    for doc in docs:
        lines = doc["content"].split("\n")
        # Use normalized line hash to save memory
        seen_in_doc = set()
        for line in lines:
            norm = line.strip()
            if not norm:
                continue
            # Normalize: collapse internal whitespace, lowercase
            norm = re.sub(r"\s+", " ", norm).lower()
            if len(norm) < 5:  # skip very short lines (braces, etc.)
                continue
            h = hashlib.md5(norm.encode("utf-8")).hexdigest()
            seen_in_doc.add(h)
        for h in seen_in_doc:
            line_doc_count[h] += 1

    # Identify boilerplate lines (appear in >= threshold documents)
    boilerplate_hashes = {h for h, c in line_doc_count.items() if c >= LINE_BOILERPLATE_THRESHOLD}
    print(f"  Boilerplate lines identified: {len(boilerplate_hashes)}")

    # Second pass: remove boilerplate lines from each document
    kept = []
    dropped_too_short = 0
    lines_removed_total = 0

    for doc in docs:
        lines = doc["content"].split("\n")
        new_lines = []
        for line in lines:
            norm = line.strip()
            if norm:
                norm_h = re.sub(r"\s+", " ", norm).lower()
                if len(norm_h) >= 5:
                    h = hashlib.md5(norm_h.encode("utf-8")).hexdigest()
                    if h in boilerplate_hashes:
                        lines_removed_total += 1
                        continue
            new_lines.append(line)

        new_content = "\n".join(new_lines).strip()
        new_line_count = len([l for l in new_lines if l.strip()])

        if new_line_count < MIN_LINES_AFTER_DEDUP or len(new_content) < MIN_CHARS_AFTER_DEDUP:
            dropped_too_short += 1
            continue

        doc = dict(doc)
        doc["content"] = new_content
        kept.append(doc)

    stats = {
        "input": len(docs),
        "kept": len(kept),
        "dropped_too_short": dropped_too_short,
        "lines_removed": lines_removed_total,
        "boilerplate_hashes": len(boilerplate_hashes),
    }
    print(f"  Input: {len(docs)} | Kept: {len(kept)} | "
          f"Dropped (too short): {dropped_too_short} | "
          f"Lines removed: {lines_removed_total}")
    return kept, stats


# ─── Main ────────────────────────────────────────────────────────────────────

def main():
    os.makedirs(DEDUP_DIR, exist_ok=True)

    print(f"Loading documents from {RAW_DOCS_PATH}...")
    docs = []
    with open(RAW_DOCS_PATH, "r", encoding="utf-8") as f:
        for line in f:
            docs.append(json.loads(line))
    print(f"  Loaded {len(docs):,} documents")
    print(f"  Total size: {sum(len(d['content']) for d in docs) / 1e6:.1f} MB")

    all_stats = {"input_count": len(docs)}

    # Stage 1: Exact dedup
    docs, s1 = stage1_exact_dedup(docs)
    all_stats["stage1_exact"] = s1

    # Stage 2: Near-dedup
    docs, s2 = stage2_near_dedup(docs)
    all_stats["stage2_near"] = s2

    # Stage 3: Line-level dedup
    docs, s3 = stage3_line_dedup(docs)
    all_stats["stage3_line"] = s3

    # Final stats
    final_size = sum(len(d["content"]) for d in docs)
    all_stats["final_count"] = len(docs)
    all_stats["final_size_mb"] = final_size / 1e6

    # Language distribution
    lang_counts = {}
    for d in docs:
        lang_counts[d["language"]] = lang_counts.get(d["language"], 0) + 1
    all_stats["language_distribution"] = lang_counts

    print("\n" + "=" * 60)
    print("DEDUP COMPLETE")
    print("=" * 60)
    print(f"  Input:   {all_stats['input_count']:,} documents")
    print(f"  Final:   {len(docs):,} documents")
    print(f"  Size:    {final_size / 1e6:.1f} MB")
    print(f"  Languages: {lang_counts}")

    # Write deduped documents
    with open(DEDUP_DOCS_PATH, "w", encoding="utf-8") as f:
        for doc in docs:
            f.write(json.dumps(doc, ensure_ascii=False) + "\n")
    print(f"\nDeduped documents written to {DEDUP_DOCS_PATH}")

    # Write stats
    with open(STATS_PATH, "w") as f:
        json.dump(all_stats, f, indent=2)
    print(f"Stats written to {STATS_PATH}")


if __name__ == "__main__":
    main()