#!/usr/bin/env python3
"""Build a deterministic punctuation-balanced fine-tuning corpus."""

from __future__ import annotations

import argparse
import hashlib
import heapq
import json
import sys
from pathlib import Path


CATEGORIES = {
    "question": lambda text: "?" in text,
    "exclamation": lambda text: "!" in text,
    "quotes": lambda text: "«" in text or "»" in text,
    "ellipsis": lambda text: "…" in text or "..." in text,
    "neutral": lambda text: not any(mark in text for mark in ("?", "!", "«", "»", "…", "...")),
}


def build(args: argparse.Namespace) -> None:
    if args.output.exists() or args.metadata.exists():
        raise FileExistsError("Refusing to overwrite output or metadata")
    excluded = set()
    if args.exclude:
        excluded = {line.strip() for line in args.exclude.read_text(encoding="utf-8-sig").splitlines() if line.strip()}
    heaps: dict[str, list[tuple[int, str]]] = {name: [] for name in CATEGORIES}
    scanned = 0
    for text in args.source.read_text(encoding="utf-8").splitlines():
        text = text.strip()
        if not text or text in excluded:
            continue
        scanned += 1
        for name, predicate in CATEGORIES.items():
            if not predicate(text):
                continue
            score = int.from_bytes(hashlib.blake2b(f"{name}\0{text}".encode("utf-8"), digest_size=8).digest(), "big")
            heap = heaps[name]
            if len(heap) < args.per_category:
                heapq.heappush(heap, (-score, text))
            elif score < -heap[0][0]:
                heapq.heapreplace(heap, (-score, text))
    selected: dict[str, set[str]] = {
        name: {text for _, text in heap} for name, heap in heaps.items()
    }
    combined = sorted(set().union(*selected.values()), key=lambda text: hashlib.sha256(text.encode("utf-8")).digest())
    args.output.write_text("\n".join(combined) + "\n", encoding="utf-8")
    metadata = {
        "source": str(args.source), "exclude": str(args.exclude) if args.exclude else None,
        "scanned": scanned, "per_category_limit": args.per_category,
        "category_counts": {name: len(values) for name, values in selected.items()},
        "unique_sentences": len(combined), "output_bytes": args.output.stat().st_size,
    }
    args.metadata.write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False, indent=2))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--exclude", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--metadata", type=Path, required=True)
    parser.add_argument("--per-category", type=int, default=4000)
    return parser.parse_args()


if __name__ == "__main__":
    sys.exit(build(parse_args()) or 0)