"""Merge numbered CSV shards in numeric filename order. Python 3.9+."""
import argparse
import csv
import os
from pathlib import Path
import re
import sys
import tempfile

NAME = re.compile(r"part-([0-9]+)\.csv")


def ordered_parts(directory):
    parts = []
    seen = set()
    for path in directory.glob("part-*.csv"):
        match = NAME.fullmatch(path.name)
        if match is None or path.is_symlink() or not path.is_file():
            raise ValueError(f"invalid shard: {path.name}")
        number = int(match[1])
        if number < 1 or number in seen:
            raise ValueError(f"nonpositive or duplicate shard number: {path.name}")
        seen.add(number)
        parts.append((number, path))
    if not parts:
        raise ValueError("no part-<positive integer>.csv files")
    return sorted(parts, key=lambda item: item[0])


def merge_csv(directory, output):
    directory = Path(directory).expanduser().resolve(strict=True)
    if not directory.is_dir():
        raise ValueError("input must be a directory")
    output = Path(output).expanduser().absolute()
    output = output.parent.resolve(strict=True) / output.name
    if output.parent == directory or output.is_symlink():
        raise ValueError("output must be outside input directory and not a symlink")
    parts = ordered_parts(directory)
    temporary = None
    count = 0
    expected = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w", encoding="utf-8", newline="", dir=output.parent,
            prefix=f".{output.name}.", suffix=".tmp", delete=False,
        ) as dest:
            temporary = Path(dest.name)
            writer = csv.writer(dest, lineterminator="\n")
            for number, path in parts:
                with path.open(encoding="utf-8-sig", newline="") as source:
                    reader = csv.reader(source, strict=True)
                    try:
                        header = next(reader, None)
                        if not header or any(not field.strip() for field in header):
                            raise ValueError(f"{path.name}: missing or empty header")
                        if len(set(header)) != len(header):
                            raise ValueError(f"{path.name}: duplicate column names")
                        if expected is None:
                            expected = header
                            writer.writerow(header)
                        elif header != expected:
                            raise ValueError(f"{path.name}: header mismatch")
                        for row in reader:
                            if len(row) != len(expected):
                                raise ValueError(
                                    f"{path.name}: physical line {reader.line_num}: "
                                    f"expected {len(expected)} fields, got {len(row)}"
                                )
                            writer.writerow(row)
                            count += 1
                    except csv.Error as exc:
                        raise ValueError(
                            f"{path.name}: physical line {reader.line_num}: {exc}"
                        ) from exc
        os.replace(temporary, output)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)
    return len(parts), count, [number for number, path in parts]


def main():
    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
    parser.add_argument("directory", type=Path, help="directory of completed shards")
    parser.add_argument("output", type=Path, help="destination outside that directory")
    args = parser.parse_args()
    try:
        files, records, order = merge_csv(args.directory, args.output)
    except (OSError, ValueError, RuntimeError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    print(f"OK: files={files} records={records} order={','.join(map(str, order))}")
    return 0


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