#!/usr/bin/env python3
"""
Split JSONL/JSONL.gz files into a desired number of pieces.

Uses a 2-pass approach:
  Pass 1 – measure cumulative output size (compressed or plain) at regular
           line intervals, without writing anything to disk.
           Uses compresslevel=1 for speed (3-5x faster than default).
  Pass 2 – split at the line boundaries that produce the most evenly-sized
           output files, never cutting a line in half.
"""

from __future__ import annotations

import argparse
import gzip
import json
import sys
from bisect import bisect_left
from pathlib import Path


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _open_func(path: Path, mode: str = "rt", **kwargs):
    """Open a file, transparently handling .gz compression."""
    suffixes = path.suffixes
    if suffixes and suffixes[-1] == ".gz":
        return gzip.open(path, mode, **kwargs)
    return open(path, mode, **kwargs)


def _output_ext(input_path: Path) -> str:
    """Determine output file extension from the input path.

    Preserves the full compound extension so that .jsonl.gz -> .jsonl.gz
    and .jsonl -> .jsonl.
    """
    suffixes = input_path.suffixes
    if len(suffixes) >= 2 and suffixes[-2:] == [".jsonl", ".gz"]:
        return ".jsonl.gz"
    return ".jsonl"


class _NullWriter:
    """A writable, flushable sink that counts bytes without storing them.

    Used as the backing store for *gzip.GzipFile* during pass 1 so we can
    learn how many compressed bytes have been produced at each sampled line
    boundary — the key to evenly-sized splitting.

    This answers the question "can we get bytes written during compression?":
    yes, by intercepting every ``write()`` call that GzipFile makes to its
    underlying file object.
    """

    def __init__(self) -> None:
        self.bytes_written: int = 0

    def write(self, data: bytes) -> int:
        self.bytes_written += len(data)
        return len(data)

    def flush(self) -> None:
        pass

    def close(self) -> None:
        pass

    def tell(self) -> int:
        return self.bytes_written


# ---------------------------------------------------------------------------
# Pass 1 — measure
# ---------------------------------------------------------------------------

def pass1_measure(
    input_path: Path,
    compressed: bool,
    sample_interval: int = 10000,
) -> tuple[int, int, list[tuple[int, int]]]:
    """Walk the input once and record cumulative byte sizes at sampled lines.

    When *compressed* is True the data is run through a *gzip.GzipFile* whose
    backing store is a :class:`_NullWriter`, so the byte counts reflect the
    real compressed output size.  Uses compresslevel=1 for speed (3-5x faster
    than default level 9), while maintaining accurate relative size estimates.

    When False, plain UTF-8 byte counts are used.

    Returns *(total_lines, total_bytes, samples)* where *samples* is a list of
    ``(line_number, cumulative_bytes)`` pairs, always ending with the final
    ``(total_lines, total_bytes)``.
    """
    samples: list[tuple[int, int]] = []

    if compressed:
        # ---------------------------------------------------------------
        # Core trick: GzipFile writes compressed bytes to _NullWriter,
        # which counts them.  After gz.flush() the count is exact.
        # compresslevel=1 is much faster than default 9, while still
        # providing accurate relative size estimates for splitting.
        # ---------------------------------------------------------------
        sink = _NullWriter()
        with gzip.GzipFile(fileobj=sink, mode="wb", compresslevel=1) as gz:
            with _open_func(input_path, "rt", encoding="utf-8") as f:
                line_no = 0
                for line_no, line in enumerate(f, 1):
                    gz.write(line.encode("utf-8"))
                    if line_no % sample_interval == 0:
                        gz.flush()                       # push bytes to sink
                        samples.append((line_no, sink.bytes_written))
                gz.flush()
                samples.append((line_no, sink.bytes_written))
        total_bytes = sink.bytes_written
    else:
        cumulative = 0
        with _open_func(input_path, "rt", encoding="utf-8") as f:
            line_no = 0
            for line_no, line in enumerate(f, 1):
                cumulative += len(line.encode("utf-8"))
                if line_no % sample_interval == 0:
                    samples.append((line_no, cumulative))
            samples.append((line_no, cumulative))
        total_bytes = cumulative

    return line_no, total_bytes, samples


# ---------------------------------------------------------------------------
# Compute split points
# ---------------------------------------------------------------------------

def compute_piece_lines(
    total_lines: int,
    total_bytes: int,
    samples: list[tuple[int, int]],
    num_pieces: int,
) -> list[int]:
    """Determine how many lines each piece should contain.

    Uses the sampled ``(line, bytes)`` data to find, for each split point,
    the line number whose cumulative size is closest to the target
    ``total_bytes / num_pieces``.  Linear interpolation is used between
    sample points for finer granularity.

    Returns a list of ``num_pieces`` line counts whose sum equals
    *total_lines*.
    """
    if num_pieces <= 1:
        return [total_lines]

    target_per_piece = total_bytes / num_pieces
    sample_bytes = [s[1] for s in samples]

    # Cumulative line count at each split boundary
    split_cum_lines: list[int] = []

    for piece in range(1, num_pieces):
        target = target_per_piece * piece

        # Binary search: first sample whose cumulative bytes >= target
        idx = bisect_left(sample_bytes, target)

        if idx == 0:
            # Target falls before the first sample — interpolate from (0, 0)
            first_line, first_bytes = samples[0]
            if first_bytes > 0:
                est_line = int(first_line * target / first_bytes)
            else:
                est_line = 1
        elif idx >= len(samples):
            # Target beyond last sample (shouldn't happen for piece < num_pieces)
            est_line = total_lines
        else:
            # Interpolate between samples[idx-1] and samples[idx]
            prev_line, prev_bytes = samples[idx - 1]
            cur_line, cur_bytes = samples[idx]
            if cur_bytes == prev_bytes:
                est_line = cur_line
            else:
                ratio = (target - prev_bytes) / (cur_bytes - prev_bytes)
                est_line = int(prev_line + ratio * (cur_line - prev_line))

        # Clamp: strictly increasing, ≥1 line per piece, leave room for rest
        min_line = (split_cum_lines[-1] + 1) if split_cum_lines else 1
        max_line = total_lines - (num_pieces - piece)
        est_line = max(min_line, min(est_line, max_line))

        split_cum_lines.append(est_line)

    # Convert cumulative boundaries → per-piece line counts
    piece_lines: list[int] = []
    prev = 0
    for cl in split_cum_lines:
        piece_lines.append(cl - prev)
        prev = cl
    piece_lines.append(total_lines - prev)

    return piece_lines


# ---------------------------------------------------------------------------
# Pass 2 — split
# ---------------------------------------------------------------------------

def pass2_split(
    input_path: Path,
    piece_lines: list[int],
    output_dir: Path,
    output_prefix: str,
    ext: str,
) -> None:
    """Write each piece with the calculated number of lines."""
    with _open_func(input_path, "rt", encoding="utf-8") as infile:
        for idx, n_lines in enumerate(piece_lines):
            output_name = f"{output_prefix}_{idx:02d}{ext}"
            output_path = output_dir / output_name

            with _open_func(output_path, "wt", encoding="utf-8") as outfile:
                for _ in range(n_lines):
                    line = infile.readline()
                    if not line:
                        break
                    outfile.write(line)  # type: ignore

            print(f"  Written {output_path} ({n_lines} lines)")


# ---------------------------------------------------------------------------
# JSON validation
# ---------------------------------------------------------------------------

def validate_jsonl(input_path: Path) -> bool:
    """Validate every line in the file is valid JSON (or blank)."""
    with _open_func(input_path, "rt", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            stripped = line.strip()
            if not stripped:
                continue
            try:
                json.loads(stripped)
            except json.JSONDecodeError as e:
                print(f"Error on line {line_num}: {e}", file=sys.stderr)
                return False
    return True


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------

def split_jsonl(
    input_path: Path,
    num_pieces: int,
    output_dir: Path,
    output_prefix: str = "part",
    validate: bool = False,
    sample_interval: int = 10000,
) -> None:
    """Split a JSONL file into evenly-sized pieces (2-pass, size-based).

    Args:
        input_path:      Path to the input .jsonl or .jsonl.gz file.
        num_pieces:      Number of pieces to produce.
        output_dir:      Directory where output files will be written.
        output_prefix:   Prefix for output filenames.
        validate:        If True, validate each line as JSON before splitting.
        sample_interval: How often (in lines) to sample size during pass 1.
    """
    if validate:
        print("Validating JSON lines...")
        if not validate_jsonl(input_path):
            sys.exit(1)
        print("Validation passed.")

    ext = _output_ext(input_path)
    compressed = ext == ".jsonl.gz"
    size_label = "compressed" if compressed else "uncompressed"

    # ---- Pass 1: measure --------------------------------------------------
    print(
        f"Pass 1: measuring {size_label} size "
        f"(sampling every {sample_interval} lines, compresslevel=1)..."
    )
    total_lines, total_bytes, samples = pass1_measure(
        input_path, compressed, sample_interval
    )
    print(f"  Total lines : {total_lines}")
    print(
        f"  Total {size_label} size : "
        f"{total_bytes:,} bytes ({total_bytes / 1048576:.2f} MiB)"
    )

    if total_lines == 0:
        print("Input file is empty. Nothing to split.", file=sys.stderr)
        sys.exit(1)

    if num_pieces > total_lines:
        print(
            f"Warning: {num_pieces} pieces requested but only {total_lines} "
            f"lines available. Creating {total_lines} pieces instead.",
            file=sys.stderr,
        )
        num_pieces = total_lines

    # ---- Compute split points ---------------------------------------------
    piece_lines = compute_piece_lines(
        total_lines, total_bytes, samples, num_pieces
    )
    target = total_bytes / num_pieces
    print(
        f"Pass 2: splitting into {num_pieces} pieces "
        f"(target ~{target / 1048576:.2f} MiB {size_label} each)..."
    )
    print(f"  Line distribution: {piece_lines}")

    # ---- Pass 2: split ----------------------------------------------------
    output_dir.mkdir(parents=True, exist_ok=True)
    pass2_split(input_path, piece_lines, output_dir, output_prefix, ext)

    print("Done.")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Split JSONL/JSONL.gz files into a desired number of pieces "
            "(2-pass, size-based)."
        ),
    )
    parser.add_argument(
        "input",
        type=Path,
        help="Input JSONL (.jsonl) or gzipped JSONL (.jsonl.gz) file",
    )
    parser.add_argument(
        "pieces",
        type=int,
        help="Number of pieces to split into",
    )
    parser.add_argument(
        "-o",
        "--output-dir",
        type=Path,
        default=None,
        help="Output directory (default: same directory as input file)",
    )
    parser.add_argument(
        "-p",
        "--prefix",
        type=str,
        default="part",
        help="Prefix for output filenames (default: 'part')",
    )
    parser.add_argument(
        "--validate",
        action="store_true",
        help="Validate every line is valid JSON before splitting",
    )
    parser.add_argument(
        "--sample-interval",
        type=int,
        default=10000,
        help=(
            "Line sampling interval for size measurement in pass 1 "
            "(default: 10000). Smaller values improve accuracy at the cost "
            "of slower pass 1."
        ),
    )

    args = parser.parse_args()

    if not args.input.exists():
        print(f"Error: Input file '{args.input}' not found.", file=sys.stderr)
        sys.exit(1)

    if not args.input.is_file():
        print(f"Error: '{args.input}' is not a file.", file=sys.stderr)
        sys.exit(1)

    if args.pieces < 1:
        print("Error: Number of pieces must be at least 1.", file=sys.stderr)
        sys.exit(1)

    if args.pieces > 100:
        print("Error: Number of pieces must not exceed 100.", file=sys.stderr)
        sys.exit(1)

    if args.sample_interval < 1:
        print("Error: Sample interval must be at least 1.", file=sys.stderr)
        sys.exit(1)

    if args.output_dir is None:
        args.output_dir = args.input.parent

    split_jsonl(
        input_path=args.input,
        num_pieces=args.pieces,
        output_dir=args.output_dir,
        output_prefix=args.prefix,
        validate=args.validate,
        sample_interval=args.sample_interval,
    )


if __name__ == "__main__":
    main()