# Adapted from browser-use/video-use (MIT, (c) 2026 Browser Use). See THIRD_PARTY_NOTICES.md.
"""Color grade presets, apply a grade to a video, or compare grades on one frame.

Usage:
    uv run python helpers/grade.py --list
    uv run python helpers/grade.py <in> -o <out> --preset vivo
    uv run python helpers/grade.py <in> -o <out> --filter "eq=contrast=1.1:saturation=1.2"
    uv run python helpers/grade.py <video> --candidates "a=eq=saturation=1.3;b=quente;original=" --frame 4.2 -o cmp.png

Imported by render_cut.py: `resolve_grade(value)` turns an EDL "grade" field (preset name or raw
ffmpeg filter) into a filter string, `check_filter(f)` validates it on a tiny test input.
Every grade runs after an 8-bit `format=yuv420p` guard (see `graded_chain`).
"""

from __future__ import annotations

import argparse
import math
import re
import subprocess
import sys
from functools import lru_cache
from pathlib import Path

from PIL import Image, ImageDraw

from common import has_ffmpeg_filter, probe_video


def _apple_log_curve() -> str:
    """Basic Apple Log -> Rec.709 curve (no LUT): decode the log curve, soft shoulder, 709 OETF."""
    r0, rt, c = -0.05641088, 0.01, 47.28711236
    beta, gamma, delta = 0.00964052, 0.08550479, 0.69336945
    p_t = gamma * math.log2(rt + beta) + delta

    def to_linear(x: float) -> float:
        if x >= p_t:
            return 2 ** ((x - delta) / gamma) - beta
        return r0 + math.sqrt(max(0.0, x) / c)

    def oetf(v: float) -> float:
        v = min(1.0, max(0.0, v))
        return 4.5 * v if v < 0.018 else 1.099 * v ** 0.45 - 0.099

    pts = []
    last = -1.0
    for x in (0.0, 0.12, 0.16, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.8, 1.0):
        lin = max(0.0, to_linear(x)) * 1.25          # exposure so 18% grey lands near 709 mid grey
        tm = lin * (1 + lin / 100.0) / (1 + lin)      # extended Reinhard, white point 10
        y = round(max(last + 0.001 if last >= 0 else 0.0, oetf(tm)), 4)
        y = min(y, 1.0)
        pts.append(f"{x:g}/{y:g}")
        last = y
    return "curves=interp=pchip:all='" + " ".join(pts) + "'"


# name -> (filter, PT-BR description). Keep names stable: they are referenced by EDLs.
PRESETS: dict[str, tuple[str, str]] = {
    "none": ("", "Sem ajuste de cor."),
    "natural": (
        "curves=interp=pchip:master='0/0 0.1/0.093 0.5/0.515 0.9/0.92 1/1',eq=saturation=1.05",
        "Contraste leve, cor fiel. Seguro para qualquer vídeo.",
    ),
    "vivo": (
        "curves=interp=pchip:master='0/0 0.12/0.1 0.5/0.52 0.88/0.915 1/1',vibrance=intensity=0.2,eq=saturation=1.05",
        "Mais contraste e cor, sem estourar a pele. Bom para anúncio.",
    ),
    "quente": (
        "colortemperature=temperature=4800:mix=0.75:pl=0.9,"
        "curves=interp=pchip:master='0/0 0.1/0.093 0.5/0.515 0.9/0.92 1/1',eq=saturation=1.04",
        "Tom mais quente e acolhedor, pele dourada.",
    ),
    "frio": (
        "colortemperature=temperature=8500:mix=0.7:pl=0.9,"
        "curves=interp=pchip:master='0/0 0.1/0.093 0.5/0.515 0.9/0.92 1/1',eq=saturation=0.97",
        "Tom mais frio e limpo, visual tech.",
    ),
    "apple_log_basico": (
        _apple_log_curve() + ",eq=saturation=1.22",
        "Só para vídeo gravado em Apple Log (imagem cinza e lavada). Converte para cor normal.",
    ),
}


def get_preset(name: str) -> str:
    if name not in PRESETS:
        raise KeyError(f"preset de cor desconhecido: {name!r}. Existem: {', '.join(PRESETS)}")
    return PRESETS[name][0]


def resolve_grade(value: str | None) -> str:
    """EDL grade field -> ffmpeg filter string ('' = no grade). Raises ValueError (PT-BR)."""
    if value is None:
        return ""
    text = str(value).strip()
    if not text or text == "none":
        return ""
    if text in PRESETS:
        return PRESETS[text][0]
    if "=" in text or "," in text:
        return text
    if re.fullmatch(r"[A-Za-z0-9_]+", text) and has_ffmpeg_filter(text):
        return text
    raise ValueError(f"grade {text!r} não é um preset ({', '.join(PRESETS)}) nem um filtro do ffmpeg")


@lru_cache(maxsize=None)
def _has_filter(name: str) -> bool:
    return has_ffmpeg_filter(name)


def hdr_to_sdr_chain(info: dict) -> str:
    """Filters that turn an HDR frame (HLG or PQ) into SDR BT.709 8-bit; '' for SDR sources.

    The one conversion used by render_cut, the grade comparison, contact sheets and timeline
    views, so the frame judged for color is the frame the cut shows. With zscale+tonemap it
    tonemaps; without them (Homebrew ffmpeg) it converts the primaries and matrix with
    `colorspace`, no extra contrast or saturation (a grade picked on the preview must not stack
    on a hidden boost).
    """
    if not info.get("is_hdr"):
        return ""
    if _has_filter("zscale") and _has_filter("tonemap"):
        return ("zscale=t=linear:npl=100,format=gbrpf32le,zscale=p=bt709,"
                "tonemap=tonemap=hable:desat=0,zscale=t=bt709:m=bt709:r=tv,format=yuv420p")
    return ("format=yuv420p10le,colorspace=all=bt709:iall=bt2020:itrc=bt2020-10:"
            "irange=tv:range=tv:fast=0:format=yuv420p")


def graded_chain(filter_str: str) -> str:
    """8-bit guard before the grade and after it, so every grade sees and returns yuv420p."""
    if not filter_str:
        return "format=yuv420p"
    return f"format=yuv420p,{filter_str},format=yuv420p"


def check_filter(filter_str: str) -> str | None:
    """Run the filter on a tiny synthetic clip. Returns None when fine, else the error text."""
    if not filter_str:
        return None
    proc = subprocess.run(
        ["ffmpeg", "-v", "error", "-nostdin", "-f", "lavfi", "-i", "testsrc2=s=64x64:d=0.1:r=10",
         "-vf", graded_chain(filter_str), "-f", "null", "-"],
        capture_output=True, text=True)
    if proc.returncode != 0:
        lines = [ln.split("] ", 1)[-1] for ln in (proc.stderr or "").splitlines() if ln.strip()]
        useful = [ln for ln in lines if "Error applying" in ln or "No such filter" in ln or "Invalid" in ln]
        return (useful or lines or ["erro desconhecido"])[0][:400]
    return None


def parse_candidates(spec: str) -> list[tuple[str, str]]:
    """'a=<filter>;b=<preset>;original=' -> [(label, filter)]."""
    out: list[tuple[str, str]] = []
    for part in spec.split(";"):
        part = part.strip()
        if not part:
            continue
        if "=" not in part:
            label, value = part, part          # bare preset name
        else:
            label, value = part.split("=", 1)
        label, value = label.strip(), value.strip()
        if not label:
            raise ValueError(f"candidato sem nome em {part!r} (use nome=filtro)")
        out.append((label, resolve_grade(value)))
    if not out:
        raise ValueError("nenhum candidato em --candidates")
    return out


def candidates_montage(video: Path, candidates: list[tuple[str, str]], t: float, out: Path,
                       tile_width: int = 0) -> Path:
    from contact_sheet import BG, DIM, FG, PANEL, extract_frame, ui_font

    info = probe_video(video)
    vertical = (info.get("display_height") or 0) >= (info.get("display_width") or 0)
    tw = tile_width or (420 if vertical else 640)
    tw -= tw % 2
    fps, dur = info.get("fps") or 30.0, info.get("duration") or 0.0
    tiles: list[Image.Image] = []
    for label, flt in candidates:
        err = check_filter(flt)
        if err:
            raise ValueError(f"o candidato '{label}' tem filtro inválido: {err}")
        # Grade at a decent size first (grades are resolution independent), then downscale.
        # HDR sources go through the same SDR conversion as render_cut, before the grade.
        work_w = min(1080, info.get("display_width") or 1080)
        work_w -= work_w % 2
        hdr = hdr_to_sdr_chain(info)
        chain = f"scale={work_w}:-2:flags=lanczos,{hdr + ',' if hdr else ''}{graded_chain(flt)}"
        tiles.append(extract_frame(video, t, width=tw, vf=chain, fps=fps, duration=dur, hdr=False))
    th = tiles[0].size[1]
    gap, pad, head, lab_h = 12, 20, (84 if info.get("is_hdr") else 60), 64
    n = len(tiles)
    W = pad * 2 + n * tw + (n - 1) * gap
    H = head + th + lab_h + pad
    sheet = Image.new("RGB", (W, H), BG)
    d = ImageDraw.Draw(sheet)
    d.text((pad, 18), f"{video.name} · comparação de cor em {t:.2f}s", fill=FG, font=ui_font(24, bold=True))
    f_lab, f_small = ui_font(20, bold=True), ui_font(14)
    if info.get("is_hdr"):
        d.text((pad, 54), "fonte HDR: convertida para SDR do mesmo jeito que no corte", fill=DIM, font=f_small)
    for i, ((label, flt), img) in enumerate(zip(candidates, tiles)):
        x = pad + i * (tw + gap)
        sheet.paste(img.resize((tw, th)) if img.size != (tw, th) else img, (x, head))
        d.rectangle((x, head + th, x + tw - 1, head + th + lab_h - 1), fill=PANEL)
        d.text((x + 10, head + th + 8), label, fill=FG, font=f_lab)
        desc = flt or "sem filtro"
        preset = next((k for k, (f, _) in PRESETS.items() if f == flt and f), None)
        if preset:
            desc = PRESETS[preset][1] if preset == label else f"preset {preset}"
        while d.textlength(desc, font=f_small) > tw - 20 and len(desc) > 6:
            desc = desc[:-4] + "…"
        d.text((x + 10, head + th + 38), desc, fill=DIM, font=f_small)
    out.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(out, "PNG", optimize=True)
    return out


def apply_grade(src: Path, dst: Path, filter_str: str) -> None:
    info = probe_video(src)
    cmd = ["ffmpeg", "-y", "-v", "error", "-stats", "-nostdin", "-i", str(src), "-map", "0:v:0"]
    if info.get("has_audio"):
        cmd += ["-map", "0:a:0", "-c:a", "copy"]
    hdr = hdr_to_sdr_chain(info)
    cmd += ["-vf", (hdr + "," if hdr else "") + graded_chain(filter_str)
            + ",setparams=range=tv:color_primaries=bt709:color_trc=bt709:colorspace=bt709",
            "-c:v", "libx264", "-profile:v", "high", "-preset", "medium", "-crf", "18",
            "-pix_fmt", "yuv420p", "-color_primaries", "bt709", "-color_trc", "bt709",
            "-colorspace", "bt709", "-color_range", "tv", "-movflags", "+faststart", str(dst)]
    dst.parent.mkdir(parents=True, exist_ok=True)
    proc = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
    if proc.returncode != 0:
        raise RuntimeError(proc.stderr[-800:])


def main() -> int:
    ap = argparse.ArgumentParser(description="Presets de cor: listar, aplicar num vídeo ou comparar num quadro.")
    ap.add_argument("input", type=Path, nargs="?", help="vídeo de entrada")
    ap.add_argument("-o", "--output", type=Path, help="vídeo de saída (ou PNG com --candidates)")
    ap.add_argument("--list", action="store_true", help="lista os presets")
    ap.add_argument("--preset", help="nome do preset (veja --list)")
    ap.add_argument("--filter", help="filtro ffmpeg cru, ex. eq=contrast=1.1")
    ap.add_argument("--candidates", help='comparação lado a lado: "a=<filtro>;b=<preset>;original="')
    ap.add_argument("--frame", default="0", help="tempo do quadro para --candidates (segundos, padrão 0)")
    ap.add_argument("--tile-width", type=int, default=0, help="largura de cada quadro na comparação")
    args = ap.parse_args()

    if args.list:
        for name, (flt, desc) in PRESETS.items():
            print(f"{name:18s} {desc}")
            if flt:
                print(f"{'':18s} {flt}")
        return 0

    if not args.input:
        ap.print_usage(sys.stderr)
        print("erro: informe o vídeo de entrada (ou use --list)", file=sys.stderr)
        return 2
    from contact_sheet import abspath, default_verify_dir

    src = abspath(args.input)
    if not src.exists():
        print(f"erro: vídeo não encontrado: {src}", file=sys.stderr)
        return 2

    if args.candidates is not None:
        from contact_sheet import parse_time
        try:
            t = parse_time(args.frame)
            cands = parse_candidates(args.candidates)
            out = abspath(args.output) if args.output else default_verify_dir(src) / f"cores_{src.stem}.png"
            candidates_montage(src, cands, t, out, args.tile_width)
        except (ValueError, RuntimeError, argparse.ArgumentTypeError) as exc:
            print(f"erro: {exc}", file=sys.stderr)
            return 2
        print(f"comparação salva: {out} ({len(cands)} candidatos)")
        return 0

    if not args.output:
        print("erro: informe a saída com -o", file=sys.stderr)
        return 2
    if args.preset and args.filter:
        print("erro: use --preset OU --filter, não os dois", file=sys.stderr)
        return 2
    try:
        flt = get_preset(args.preset) if args.preset else resolve_grade(args.filter or "")
    except (KeyError, ValueError) as exc:
        print(f"erro: {exc.args[0] if exc.args else exc}", file=sys.stderr)
        return 2
    err = check_filter(flt)
    if err:
        print(f"erro: filtro inválido: {err}", file=sys.stderr)
        return 2
    print(f"aplicando cor em {src.name}: {flt or 'sem filtro (só recodifica)'}")
    try:
        apply_grade(src, abspath(args.output), flt)
    except RuntimeError as exc:
        print(f"erro: o ffmpeg falhou: {exc}", file=sys.stderr)
        return 1
    print(f"pronto: {abspath(args.output)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())