#!/usr/bin/env python3
import argparse
import csv
import json
import math
import re
import subprocess
from collections import defaultdict
from datetime import datetime
from pathlib import Path


CADENCES = (6, 8, 10)


def number(name):
    return int(re.search(r"(\d+)(?=\.[^.]+$)", name).group(1))


def seconds(item, raw_root):
    if not item.get("DateTimeOriginal"):
        raw = raw_root / f"{Path(item['FileName']).stem}.CR3"
        recovered = json.loads(subprocess.check_output([
            "exiftool", "-j", "-DateTimeOriginal", "-SubSecTimeOriginal", str(raw)
        ]))[0]
        item = {**item, **recovered}
    base = datetime.strptime(item["DateTimeOriginal"], "%Y:%m:%d %H:%M:%S").timestamp()
    sub = str(item.get("SubSecTimeOriginal", "0"))
    return base + int(sub) / (10 ** len(sub))


def cadence(segment):
    gaps = [segment[i + 1][1] - segment[i][1] for i in range(len(segment) - 1)]
    useful = sorted(gap for gap in gaps if 0.04 <= gap <= 0.25)
    if not useful:
        return 10
    measured = 1 / useful[len(useful) // 2]
    return min(CADENCES, key=lambda fps: abs(fps - measured))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("groups", type=Path)
    parser.add_argument("raw_root", type=Path)
    parser.add_argument("--gap", type=float, default=0.25)
    parser.add_argument("--minimum-output", type=float, default=3.0)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    items = json.loads(subprocess.check_output([
        "exiftool", "-j", "-r", "-ext", "jpg", "-ext", "jpeg",
        "-DateTimeOriginal", "-SubSecTimeOriginal", "-FileName", "-Directory",
        str(args.groups),
    ]))
    by_dir = defaultdict(list)
    for item in items:
        path = Path(item["Directory"])
        if path.name.startswith("Burst"):
            by_dir[path].append(item)

    rows = []
    for directory in sorted(by_dir, key=lambda path: int(path.name[5:7])):
        ordered = sorted(by_dir[directory], key=lambda item: number(item["FileName"]))
        timed = [(item, seconds(item, args.raw_root)) for item in ordered]
        segments = [[]]
        boundary_gap = None
        boundaries = []
        for pair in timed:
            if segments[-1] and pair[1] - segments[-1][-1][1] > args.gap:
                boundary_gap = pair[1] - segments[-1][-1][1]
                boundaries.append(boundary_gap)
                segments.append([])
            segments[-1].append(pair)
        for index, segment in enumerate(segments, 1):
            fps = cadence(segment)
            frames = len(segment)
            cycle_frames = 1 if frames == 1 else 2 * frames - 2
            cycle_seconds = cycle_frames / fps
            cycles = max(1, math.ceil(args.minimum_output / cycle_seconds))
            rows.append({
                "parent": directory.name,
                "segment": index,
                "first_file": segment[0][0]["FileName"],
                "last_file": segment[-1][0]["FileName"],
                "frames": frames,
                "capture_span_seconds": round(segment[-1][1] - segment[0][1], 3),
                "fps": fps,
                "cycle_seconds": round(cycle_seconds, 3),
                "cycles_for_minimum": cycles,
                "output_seconds": round(cycles * cycle_seconds, 3),
                "split_boundaries_in_parent": len(boundaries),
            })

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=rows[0])
        writer.writeheader()
        writer.writerows(rows)

    print(f"Parents: {len(by_dir)}")
    print(f"True burst segments: {len(rows)}")
    print(f"Parents split: {sum(1 for path in by_dir if sum(row['parent'] == path.name for row in rows) > 1)}")
    sizes = [row["frames"] for row in rows]
    print(f"Frame range: {min(sizes)}-{max(sizes)}")
    print(f"Segments under 5 frames: {sum(size < 5 for size in sizes)}")
    print(f"Segments 5+ frames: {sum(size >= 5 for size in sizes)}")
    print(f"Plan: {args.output}")


if __name__ == "__main__":
    main()