#!/usr/bin/env python3
"""Colour grade videos and images with a .cube LUT or a built-in look, plus exposure / contrast / saturation /
warmth / tint. The whole grade is baked into one 33-point .cube, so stills (Pillow) and video (ffmpeg lut3d) match,
and the .cube can be handed to Premiere / Resolve / CapCut.

Examples
  python grade.py --input clip.mp4 --looks-sheet                       # all looks on one frame: pick one
  python grade.py --input clip.mp4 --look teal-orange --intensity 0.7 --contrast 1.05
  python grade.py --input a.jpg b.jpg --lut brand.cube --intensity 0.8 --warmth 10
  python grade.py --look warm-film --export-cube kenerate-out/color-grade/warm-film.cube   # just the LUT
"""
from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import kmedia as K  # noqa: E402


def compare_still(before, after, dst: Path, label=True) -> Path:
    """Left half before, right half after, with a divider."""
    from PIL import Image, ImageDraw
    W, H = before.size
    out = before.convert("RGB").copy()
    out.paste(after.convert("RGB").crop((W // 2, 0, W, H)), (W // 2, 0))
    d = ImageDraw.Draw(out)
    d.rectangle((W // 2 - max(1, W // 400), 0, W // 2 + max(1, W // 400), H), fill=(255, 255, 255))
    if label:
        f = K.load_font(max(14, H // 22), "bold")
        for txt, x in (("BEFORE", int(W * 0.03)), ("AFTER", W // 2 + int(W * 0.03))):
            d.text((x, int(H * 0.03)), txt, font=f, fill=(255, 255, 255), stroke_width=max(1, H // 300), stroke_fill=(0, 0, 0))
    out.save(dst, quality=88)
    return dst


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--input", nargs="*", default=[], help="videos and / or images (URLs or files)")
    ap.add_argument("--look", choices=list(K.LOOKS), help="built-in look")
    ap.add_argument("--lut", help=".cube 3D LUT file (instead of --look)")
    ap.add_argument("--intensity", type=float, default=1.0, help="0..1 strength of the look / LUT (default 1)")
    ap.add_argument("--exposure", type=float, default=0.0, help="stops, e.g. +0.3 or -0.5")
    ap.add_argument("--contrast", type=float, default=1.0)
    ap.add_argument("--saturation", type=float, default=1.0)
    ap.add_argument("--warmth", type=float, default=0.0, help="-100 (cool) .. 100 (warm)")
    ap.add_argument("--tint", type=float, default=0.0, help="-100 (green) .. 100 (magenta)")
    ap.add_argument("--looks-sheet", action="store_true", help="render every built-in look on a frame of the first input and stop")
    ap.add_argument("--export-cube", help="also write the combined grade as this .cube file")
    ap.add_argument("--at", type=float, default=-1, help="video: time of the preview frame (default: middle)")
    ap.add_argument("--out-dir", default="kenerate-out/color-grade")
    ap.add_argument("--suffix", default="_graded")
    a = ap.parse_args()

    st = K.Status(5, "color-grade")
    K.require(status=st)
    out_dir = Path(a.out_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    tmp = K.workdir()
    if a.lut and not Path(a.lut).exists():
        K.fail(f"LUT not found: {a.lut}")
    look = a.lut or a.look
    neutral = not look and a.exposure == 0 and a.contrast == 1 and a.saturation == 1 and a.warmth == 0 and a.tint == 0
    if neutral and not a.looks_sheet:
        K.fail("Nothing to do: pass --look, --lut or an adjustment (--exposure, --contrast, --saturation, --warmth, --tint).",
               "Run with --looks-sheet first to compare the built-in looks.")

    st.step("Building the grade (.cube)")
    cube = tmp / "grade.cube"
    if not a.looks_sheet:
        try:
            K.write_cube(K.grade_fn(look, a.intensity, a.exposure, a.contrast, a.saturation, a.warmth, a.tint), cube,
                         title=f"Kenerate {Path(look).stem if look else 'grade'}")
        except ValueError as e:
            K.fail(str(e), "Use a 3D .cube LUT (LUT_3D_SIZE) or one of the built-in looks.")
        st.info(f"{look or 'adjustments only'} at {a.intensity:.2f}; exposure {a.exposure:+.2f}, contrast {a.contrast}, "
                f"saturation {a.saturation}, warmth {a.warmth:+.0f}, tint {a.tint:+.0f}")
        if a.export_cube:
            Path(a.export_cube).parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(cube, a.export_cube)
            st.info(f"LUT written: {Path(a.export_cube).resolve()}")
    if not a.input:
        if a.export_cube:
            K.finish({"cube": str(Path(a.export_cube).resolve())})
            return
        K.fail("No --input given.", "Pass videos / images with --input, or only --export-cube to write the LUT.")

    st.step(f"Fetching {len(a.input)} input(s)")
    srcs = K.fetch_all(a.input, tmp / "in", st, prefix="in")

    def ref_frame(p: Path):
        if K.kind_of(str(p)) == "image":
            return K.open_image(p).convert("RGB")
        info = K.probe(p)
        t = a.at if a.at >= 0 else info["duration"] / 2
        return K.open_image(K.extract_frame(p, t, tmp / f"ref_{p.stem}.jpg")).convert("RGB")

    if a.looks_sheet:
        st.step(f"Rendering {len(K.LOOKS)} looks on one frame")
        fr = ref_frame(srcs[0])
        fr.thumbnail((640, 640))
        cells, labels = [tmp / "look_original.jpg"], ["original"]
        fr.save(cells[0], quality=90)
        for name in K.LOOKS:
            c = K.write_cube(K.grade_fn(name, a.intensity), tmp / f"{name}.cube", size=17)
            p = tmp / f"look_{name}.jpg"
            K.apply_look_image(fr, str(c)).save(p, quality=90)
            cells.append(p)
            labels.append(name if a.intensity >= 0.999 else f"{name} {a.intensity:.1f}")
        sheet = K.contact_sheet(cells, out_dir / "looks_sheet.jpg", cols=4, cell=320, labels=labels)
        st.info("show the sheet to the user and ask which look (and how strong) before grading")
        K.finish({"looks_sheet": str(sheet), "looks": list(K.LOOKS)})
        return

    st.step("Grading")
    results, previews = [], []
    for i, p in enumerate(srcs, 1):
        stem = K.safe_name(p.stem)
        if K.kind_of(str(p)) == "image":
            im = K.open_image(p)
            g = K.apply_look_image(im, str(cube))
            ext = ".png" if im.mode == "RGBA" or p.suffix.lower() == ".png" else ".jpg"
            dst = out_dir / f"{stem}{a.suffix}{ext}"
            g.save(dst, **({"quality": 93} if ext == ".jpg" else {}))
            before, after = im.convert("RGB"), g.convert("RGB")
        else:
            info = K.probe(p)
            dst = out_dir / f"{stem}{a.suffix}.mp4"
            st.info(f"video {i}/{len(srcs)}: {p.name} ({info['duration']:.1f} s)")
            K.run_ffmpeg(["-i", str(p.resolve()), "-vf", "lut3d=file=grade.cube:interp=tetrahedral,format=yuv420p",
                          "-map", "0:v:0", "-map", "0:a?"] + K.video_codec_args(18) + ["-c:a", "copy", str(dst)],
                         duration=info["duration"], status=st, cwd=tmp)
            t = a.at if a.at >= 0 else info["duration"] / 2
            before = ref_frame(p)
            after = K.open_image(K.extract_frame(dst, t, tmp / f"after_{stem}.jpg")).convert("RGB")
        cmp_path = compare_still(before, after.resize(before.size), out_dir / f"{stem}{a.suffix}_compare.jpg")
        previews.append(cmp_path)
        results.append({"input": str(p), "output": K.file_info(dst), "compare": str(cmp_path)})
        st.progress(i * 100 / len(srcs), f"{i}/{len(srcs)}")
    st.step("Preview sheet")
    sheet = K.contact_sheet(previews, out_dir / "grade_preview.jpg", cols=min(3, len(previews)), cell=420) if previews else None
    K.finish({"outputs": results, "preview": str(sheet) if sheet else None,
              "cube": str(Path(a.export_cube).resolve()) if a.export_cube else None, "elapsed_s": st.elapsed()})


if __name__ == "__main__":
    K.main_guard(main)