"""
generate_splits_txt.py

Generates a splits.txt file listing all label filenames across train/test/val
splits in the format:

    split/labels/<filename>.txt

One entry per line, ordered train → test → val. This file can be consumed by a
shell script to recreate the directory structure and create symlinks, e.g.:

    while IFS= read -r entry; do
      split=$(echo "$entry" | cut -d/ -f1)
      fname=$(basename "$entry" .txt)
      ln -s "$SOURCE/labels/$fname.txt" "splits/$split/labels/$fname.txt"
      ln -s "$SOURCE/images/$fname.jpg" "splits/$split/images/$fname.jpg"
    done < splits.txt
"""

import argparse
from pathlib import Path


SPLITS = ["train", "test", "val"]


def generate_splits_txt(splits_dir: Path, output_file: Path) -> None:
    entries = []
    counts = {}

    for split in SPLITS:
        labels_dir = splits_dir / split / "labels"
        if not labels_dir.is_dir():
            print(f"Warning: {labels_dir} not found, skipping.")
            counts[split] = 0
            continue

        files = sorted(labels_dir.glob("*.txt"))
        for f in files:
            entries.append(f"{split}/labels/{f.name}")
        counts[split] = len(files)

    output_file.write_text("\n".join(entries) + "\n")

    total = sum(counts.values())
    print(f"Written {total} entries to {output_file}")
    for split in SPLITS:
        print(f"  {split}: {counts.get(split, 0)}")


def main():
    parser = argparse.ArgumentParser(description="Generate splits.txt from a splits directory.")
    parser.add_argument(
        "--input_dir",
        type=Path,
        help="Path to the input folder containing train/val/test splits",
    )
    parser.add_argument(
        "--output_dir",
        type=Path,
        default=Path(__file__).resolve().parents[2] / "data" / "splits",
        help="Path to the output folder (default: <project_root>/data/splits)",
    )
    args = parser.parse_args()

    splits_dir = args.input_dir
    output_dir = args.output_dir

    if not splits_dir.is_dir():
        raise SystemExit(f"Error: input folder not found: {splits_dir}")

    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / f"{splits_dir.name}_split.txt"

    generate_splits_txt(splits_dir, output_file)


if __name__ == "__main__":
    main()