#!/usr/bin/env python3
"""Create DaVinci-ready H.264 ping-pong clips from YellowMellow burst folders."""

import argparse
import csv
import json
import math
import os
import re
import shutil
import subprocess
import tempfile
from datetime import datetime
from pathlib import Path


DEFAULT_SOURCE = Path(
    "/Users/Sam/Sam Storage/Photography/Photography (2026)/"
    "YellowMellow SummerFall 2026 /JPEGs Sorted"
)
DEFAULT_OUTPUT = Path(__file__).resolve().parent / "YellowMellow_Boomerang_H264"
CADENCES = (6, 8, 10)


def capture_seconds(item, raw_root):
    if "DateTimeOriginal" not in item:
        raw = raw_root / (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 frame_number(path):
    return int(re.search(r"(\d+)$", path.stem).group(1))


def choose_fps(items, raw_root):
    ordered = sorted(items, key=lambda item: frame_number(Path(item["FileName"])))
    times = [capture_seconds(item, raw_root) for item in ordered]
    gaps = [times[i + 1] - times[i] for i in range(len(times) - 1)]
    drive_gaps = sorted(gap for gap in gaps if 0.04 <= gap <= 0.25)
    if not drive_gaps:
        return 10
    median = drive_gaps[len(drive_gaps) // 2]
    measured = 1 / median
    return min(CADENCES, key=lambda fps: abs(fps - measured))


def render_burst(burst, items, raw_root, output_dir, target_seconds, overwrite):
    frames = sorted(burst.glob("*.jpg"), key=frame_number)
    metadata = {item["FileName"]: item for item in items}
    fps = choose_fps(items, raw_root)
    ping_pong = frames + frames[-2:0:-1]
    cycle_duration = len(ping_pong) / fps
    cycles = max(1, math.ceil(target_seconds / cycle_duration))
    output_frames = len(ping_pong) * cycles
    output = output_dir / f"{burst.name}_Boomerang_{fps}fps.mp4"
    if output.exists() and not overwrite:
        print(f"SKIP {output.name}", flush=True)
        return {
            "burst": burst.name, "source_frames": len(frames), "fps": fps,
            "cycles": cycles, "output_frames": output_frames,
            "duration_seconds": round(output_frames / fps, 3),
            "output_file": output.name, "file_bytes": output.stat().st_size,
            "corrections": "existing file retained",
        }

    corrections = set()
    with tempfile.TemporaryDirectory(prefix="yellowmellow_boomerang_") as temp_name:
        temp_dir = Path(temp_name)
        for frame in frames:
            if not metadata[frame.name].get("Orientation"):
                corrections.add(f"sequence-level rotation restored for {frame.stem}.jpg")

        sequence = ping_pong * cycles
        for index, frame in enumerate(sequence):
            (temp_dir / f"frame{index:06d}.jpg").symlink_to(frame)

        command = [
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
            "-noautorotate", "-framerate", str(fps), "-start_number", "0",
            "-i", str(temp_dir / "frame%06d.jpg"),
            "-vf", (
                "transpose=cclock,scale=1366:2048:force_original_aspect_ratio=decrease:"
                "force_divisible_by=2:flags=lanczos:in_range=full:out_range=tv,"
                "pad=1366:2048:(ow-iw)/2:(oh-ih)/2:black,setsar=1,format=yuv420p"
            ),
            "-c:v", "libx264", "-preset", "medium", "-crf", "15",
            "-profile:v", "high", "-level:v", "5.1", "-fps_mode", "cfr",
            "-colorspace", "bt709", "-color_primaries", "bt709",
            "-color_trc", "bt709", "-movflags", "+faststart",
            "-metadata", f"title={burst.name} ping-pong burst",
            str(output),
        ]
        subprocess.run(command, check=True)

    print(
        f"DONE {output.name} | {len(frames)} source frames | {fps} fps | "
        f"{cycles} cycle(s) | {output_frames / fps:.2f}s",
        flush=True,
    )
    return {
        "burst": burst.name, "source_frames": len(frames), "fps": fps,
        "cycles": cycles, "output_frames": output_frames,
        "duration_seconds": round(output_frames / fps, 3),
        "output_file": output.name, "file_bytes": output.stat().st_size,
        "corrections": "; ".join(sorted(corrections)),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--target-seconds", type=float, default=12.0)
    parser.add_argument("--only", help="Render one folder, e.g. Burst19")
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    for tool in ("exiftool", "ffmpeg"):
        if not shutil.which(tool):
            raise SystemExit(f"Required tool not found: {tool}")
    if not args.source.is_dir():
        raise SystemExit(f"Source folder not found: {args.source}")
    args.output.mkdir(parents=True, exist_ok=True)

    metadata = json.loads(subprocess.check_output([
        "exiftool", "-j", "-r", "-ext", "jpg", "-DateTimeOriginal",
        "-SubSecTimeOriginal", "-Orientation", "-FileName", "-Directory",
        str(args.source),
    ]))
    by_directory = {}
    for item in metadata:
        by_directory.setdefault(item["Directory"], []).append(item)

    bursts = sorted(args.source.glob("Burst*"), key=lambda path: int(path.name[5:7]))
    if args.only:
        bursts = [burst for burst in bursts if burst.name.startswith(args.only)]
        if not bursts:
            raise SystemExit(f"No burst folder matched: {args.only}")

    raw_root = args.source.parent
    rows = []
    for burst in bursts:
        rows.append(render_burst(
            burst, by_directory[str(burst)], raw_root, args.output,
            args.target_seconds, args.overwrite,
        ))

    manifest = args.output / "burst_render_manifest.csv"
    with manifest.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=rows[0].keys())
        writer.writeheader()
        writer.writerows(rows)
    print(f"MANIFEST {manifest}", flush=True)


if __name__ == "__main__":
    main()