#!/usr/bin/env python3
import argparse
import csv
import re
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("plan", type=Path)
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--minimum-frames", type=int, default=5)
    args = parser.parse_args()

    args.output.mkdir(parents=True, exist_ok=True)
    if any(args.output.iterdir()):
        raise SystemExit(f"Output must be empty: {args.output}")

    with args.plan.open(newline="", encoding="utf-8") as handle:
        rows = list(csv.DictReader(handle))

    kept = []
    review = []
    sequence = 0
    for row in rows:
        frames = int(row["frames"])
        destination_list = kept if frames >= args.minimum_frames else review
        if frames < args.minimum_frames:
            destination_list.append({**row, "output_group": ""})
            continue
        sequence += 1
        first = Path(row["first_file"]).stem
        parent_number = int(re.search(r"Burst(\d+)", row["parent"]).group(1))
        segment_number = int(row["segment"])
        group_name = (
            f"Burst{sequence:02d}_{first}_{frames}f_"
            f"P{parent_number:02d}S{segment_number:02d}"
        )
        group = args.output / group_name
        group.mkdir()
        first_number = int(re.search(r"(\d+)$", first).group(1))
        last_number = int(re.search(r"(\d+)$", Path(row["last_file"]).stem).group(1))
        parent = args.source / row["parent"]
        for number in range(first_number, last_number + 1):
            matches = list(parent.glob(f"*{number}.*"))
            matches = [path for path in matches if path.suffix.lower() in (".jpg", ".jpeg")]
            if len(matches) != 1:
                raise RuntimeError(f"Expected one JPEG for frame {number} in {parent}, got {matches}")
            (group / matches[0].name).symlink_to(matches[0].resolve())
        destination_list.append({**row, "output_group": group_name})

    fields = list(rows[0]) + ["output_group"]
    for name, data in (("kept_segments.csv", kept), ("review_fragments.csv", review)):
        with (args.output / name).open("w", newline="", encoding="utf-8") as handle:
            writer = csv.DictWriter(handle, fieldnames=fields)
            writer.writeheader()
            writer.writerows(data)
    print(f"Kept groups: {len(kept)}")
    print(f"Review fragments: {len(review)}")
    print(f"Output: {args.output}")


if __name__ == "__main__":
    main()