"""Apply a color grade to a video via ffmpeg filter chain.

Two modes:

  1. Preset mode — pick a named preset (e.g. `warm_cinematic`, `neutral_punch`).
     Simple fixed filter chain applied uniformly.

  2. Auto mode (DEFAULT) — analyze the clip mathematically and emit a subtle
     per-clip correction. Samples N frames via ffmpeg, computes mean brightness,
     RMS contrast, saturation. Emits a bounded filter string that corrects
     under-exposure, flatness, and mild desaturation without applying any
     creative color shift. All adjustments capped at ±8% on any axis.

     The goal is "make it look clean without looking graded". Never applies
     creative LUTs, teal/orange splits, or filmic curves. For creative looks,
     use `--preset warm_cinematic` explicitly.

Usage:
    python helpers/grade.py <input> -o <output>                   # auto mode
    python helpers/grade.py <input> -o <output> --preset warm_cinematic
    python helpers/grade.py <input> -o <output> --filter 'eq=contrast=1.1'
    python helpers/grade.py --print-preset warm_cinematic         # print filter only
    python helpers/grade.py --analyze <input>                     # print auto-grade analysis

Can also be imported by render.py: `get_preset(name)` and `auto_grade_for_clip(path, edl_range)`.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
from pathlib import Path


PRESETS: dict[str, str] = {
    # Subtle baseline — barely perceptible cleanup. No color shift.
    # Use when auto-analysis isn't available or when you want a safe floor.
    "subtle": "eq=contrast=1.03:saturation=0.98",

    # Minimal corrective grade: light contrast + subtle S-curve, no color shifts.
    "neutral_punch": (
        "eq=contrast=1.06:brightness=0.0:saturation=1.0,"
        "curves=master='0/0 0.25/0.23 0.75/0.77 1/1'"
    ),

    # OPT-IN creative preset for retro/cinematic looks ONLY. Not a default.
    # +12% contrast, crushed blacks, -12% sat, warm shadows + cool highs, filmic curve.
    # Originally from HEURISTICS §6 — too aggressive for standard launch content.
    "warm_cinematic": (
        "eq=contrast=1.12:brightness=-0.02:saturation=0.88,"
        "colorbalance="
        "rs=0.02:gs=0.0:bs=-0.03:"
        "rm=0.04:gm=0.01:bm=-0.02:"
        "rh=0.08:gh=0.02:bh=-0.05,"
        "curves=master='0/0 0.25/0.22 0.75/0.78 1/1'"
    ),

    # Flat — no grade. Useful as a sentinel for "skip grading this source".
    "none": "",
}


def get_preset(name: str) -> str:
    """Return the ffmpeg filter string for a preset name. Empty string for 'none'."""
    if name not in PRESETS:
        raise KeyError(
            f"unknown preset '{name}'. Available: {', '.join(sorted(PRESETS))}"
        )
    return PRESETS[name]


# -------- Auto grade (data-driven, per-clip) --------------------------------


def _sample_frame_stats(
    video: Path,
    start: float,
    duration: float,
    n_samples: int = 10,
) -> dict[str, float]:
    """Sample N frames from a range and compute brightness/contrast/saturation stats.

    Uses ffmpeg's `signalstats` filter which gives us YMIN, YMAX, YAVG, SATAVG
    etc. in the metadata. We average across the sample range.

    Returns:
        {
          "y_mean":   mean Y (luma) in 0..1,
          "y_std":    approximate stddev of Y across samples (0..1),
          "sat_mean": mean saturation in 0..1,
        }
    """
    # Use signalstats + metadata=print to get per-frame stats
    # Sample fps = n_samples / duration, clamped so we don't over-sample short clips
    fps = max(0.5, min(n_samples / max(duration, 0.1), 10.0))

    with tempfile.NamedTemporaryFile(mode="w+", suffix=".txt", delete=False) as f:
        metadata_path = f.name

    try:
        cmd = [
            "ffmpeg", "-y", "-hide_banner", "-nostats",
            "-ss", f"{start:.3f}",
            "-i", str(video),
            "-t", f"{duration:.3f}",
            "-vf", f"fps={fps:.2f},signalstats,metadata=print:file={metadata_path}",
            "-f", "null", "-",
        ]
        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        # Parse signalstats metadata. Signalstats reports values in the NATIVE
        # bit depth of the decoded frame (8-bit → 0-255, 10-bit → 0-1023). We
        # read YBITDEPTH and normalize by (2^depth - 1) so downstream math is
        # in 0..1 regardless of source bit depth.
        y_avgs: list[float] = []
        y_mins: list[float] = []
        y_maxs: list[float] = []
        sat_avgs: list[float] = []
        bit_depth: int = 8

        def _parse_value(line: str) -> float | None:
            try:
                return float(line.rsplit("=", 1)[1])
            except (ValueError, IndexError):
                return None

        with open(metadata_path) as f:
            for line in f:
                line = line.strip()
                if "lavfi.signalstats.YBITDEPTH" in line:
                    v = _parse_value(line)
                    if v is not None:
                        bit_depth = int(v)
                elif "lavfi.signalstats.YAVG" in line:
                    v = _parse_value(line)
                    if v is not None:
                        y_avgs.append(v)
                elif "lavfi.signalstats.YMIN" in line:
                    v = _parse_value(line)
                    if v is not None:
                        y_mins.append(v)
                elif "lavfi.signalstats.YMAX" in line:
                    v = _parse_value(line)
                    if v is not None:
                        y_maxs.append(v)
                elif "lavfi.signalstats.SATAVG" in line:
                    v = _parse_value(line)
                    if v is not None:
                        sat_avgs.append(v)

        if not y_avgs:
            # Analysis failed — return neutral defaults (no correction)
            return {"y_mean": 0.5, "y_std": 0.18, "sat_mean": 0.25}

        # Normalize by native bit-depth max value
        max_val = (2 ** bit_depth) - 1

        y_mean = (sum(y_avgs) / len(y_avgs)) / max_val
        y_range = (
            ((sum(y_maxs) / len(y_maxs)) - (sum(y_mins) / len(y_mins))) / max_val
            if y_maxs and y_mins
            else 0.7
        )
        sat_mean = ((sum(sat_avgs) / len(sat_avgs)) / max_val) if sat_avgs else 0.25

        return {
            "y_mean": y_mean,
            "y_std": y_range / 4.0,  # range ÷ 4 ≈ stddev for normal-ish distributions
            "sat_mean": sat_mean,
        }
    finally:
        Path(metadata_path).unlink(missing_ok=True)


def auto_grade_for_clip(
    video: Path,
    start: float = 0.0,
    duration: float | None = None,
    verbose: bool = False,
) -> tuple[str, dict[str, float]]:
    """Analyze a clip range and emit a subtle per-clip correction filter.

    Returns (filter_string, stats_dict). The filter is bounded to ±8% on any axis
    and applies NO color shift. It only addresses:
      - Underexposure (lift gamma slightly if too dark)
      - Flatness (tiny contrast boost if range is narrow)
      - Desaturation (tiny sat boost if extremely flat)

    If the clip is already well-balanced, returns the baseline `subtle` preset.
    """
    if duration is None:
        # Probe duration
        probe_cmd = [
            "ffprobe", "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1",
            str(video),
        ]
        try:
            duration = float(subprocess.check_output(probe_cmd).decode().strip())
        except Exception:
            duration = 10.0

    stats = _sample_frame_stats(video, start, duration)

    y_mean = stats["y_mean"]
    y_range = stats["y_std"] * 4.0  # back to range
    sat_mean = stats["sat_mean"]

    # ------ Decision rules ---------------------------------------------------
    # All caps bounded to ±8%. Target "clean, not graded".

    # Contrast: target y_range ≈ 0.72. Boost gently if flat, never reduce.
    contrast_adj = 1.0
    if y_range < 0.65:
        # Map [0.50, 0.65] → [1.08, 1.03]
        t = max(0.0, min(1.0, (y_range - 0.50) / 0.15))
        contrast_adj = 1.08 - 0.05 * t
    else:
        contrast_adj = 1.03  # subtle baseline

    # Gamma: target y_mean ≈ 0.48. Lift gently if too dark.
    gamma_adj = 1.0
    if y_mean < 0.42:
        # Map [0.30, 0.42] → [1.10, 1.02]
        t = max(0.0, min(1.0, (y_mean - 0.30) / 0.12))
        gamma_adj = 1.10 - 0.08 * t
    elif y_mean > 0.60:
        # Slightly overexposed — tiny pullback
        gamma_adj = 0.97

    # Saturation: target sat_mean ≈ 0.25. Never desaturate aggressively;
    # modest boost if very flat. Default to 0.98 (tiny pullback — most digital
    # video is slightly over-saturated on consumer displays).
    sat_adj = 0.98
    if sat_mean < 0.18:
        # Very flat — tiny boost
        sat_adj = 1.04
    elif sat_mean > 0.38:
        # Already punchy — hold
        sat_adj = 0.96

    # Clamp all adjustments hard
    contrast_adj = max(0.94, min(1.08, contrast_adj))
    gamma_adj = max(0.94, min(1.10, gamma_adj))
    sat_adj = max(0.94, min(1.06, sat_adj))

    # Build filter string
    eq_parts = []
    if abs(contrast_adj - 1.0) > 0.005:
        eq_parts.append(f"contrast={contrast_adj:.3f}")
    if abs(gamma_adj - 1.0) > 0.005:
        eq_parts.append(f"gamma={gamma_adj:.3f}")
    if abs(sat_adj - 1.0) > 0.005:
        eq_parts.append(f"saturation={sat_adj:.3f}")

    if not eq_parts:
        filter_string = ""
    else:
        filter_string = "eq=" + ":".join(eq_parts)

    if verbose:
        print(f"  auto-grade stats:")
        print(f"    y_mean={y_mean:.3f}  y_range={y_range:.3f}  sat_mean={sat_mean:.3f}")
        print(f"    → contrast={contrast_adj:.3f}  gamma={gamma_adj:.3f}  sat={sat_adj:.3f}")
        print(f"    → filter: {filter_string or '(empty)'}")

    return filter_string, stats


def apply_grade(input_path: Path, output_path: Path, filter_string: str) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    if not filter_string:
        cmd = [
            "ffmpeg", "-y", "-i", str(input_path),
            "-c", "copy", str(output_path),
        ]
    else:
        cmd = [
            "ffmpeg", "-y", "-i", str(input_path),
            "-vf", filter_string,
            "-c:v", "libx264", "-preset", "fast", "-crf", "18",
            "-pix_fmt", "yuv420p",
            "-c:a", "copy",
            "-movflags", "+faststart",
            str(output_path),
        ]
    subprocess.run(cmd, check=True)


def main() -> None:
    ap = argparse.ArgumentParser(description="Apply a color grade via ffmpeg filter chain")
    ap.add_argument("input", type=Path, nargs="?", help="Input video")
    ap.add_argument("-o", "--output", type=Path, help="Output video")
    ap.add_argument(
        "--preset",
        type=str,
        default=None,
        choices=list(PRESETS.keys()),
        help="Grade preset. Omit for auto mode (default).",
    )
    ap.add_argument(
        "--filter",
        type=str,
        default=None,
        help="Raw ffmpeg filter string. Overrides --preset.",
    )
    ap.add_argument(
        "--analyze",
        type=Path,
        default=None,
        help="Analyze a clip and print the auto-grade filter it would produce. No output written.",
    )
    ap.add_argument(
        "--print-preset",
        type=str,
        default=None,
        help="Print the filter string for a preset and exit. No input/output needed.",
    )
    ap.add_argument(
        "--list-presets",
        action="store_true",
        help="List available presets and exit.",
    )
    args = ap.parse_args()

    if args.list_presets:
        for name, f in PRESETS.items():
            print(f"{name}:")
            print(f"  {f}" if f else "  (no filter)")
            print()
        return

    if args.print_preset is not None:
        print(get_preset(args.print_preset))
        return

    if args.analyze is not None:
        if not args.analyze.exists():
            sys.exit(f"input not found: {args.analyze}")
        filter_string, stats = auto_grade_for_clip(args.analyze, verbose=True)
        print(f"\nfilter: {filter_string or '(none)'}")
        print(f"stats:  {json.dumps(stats, indent=2)}")
        return

    if not args.input or not args.output:
        ap.error("input and -o/--output are required unless using --analyze/--print-preset/--list-presets")

    if not args.input.exists():
        sys.exit(f"input not found: {args.input}")

    # Decide filter string
    if args.filter is not None:
        filter_string = args.filter
    elif args.preset is not None:
        filter_string = get_preset(args.preset)
    else:
        # Auto mode (default)
        filter_string, _ = auto_grade_for_clip(args.input, verbose=True)

    print(f"grading {args.input.name} → {args.output.name}")
    if filter_string:
        print(f"  filter: {filter_string[:120]}{'...' if len(filter_string) > 120 else ''}")
    else:
        print("  filter: (none — copy)")

    apply_grade(args.input, args.output, filter_string)
    print(f"done: {args.output}")


if __name__ == "__main__":
    main()