#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations

import argparse
import gzip
import io
import json
import os
import random
import sys
from dataclasses import dataclass
from typing import Iterable, List, Optional, Tuple

@dataclass
class ShardConfig:
    input_jsonl: str
    output_dir: str
    samples_per_shard: int = 100_000
    compression: str = "none"  # none | gz
    shuffle_buffer: int = 0  # 0 disables buffered shuffle
    seed: int = 42

def ensure_dir(path: str) -> None:
    os.makedirs(path, exist_ok=True)

def open_input(path: str) -> Iterable[str]:
    if path.endswith(".gz"):
        return (line.decode("utf-8") for line in gzip.open(path, "rb"))
    return open(path, "r", encoding="utf-8")

def open_output(path: str, compression: str):
    if compression == "gz":
        return gzip.open(path + ".gz", "wb")
    return open(path, "wb")

def write_line_binary(fh, line: str) -> None:
    if isinstance(fh, io.BufferedWriter) or isinstance(fh, gzip.GzipFile):
        fh.write(line.encode("utf-8"))
        return
    # Fallback
    fh.write(line)

def buffered_shuffle_stream(lines: Iterable[str], buffer_size: int, rng: random.Random) -> Iterable[str]:
    if buffer_size <= 0:
        for line in lines:
            yield line
        return

    buffer: List[str] = []
    for line in lines:
        buffer.append(line)
        if len(buffer) >= buffer_size:
            # Pop a random element
            idx = rng.randrange(len(buffer))
            buffer[idx], buffer[-1] = buffer[-1], buffer[idx]
            yield buffer.pop()
    # Flush remaining in random order
    while buffer:
        idx = rng.randrange(len(buffer))
        buffer[idx], buffer[-1] = buffer[-1], buffer[idx]
        yield buffer.pop()

def shard_writer(cfg: ShardConfig) -> Tuple[List[str], List[int]]:
    ensure_dir(cfg.output_dir)
    rng = random.Random(cfg.seed)

    shard_paths: List[str] = []
    shard_counts: List[int] = []

    base = os.path.splitext(os.path.basename(cfg.input_jsonl))[0]
    shard_index = 0
    count_in_shard = 0
    current_fh = None

    def _start_new_shard() -> Tuple[object, str]:
        nonlocal shard_index
        shard_name = f"{base}-{shard_index:06d}.jsonl"
        shard_path = os.path.join(cfg.output_dir, shard_name)
        fh = open_output(shard_path, cfg.compression)
        shard_index += 1
        final_path = shard_path + (".gz" if cfg.compression == "gz" else "")
        shard_paths.append(os.path.abspath(final_path))
        shard_counts.append(0)
        return fh, final_path

    try:
        current_fh, _ = _start_new_shard()
        with open_input(cfg.input_jsonl) as fin:  # type: ignore[arg-type]
            # Optionally apply buffered shuffle to reduce ordering bias
            stream = buffered_shuffle_stream(fin, cfg.shuffle_buffer, rng)
            for raw in stream:
                line = raw.rstrip("\n")
                if not line:
                    continue
                # Validate JSON once to avoid writing corrupted records
                try:
                    _ = json.loads(line)
                except Exception as e:
                    print(f"⚠️  Skip bad JSON line: {e}", file=sys.stderr)
                    continue

                if count_in_shard >= cfg.samples_per_shard:
                    current_fh.close()
                    count_in_shard = 0
                    current_fh, _ = _start_new_shard()

                write_line_binary(current_fh, line + "\n")
                count_in_shard += 1
                shard_counts[-1] += 1
    finally:
        if current_fh is not None:
            current_fh.close()

    # Write an index file for convenience
    index_path = os.path.join(cfg.output_dir, f"{base}.index.json")
    with open(index_path, "w", encoding="utf-8") as f:
        json.dump({"shards": shard_paths, "counts": shard_counts}, f, ensure_ascii=False, indent=2)

    return shard_paths, shard_counts

def parse_args() -> ShardConfig:
    p = argparse.ArgumentParser(description="Pack a large JSONL into sharded JSONL files for efficient training")
    p.add_argument("--input_jsonl", required=True, help="Path to the source JSONL (optionally .gz)")
    p.add_argument("--output_dir", required=True, help="Directory to write shard files")
    p.add_argument("--samples_per_shard", type=int, default=100_000)
    p.add_argument("--compression", choices=["none", "gz"], default="none")
    p.add_argument("--shuffle_buffer", type=int, default=0, help="Buffered shuffle size; 0 disables")
    p.add_argument("--seed", type=int, default=42)
    args = p.parse_args()
    return ShardConfig(
        input_jsonl=args.input_jsonl,
        output_dir=args.output_dir,
        samples_per_shard=args.samples_per_shard,
        compression=args.compression,
        shuffle_buffer=args.shuffle_buffer,
        seed=args.seed,
    )

def main() -> None:
    cfg = parse_args()
    if not os.path.isfile(cfg.input_jsonl):
        print(f"❌ Input JSONL not found: {cfg.input_jsonl}", file=sys.stderr)
        sys.exit(1)
    ensure_dir(cfg.output_dir)
    shards, counts = shard_writer(cfg)
    total = sum(counts)
    print("✅ Sharding completed.")
    print(f"   Shards: {len(shards)}")
    print(f"   Samples: {total}")
    print(f"   First shard: {shards[0] if shards else 'N/A'}")
    print(f"   Output dir: {os.path.abspath(cfg.output_dir)}")

if __name__ == "__main__":
    main()