#!/usr/bin/env python3
"""Develop grouped CR3 bursts and render DaVinci-ready boomerang clips."""

import argparse
import csv
import json
import math
import re
import shutil
import subprocess
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path


CADENCES = (6, 8, 10)
PRORES_PROFILES = {"lt": "1", "422": "2", "hq": "3"}


def frame_number(path):
    match = re.search(r"(\d+)$", path.stem)
    if not match:
        raise ValueError(f"No numeric frame number in {path.name}")
    return int(match.group(1))


def capture_seconds(item):
    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 choose_fps(items):
    timed = [item for item in items if item.get("DateTimeOriginal")]
    timed.sort(key=lambda item: frame_number(Path(item["FileName"])))
    times = [capture_seconds(item) for item in timed]
    gaps = sorted(
        times[index + 1] - times[index]
        for index in range(len(times) - 1)
        if 0.04 <= times[index + 1] - times[index] <= 0.25
    )
    if not gaps:
        return 10
    measured = 1 / gaps[len(gaps) // 2]
    return min(CADENCES, key=lambda fps: abs(fps - measured))


def raw_orientation(item):
    label = str(item.get("Orientation", ""))
    if "270" in label:
        return "270"
    if "90" in label:
        return "90"
    if "180" in label:
        return "180"
    return "0"


def prepare_frame(raw, output, mode, width, height, orientation):
    if mode == "developed":
        intermediate = output.with_suffix(".developed.tiff")
        try:
            subprocess.run([
                "dcraw_emu", "-w", "-H", "2", "-q", "3", "-6", "-o", "1",
                "-T", "-Z", str(intermediate), str(raw),
            ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            command = ["magick", str(intermediate)]
            # dcraw_emu applies the CR3 orientation itself.
        finally:
            pass
    else:
        intermediate = output.with_suffix(".embedded.jpg")
        with intermediate.open("wb") as handle:
            subprocess.run(
                ["exiftool", "-b", "-JpgFromRaw", str(raw)],
                check=True, stdout=handle, stderr=subprocess.DEVNULL,
            )
        command = ["magick", str(intermediate)]
        if orientation != "0":
            command += ["-rotate", orientation]

    try:
        command += [
            "-resize", f"{width}x{height}>", "-gravity", "center",
            "-background", "black", "-extent", f"{width}x{height}",
            "-colorspace", "sRGB", "-depth", "16",
            "-define", "png:color-type=2", "-define", "png:compression-level=1",
            str(output),
        ]
        subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    finally:
        if intermediate.exists():
            intermediate.unlink()
    return output


def render_one(burst, jpeg_items, raw_items, args):
    guide_frames = sorted([
        path for path in burst.iterdir() if path.suffix.lower() in (".jpg", ".jpeg")
    ], key=frame_number)
    if len(guide_frames) < 2:
        raise ValueError(f"{burst.name} has fewer than two guide frames")
    source_fps = args.fps or choose_fps(jpeg_items or raw_items)
    fps, retimed = source_fps, False
    cycle_frames = 2 * len(guide_frames) - 2
    base_cycle = cycle_frames / fps
    if len(guide_frames) < args.minimum_frames and not args.include_review_fragments:
        print(f"REVIEW {burst.name} | only {len(guide_frames)} source frames", flush=True)
        return {
            "burst": burst.name, "status": "fragment_review", "source_frames": len(guide_frames),
            "capture_fps": source_fps, "playback_fps": fps, "retimed": retimed,
            "cycles": 0, "base_cycle_seconds": round(base_cycle, 3), "output_frames": 0,
            "duration_seconds": 0, "output_file": "", "file_bytes": 0,
        }

    cycles = max(1, math.ceil(args.target_seconds / base_cycle))
    output_frames = cycle_frames * cycles
    suffix = ".mov" if args.codec == "prores" else ".mp4"
    output = args.output / f"{burst.name}_CR3_Boomerang_{fps}fps_{args.codec}{suffix}"
    if output.exists() and not args.overwrite:
        print(f"SKIP {output.name}", flush=True)
        return {
            "burst": burst.name, "status": "existing", "source_frames": len(guide_frames),
            "capture_fps": source_fps, "playback_fps": fps, "retimed": retimed,
            "cycles": cycles, "base_cycle_seconds": round(base_cycle, 3),
            "output_frames": output_frames, "duration_seconds": round(output_frames / fps, 3),
            "output_file": output.name, "file_bytes": output.stat().st_size,
        }

    raw_by_name = {Path(item["FileName"]).stem: item for item in raw_items}
    with tempfile.TemporaryDirectory(prefix=f"{burst.name}_", dir=args.temp_root) as temp_name:
        temp_dir = Path(temp_name)
        prepared = {}

        def task(guide):
            raw = args.raw_root / f"{guide.stem}.CR3"
            if not raw.exists():
                raise FileNotFoundError(raw)
            item = raw_by_name.get(guide.stem, {})
            normalized = temp_dir / f"source_{guide.stem}.png"
            return guide, prepare_frame(
                raw, normalized, args.raw_mode, args.width, args.height, raw_orientation(item)
            )

        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            futures = [pool.submit(task, guide) for guide in guide_frames]
            for future in as_completed(futures):
                guide, normalized = future.result()
                prepared[guide] = normalized

        ping_pong = guide_frames + guide_frames[-2:0:-1]
        sequence = ping_pong * cycles
        for index, guide in enumerate(sequence):
            (temp_dir / f"frame{index:06d}.png").symlink_to(prepared[guide])

        if args.codec == "prores":
            codec_args = [
                "-c:v", "prores_ks", "-profile:v", PRORES_PROFILES[args.prores_profile],
                "-vendor", "apl0", "-pix_fmt", "yuv422p10le",
            ]
        else:
            codec_args = [
                "-c:v", "libx264", "-preset", "medium", "-crf", "14",
                "-profile:v", "high", "-pix_fmt", "yuv420p", "-movflags", "+faststart",
            ]
        subprocess.run([
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
            "-framerate", str(fps), "-start_number", "0",
            "-i", str(temp_dir / "frame%06d.png"),
            *codec_args, "-fps_mode", "cfr", "-colorspace", "bt709",
            "-color_primaries", "bt709", "-color_trc", "bt709",
            "-metadata", f"title={burst.name} CR3 ping-pong burst", str(output),
        ], check=True)

    print(
        f"DONE {output.name} | {len(guide_frames)} RAWs | {fps} fps | "
        f"{cycles} cycle(s) | {output_frames / fps:.2f}s",
        flush=True,
    )
    return {
        "burst": burst.name, "status": "rendered", "source_frames": len(guide_frames),
        "capture_fps": source_fps, "playback_fps": fps, "retimed": retimed,
        "cycles": cycles, "base_cycle_seconds": round(base_cycle, 3),
        "output_frames": output_frames, "duration_seconds": round(output_frames / fps, 3),
        "output_file": output.name, "file_bytes": output.stat().st_size,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--groups", type=Path, required=True)
    parser.add_argument("--raw-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--temp-root", type=Path)
    parser.add_argument("--raw-mode", choices=("developed", "embedded"), default="developed")
    parser.add_argument("--codec", choices=("prores", "h264"), default="prores")
    parser.add_argument("--prores-profile", choices=tuple(PRORES_PROFILES), default="hq")
    parser.add_argument("--width", type=int, default=2560)
    parser.add_argument("--height", type=int, default=3840)
    parser.add_argument("--target-seconds", type=float, default=3.0)
    parser.add_argument("--minimum-frames", type=int, default=5)
    parser.add_argument("--minimum-cycle", type=float, help=argparse.SUPPRESS)
    parser.add_argument(
        "--include-review-fragments", "--include-microbursts",
        dest="include_review_fragments", action="store_true",
    )
    parser.add_argument("--fps", type=int)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--only")
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    for tool in ("dcraw_emu", "exiftool", "magick", "ffmpeg"):
        if not shutil.which(tool):
            raise SystemExit(f"Required tool not found: {tool}")
    if not args.groups.is_dir() or not args.raw_root.is_dir():
        raise SystemExit("Both --groups and --raw-root must be directories")
    if args.width % 2 or args.height % 2:
        raise SystemExit("Output width and height must be even")
    args.groups = args.groups.resolve()
    args.raw_root = args.raw_root.resolve()
    args.output = args.output.resolve()
    args.output.mkdir(parents=True, exist_ok=True)
    if args.temp_root:
        args.temp_root = args.temp_root.resolve()
        args.temp_root.mkdir(parents=True, exist_ok=True)

    jpeg_metadata = json.loads(subprocess.check_output([
        "exiftool", "-j", "-r", "-ext", "jpg", "-ext", "jpeg",
        "-DateTimeOriginal", "-SubSecTimeOriginal", "-Orientation",
        "-FileName", "-Directory", str(args.groups),
    ]))
    raw_metadata = json.loads(subprocess.check_output([
        "exiftool", "-fast2", "-j", "-ext", "cr3", "-DateTimeOriginal",
        "-SubSecTimeOriginal", "-Orientation", "-FileName", str(args.raw_root),
    ]))
    jpeg_by_dir = {}
    for item in jpeg_metadata:
        jpeg_by_dir.setdefault(item["Directory"], []).append(item)
    raw_by_stem = {Path(item["FileName"]).stem: item for item in raw_metadata}

    bursts = sorted(args.groups.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("No burst folders matched")

    rows = []
    for burst in bursts:
        guides = [path for path in burst.iterdir() if path.suffix.lower() in (".jpg", ".jpeg")]
        raw_items = [raw_by_stem.get(path.stem, {}) for path in guides]
        rows.append(render_one(burst, jpeg_by_dir.get(str(burst), []), raw_items, args))

    manifest = args.output / "cr3_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}")


if __name__ == "__main__":
    main()