import os
import json
import argparse


def get_filenames_from_dir(directory):
    """
    Recursively collect all filenames from directory.
    """
    filenames = set()
    for root, _, files in os.walk(directory):
        for f in files:
            filenames.add(f)
    return filenames


def filter_and_split(input_jsonl, directory, output_dir, batch_size=5000):
    os.makedirs(output_dir, exist_ok=True)

    existing_files = get_filenames_from_dir(directory)

    batch = []
    batch_start_idx = 1
    total_kept = 0
    total_removed = 0

    def write_batch(batch, start_idx):
        end_idx = start_idx + len(batch) - 1
        filename = f"batch_input_remaining_{start_idx}_{end_idx}.jsonl"
        filepath = os.path.join(output_dir, filename)

        with open(filepath, "w") as f:
            for item in batch:
                f.write(json.dumps(item) + "\n")

        print(f"Written: {filepath}")

    with open(input_jsonl, "r") as infile:
        for line in infile:
            line = line.strip()
            if not line:
                continue

            try:
                data = json.loads(line)
            except json.JSONDecodeError:
                print("Skipping invalid JSON line")
                continue

            custom_id = data.get("custom_id")
            if custom_id is None:
                continue

            if custom_id.removesuffix(".png") + ".json" not in existing_files:
                batch.append(data)
                total_kept += 1

                if len(batch) == batch_size:
                    write_batch(batch, batch_start_idx)
                    batch_start_idx += batch_size
                    batch = []
            else:
                total_removed += 1

    # Write remaining
    if batch:
        write_batch(batch, batch_start_idx)

    print(f"\nDone.")
    print(f"Kept: {total_kept}")
    print(f"Removed: {total_removed}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Filter JSONL and split into batches")
    parser.add_argument("--input", required=True, help="Input JSONL file")
    parser.add_argument("--dir", required=True, help="Directory with files")
    parser.add_argument("--output_dir", required=True, help="Directory to save output JSONL files")
    parser.add_argument("--batch_size", type=int, default=5000, help="Max lines per output file")

    args = parser.parse_args()

    filter_and_split(args.input, args.dir, args.output_dir, args.batch_size)