import os

input_file = "batchapi/batch_input.jsonl"
output_dir = "batchapi/batches"
batch_size = 5000

os.makedirs(output_dir, exist_ok=True)

with open(input_file, "r", encoding="utf-8") as f:
    batch = []
    start_line = 1

    for i, line in enumerate(f, start=1):
        batch.append(line)

        if len(batch) == batch_size:
            end_line = start_line + batch_size - 1
            output_file = os.path.join(
                output_dir, f"batch_input_{start_line}_{end_line}.jsonl"
            )

            with open(output_file, "w", encoding="utf-8") as out:
                out.writelines(batch)

            batch = []
            start_line = i + 1

    # Handle remaining lines
    if batch:
        end_line = start_line + len(batch) - 1
        output_file = os.path.join(
            output_dir, f"batch_input_{start_line}_{end_line}.jsonl"
        )

        with open(output_file, "w", encoding="utf-8") as out:
            out.writelines(batch)