"""Python 3: isolate invalid JSONL lines; export accepted objects as CSV.
Run: python quarantine-jsonl.py input.jsonl new-output-directory
Output directory must not exist. This example retains accepted rows in memory.
"""
import base64
import csv
import json
import sys
from pathlib import Path

def reject_constant(token):
    raise ValueError("non-JSON constant: " + token)

def unique_keys(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError("duplicate key: " + key)
        result[key] = value
    return result

def cell(value):
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False, allow_nan=False,
                      separators=(",", ":"))

source, destination = map(Path, sys.argv[1:3])
destination.mkdir(parents=True, exist_ok=False)
accepted, rejected, fields = [], [], []
line_count = 0
with source.open("rb") as stream:
    for line_number, raw in enumerate(stream, 1):
        line_count = line_number
        try:
            text = raw.decode("utf-8")
            if not text.strip():
                raise ValueError("blank line")
            record = json.loads(text, parse_constant=reject_constant,
                                object_pairs_hook=unique_keys)
            if not isinstance(record, dict):
                raise ValueError("valid JSON, but expected an object")
        except (UnicodeDecodeError, ValueError) as error:
            rejected.append({"source_line": line_number,
                             "reason": str(error),
                             "raw_base64": base64.b64encode(raw).decode("ascii")})
            continue
        accepted.append((line_number, record))
        for field in record:
            if field not in fields:
                fields.append(field)

# Prefix source keys so they cannot collide with the provenance column.
with (destination / "accepted.csv").open("w", encoding="utf-8", newline="") as stream:
    writer = csv.writer(stream)
    writer.writerow(["source_line"] + ["data." + key for key in fields])
    for line_number, record in accepted:
        writer.writerow([line_number] + [cell(record[key]) if key in record else ""
                                        for key in fields])
with (destination / "rejected.jsonl").open("w", encoding="utf-8", newline="\n") as stream:
    for entry in rejected:
        stream.write(json.dumps(entry, ensure_ascii=False) + "\n")
report = {"physical_lines": line_count, "accepted": len(accepted),
          "rejected": len(rejected),
          "balanced": line_count == len(accepted) + len(rejected)}
(destination / "report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
print(json.dumps(report))