"""Sort UTF-8 CSV by one signed 64-bit integer column. Python 3.9+."""
import argparse
from contextlib import closing
import csv
import json
import os
from pathlib import Path
import re
import sqlite3
import sys
import tempfile


def records(reader, key_index, width):
    for seq, row in enumerate(reader, 1):
        if len(row) != width:
            raise ValueError(f"record {seq}: expected {width} fields, got {len(row)}")
        text = row[key_index]
        if not re.fullmatch(r"[+-]?[0-9]+", text):
            raise ValueError(f"record {seq}: key must be an integer, got {text!r}")
        number = int(text)
        if not -(2**63) <= number < 2**63:
            raise ValueError(f"record {seq}: key outside signed 64-bit range")
        yield seq, number, json.dumps(row, ensure_ascii=True)


def sort_csv(source, output, key, work_dir):
    if source.resolve() == output.resolve():
        raise ValueError("input and output must differ")
    if os.path.lexists(output):
        raise FileExistsError(f"output already exists: {output}")
    staged = None
    try:
        with tempfile.TemporaryDirectory(prefix="csv-sort-", dir=work_dir) as tmp:
            with closing(sqlite3.connect(str(Path(tmp) / "rows.sqlite"))) as con:
                con.execute("PRAGMA temp_store=FILE")
                con.execute("PRAGMA cache_size=-8192")
                con.execute("CREATE TABLE items (seq INTEGER PRIMARY KEY, "
                            "sort_key INTEGER NOT NULL, payload TEXT NOT NULL)")
                with source.open(encoding="utf-8-sig", newline="") as file:
                    reader = csv.reader(file, strict=True)
                    header = next(reader, None)
                    if not header or any(not name.strip() for name in header):
                        raise ValueError("a nonempty header is required")
                    if len(set(header)) != len(header):
                        raise ValueError("duplicate header names")
                    if key not in header:
                        raise ValueError(f"missing key column: {key!r}")
                    with con:
                        con.executemany(
                            "INSERT INTO items VALUES (?, ?, ?)",
                            records(reader, header.index(key), len(header)),
                        )
                count = con.execute("SELECT count(*) FROM items").fetchone()[0]
                with tempfile.NamedTemporaryFile(
                    mode="w", encoding="utf-8", newline="", delete=False,
                    prefix=".csv-sort-", suffix=".tmp", dir=output.parent,
                ) as file:
                    staged = Path(file.name)
                    writer = csv.writer(file, lineterminator="\n")
                    writer.writerow(header)
                    for (payload,) in con.execute(
                        "SELECT payload FROM items ORDER BY sort_key, seq"
                    ):
                        writer.writerow(json.loads(payload))
        # Same-directory hard link: publish a complete file without overwriting.
        os.link(staged, output)
        return count
    finally:
        if staged is not None:
            staged.unlink(missing_ok=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
    parser.add_argument("input", type=Path)
    parser.add_argument("output", type=Path, help="new file in an existing directory")
    parser.add_argument("--key", required=True, help="exact integer column name")
    parser.add_argument("--work-dir", type=Path, help="existing SQLite workspace parent")
    args = parser.parse_args()
    csv.field_size_limit(1024 * 1024)
    try:
        count = sort_csv(args.input, args.output, args.key, args.work_dir)
    except (OSError, ValueError, csv.Error, sqlite3.Error) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    print(f"OK: rows={count} key={args.key} output={args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())