#!/usr/bin/env python3
"""Quality filtering for polished folksy sayings.

Reads corpus_polished.jsonl, applies quality filters, outputs filtered corpus
and discard analysis.

Usage:
  python scripts/filter_corpus.py
  python scripts/filter_corpus.py --input corpus/corpus_polished.jsonl --output corpus/corpus_filtered.jsonl
"""

import argparse
import csv
import json
import sys
from difflib import SequenceMatcher
from pathlib import Path

SCRIPT_DIR = Path(__file__).parent
PROJECT_DIR = SCRIPT_DIR.parent
CORPUS_DIR = PROJECT_DIR / "corpus"


def quality_filter(entry):
    """Apply quality filters to a polished entry.

    Returns (passed, reason) tuple.
    """
    text = entry.get("polished_text", "")
    if not text:
        return False, "no_polished_text"

    words = text.split()

    # Length check
    if len(words) > 25:
        return False, "too_long"
    if len(words) < 5:
        return False, "too_short"

    # Must contain at least 2 of the original slot-fill nouns
    slot_words = set(entry.get("slots", {}).values())
    words_present = sum(1 for w in slot_words if w.lower() in text.lower())
    if words_present < 2:
        return False, "lost_key_nouns"

    # No raw ConceptNet artifacts (multi-word underscore phrases)
    if "_" in text:
        return False, "conceptnet_artifact"

    # No broken templates (unfilled slots)
    if "{" in text or "}" in text:
        return False, "unfilled_slot"

    return True, "pass"


def is_near_duplicate(text_a, text_b, threshold=0.75):
    """Check if two texts are near-duplicates."""
    return SequenceMatcher(None, text_a.lower(), text_b.lower()).ratio() > threshold


def deduplicate_within_family(entries):
    """Remove near-duplicates within each meta-template family.

    Returns (kept, removed) lists.
    """
    by_family = {}
    for entry in entries:
        family = entry.get("meta_template", "unknown")
        by_family.setdefault(family, []).append(entry)

    kept = []
    removed = []

    for family, family_entries in by_family.items():
        family_kept = []
        for entry in family_entries:
            text = entry.get("polished_text", "")
            is_dup = False
            for existing in family_kept:
                if is_near_duplicate(text, existing.get("polished_text", "")):
                    is_dup = True
                    break
            if is_dup:
                removed.append((entry, "near_duplicate"))
            else:
                family_kept.append(entry)
        kept.extend(family_kept)

    return kept, removed


def main():
    parser = argparse.ArgumentParser(description="Quality filtering for polished folksy sayings.")
    parser.add_argument("--input", default=str(CORPUS_DIR / "corpus_polished.jsonl"),
                        help="Input polished JSONL file")
    parser.add_argument("--output", default=str(CORPUS_DIR / "corpus_filtered.jsonl"),
                        help="Output filtered JSONL file")
    parser.add_argument("--discard-analysis", default=str(CORPUS_DIR / "discard_analysis.csv"),
                        help="Discard analysis CSV file")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output)
    discard_path = Path(args.discard_analysis)

    if not input_path.exists():
        print(f"Error: {input_path} not found.", file=sys.stderr)
        sys.exit(1)

    # Load polished entries (only those with status=polished)
    all_entries = []
    already_discarded = 0
    with open(input_path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            entry = json.loads(line)
            if entry.get("status") == "polished":
                all_entries.append(entry)
            elif entry.get("status") == "discarded":
                already_discarded += 1

    print(f"Loaded {len(all_entries)} polished entries ({already_discarded} already discarded by LLM)")

    # Apply quality filters
    passed = []
    discards = []  # (entry, reason)

    for entry in all_entries:
        ok, reason = quality_filter(entry)
        if ok:
            passed.append(entry)
        else:
            discards.append((entry, reason))

    print(f"Quality filter: {len(passed)} passed, {len(discards)} discarded")

    # Show discard breakdown
    from collections import Counter
    reason_counts = Counter(r for _, r in discards)
    for reason, count in reason_counts.most_common():
        print(f"  {reason}: {count}")

    # Near-duplicate detection within template families
    kept, dup_removed = deduplicate_within_family(passed)
    discards.extend(dup_removed)

    print(f"Near-duplicate removal: {len(dup_removed)} removed, {len(kept)} remaining")

    # Write filtered output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        for entry in kept:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")

    print(f"\nFiltered corpus: {len(kept)} entries -> {output_path}")

    # Write discard analysis
    with open(discard_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["raw_text", "meta_template", "discard_stage", "discard_reason"])
        for entry, reason in discards:
            writer.writerow([
                entry.get("raw_text", ""),
                entry.get("meta_template", ""),
                "llm_polish" if reason == "no_polished_text" else "quality_filter",
                reason,
            ])

    print(f"Discard analysis: {len(discards)} entries -> {discard_path}")


if __name__ == "__main__":
    main()
