#!/usr/bin/env python3
"""Sample CSV records in one pass, then write a new output file."""
import argparse
import csv
import json
from pathlib import Path
import random


def sample_csv(source, size, seed):
    if size < 1:
        raise ValueError("sample size must be positive")
    rng = random.Random(seed)
    reservoir = []
    seen = 0
    with source.open(encoding="utf-8-sig", newline="") as handle:
        reader = csv.reader(handle, strict=True)
        header = next(reader, None)
        if not header or any(not name.strip() for name in header):
            raise ValueError("a header with nonempty column names is required")
        if len(set(header)) != len(header):
            raise ValueError("duplicate column names")
        for seen, row in enumerate(reader, start=1):
            if len(row) != len(header):
                raise ValueError(
                    f"record {seen}: expected {len(header)} fields, got {len(row)}"
                )
            item = (seen, row)
            if seen <= size:
                reservoir.append(item)
            else:
                slot = rng.randrange(seen)
                if slot < size:
                    reservoir[slot] = item
    reservoir.sort(key=lambda item: item[0])
    return header, [row for _, row in reservoir], seen


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--size", required=True, type=int)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    if args.size < 1:
        parser.error("--size must be positive")
    try:
        header, rows, total = sample_csv(args.input, args.size, args.seed)
        # Exclusive creation refuses to overwrite any existing destination.
        with args.output.open("x", encoding="utf-8", newline="") as handle:
            writer = csv.writer(handle)
            writer.writerow(header)
            writer.writerows(rows)
    except (OSError, UnicodeError, csv.Error, ValueError) as exc:
        parser.exit(1, f"ERROR: {exc}\n")
    print(json.dumps({
        "input_records": total,
        "requested_records": args.size,
        "sampled_records": len(rows),
        "seed": args.seed,
    }, indent=2))


if __name__ == "__main__":
    main()