#!/usr/bin/env python3
"""Apply one colour grade, and prove what it did with a before/after sheet.

Grading is the step most often applied twice — once in the base render and again on
the master — which doubles the contrast and buries the shadows. So this reports the
measured change: average luma, contrast and saturation before and after, plus how
much of the frame is clipped at either end. A grade that crushes 2% of the picture
to pure black is visible as a number here rather than as a complaint later.

The order of operations matters and is enforced by the report, not just documented:
grade the base, then burn captions. Grading over burned-in text shifts the caption
colour away from the value the preset was designed with.
"""
from __future__ import annotations

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

sys.path.insert(0, str(Path(__file__).resolve().parent))

from media_probe import probe  # noqa: E402
from render_core import PROFILES, capabilities, output_arguments  # noqa: E402
from skill_config import GRADES, emit, require_binary, utf8_stdout  # noqa: E402

PRESETS = GRADES / "grade_presets.json"


def library() -> dict:
    if not PRESETS.exists():
        raise SystemExit(f"Grade presets not found: {PRESETS}")
    return json.loads(PRESETS.read_text(encoding="utf-8"))


def measure(video: Path, at: float) -> dict | None:
    """Luma statistics of one frame, read straight from the decoded pixels."""
    try:
        import numpy
    except ImportError:
        return None
    completed = subprocess.run(
        ["ffmpeg", "-v", "error", "-nostdin", "-ss", f"{at:.3f}", "-i", str(video),
         "-frames:v", "1", "-vf", "scale=320:-2", "-f", "rawvideo",
         "-pix_fmt", "rgb24", "-"],
        capture_output=True,
    )
    if completed.returncode or not completed.stdout:
        return None
    data = numpy.frombuffer(completed.stdout, dtype=numpy.uint8).astype(numpy.float32)
    if data.size < 3:
        return None
    pixels = data[: (data.size // 3) * 3].reshape(-1, 3)
    luma = pixels @ numpy.array([0.2126, 0.7152, 0.0722], dtype=numpy.float32)
    maximum = pixels.max(axis=1)
    minimum = pixels.min(axis=1)
    saturation = numpy.where(maximum > 0, (maximum - minimum) / numpy.maximum(maximum, 1), 0)
    return {
        "mean_luma": round(float(luma.mean()), 2),
        "contrast_stdev": round(float(luma.std()), 2),
        "mean_saturation": round(float(saturation.mean()), 4),
        "clipped_black_percent": round(float((luma < 2).mean() * 100), 3),
        "clipped_white_percent": round(float((luma > 253).mean() * 100), 3),
        "percentile_5": round(float(numpy.percentile(luma, 5)), 1),
        "percentile_95": round(float(numpy.percentile(luma, 95)), 1),
    }


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("input", nargs="?")
    parser.add_argument("output", nargs="?")
    parser.add_argument("--grade", default="neutral_clean")
    parser.add_argument("--extras", default="",
                        help="Comma-separated extras: vignette,grain_light,sharpen_web,...")
    parser.add_argument("--strength", type=float, default=1.0,
                        help="Blend the graded picture back over the original, 0-1")
    parser.add_argument("--profile", default="master", choices=list(PROFILES))
    parser.add_argument("--codec", default="h264", choices=["h264", "hevc"])
    parser.add_argument("--cpu", action="store_true")
    parser.add_argument("--contact-sheet", help="Write a before/after comparison here")
    parser.add_argument("--sample-at", type=float, help="Frame time used for the comparison")
    parser.add_argument("--list", action="store_true")
    args = parser.parse_args()

    data = library()
    grades = {item["id"]: item for item in data["grades"]}
    if args.list:
        emit({
            "how_to_use": data["how_to_use"],
            "order_of_operations": data["order_of_operations"],
            "grades": {key: {"label": item["label"], "use": item["use"],
                             "warnings": item["warnings"]}
                       for key, item in grades.items()},
            "extras": {key: item["use"] for key, item in data["extras"].items()},
        })
        return
    if not args.input or not args.output:
        parser.error("input and output are required unless --list is used")
    if args.grade not in grades:
        raise SystemExit(f"Unknown grade '{args.grade}'. Known: {', '.join(grades)}")

    require_binary("ffmpeg")
    source = Path(args.input).resolve()
    if not source.is_file():
        raise SystemExit(f"Not found: {source}")
    info = probe(source)
    duration = float(info.get("duration") or 0.0)
    sample_at = args.sample_at if args.sample_at is not None else min(max(duration * 0.4, 0.5),
                                                                     max(duration - 0.5, 0.5))

    chain = grades[args.grade]["filter"]
    extras = [name.strip() for name in args.extras.split(",") if name.strip()]
    for name in extras:
        if name not in data["extras"]:
            raise SystemExit(f"Unknown extra '{name}'. Known: {', '.join(data['extras'])}")
        chain += "," + data["extras"][name]["filter"]

    strength = max(0.0, min(1.0, args.strength))
    if strength >= 0.999:
        video_filter = chain
    else:
        # Blending the graded picture back over the original is how a preset gets
        # dialled down without re-authoring every curve in it.
        video_filter = (
            f"split=2[grade_a][grade_b];[grade_b]{chain}[graded];"
            f"[grade_a][graded]blend=all_mode=normal:all_opacity={strength:.3f}"
        )

    caps = capabilities()
    encoder_arguments, selection = output_arguments(
        args.codec, args.profile, prefer_gpu=not args.cpu, caps=caps
    )
    output = Path(args.output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    command = [
        "ffmpeg", "-y", "-v", "error", "-nostdin", "-i", str(source),
        "-filter_complex", f"[0:v]{video_filter}[vout]",
        "-map", "[vout]", "-map", "0:a?",
        *encoder_arguments, "-c:a", "copy",
        str(output),
    ]
    completed = subprocess.run(command, capture_output=True, text=True, errors="replace")
    if completed.returncode:
        # `-c:a copy` fails when the input audio codec is not MP4-compatible.
        command[command.index("-c:a") + 1] = "aac"
        completed = subprocess.run(command, capture_output=True, text=True, errors="replace")
    if completed.returncode:
        emit({"status": "failed", "grade": args.grade,
              "tail": (completed.stderr or "").strip()[-1500:]})
        raise SystemExit(completed.returncode)

    before = measure(source, sample_at)
    after = measure(output, sample_at)
    sheet = None
    if args.contact_sheet:
        sheet_path = Path(args.contact_sheet)
        sheet_path.parent.mkdir(parents=True, exist_ok=True)
        sheet_command = [
            "ffmpeg", "-y", "-v", "error",
            "-ss", f"{sample_at:.3f}", "-i", str(source),
            "-ss", f"{sample_at:.3f}", "-i", str(output),
            "-filter_complex",
            "[0:v]scale=420:-2[a];[1:v]scale=420:-2[b];[a][b]hstack=inputs=2",
            "-frames:v", "1", str(sheet_path),
        ]
        result = subprocess.run(sheet_command, capture_output=True, text=True, errors="replace")
        sheet = str(sheet_path.resolve()) if result.returncode == 0 else None

    warnings = list(grades[args.grade]["warnings"])
    if after and after["clipped_black_percent"] > 1.5:
        warnings.append(
            f"{after['clipped_black_percent']}% кадра ушло в чистый чёрный — "
            f"тени потеряны. Уменьши --strength или возьми мягче грейд."
        )
    if after and after["clipped_white_percent"] > 1.0:
        warnings.append(
            f"{after['clipped_white_percent']}% кадра выбито в белый — света срезаны."
        )
    # A ratio alone is misleading: the metric is (max-min)/max on RGB, so it rises
    # with contrast, and on a desaturated source any contrast grade trips 1.6x.
    # The warning only matters once the absolute level is high enough to look
    # artificial.
    if before and after and after["mean_saturation"] > 0.45 and             after["mean_saturation"] > before["mean_saturation"] * 1.5:
        warnings.append(
            f"Насыщенность {after['mean_saturation']:.2f} и выросла в "
            f"{after['mean_saturation'] / max(before['mean_saturation'], 1e-6):.1f} раза — "
            f"проверь тон кожи крупным планом."
        )

    emit({
        "status": "ok",
        "grade": args.grade,
        "label": grades[args.grade]["label"],
        "use": grades[args.grade]["use"],
        "extras": extras,
        "strength": strength,
        "input": str(source),
        "output": str(output),
        "sampled_at": round(sample_at, 3),
        "before": before,
        "after": after,
        "measured_change": None if not (before and after) else {
            "mean_luma": round(after["mean_luma"] - before["mean_luma"], 2),
            "contrast_stdev": round(after["contrast_stdev"] - before["contrast_stdev"], 2),
            "mean_saturation": round(after["mean_saturation"] - before["mean_saturation"], 4),
        },
        "contact_sheet": sheet,
        "encoder": selection,
        "warnings": warnings,
        "order_reminder": data["order_of_operations"],
    })


if __name__ == "__main__":
    main()