#!/usr/bin/env python3
"""Before / after comparison video (and still) from two images or two videos: an animated wipe slider, a static
split, side by side, stacked (great for 9:16), or a toggle that cuts between them. Labels, divider handle and a
still compare PNG for posts. Use after an upscale, restyle, colour grade, retouch, background swap or remaster.

Examples
  python beforeafter.py --before old.jpg --after restored.jpg --mode wipe --seconds 6
  python beforeafter.py --before raw.mp4 --after graded.mp4 --mode split --aspect 16:9
  python beforeafter.py --before a.png --after b.png --mode stack --aspect 9:16 --labels "ORIGINAL,KENERATE 4K"
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

MODES = ("wipe", "split", "side", "stack", "toggle")


def label_png(W: int, H: int, labels, mode: str, font, dst: Path, accent) -> Path:
    from PIL import Image
    img = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    s = min(W, H)
    m = int(s * 0.035)
    for i, txt in enumerate(labels):
        if not txt:
            continue
        spr = K.text_img(txt, int(s * 0.04), (255, 255, 255, 255), "bold", font)
        pad = int(s * 0.014)
        box = Image.new("RGBA", (spr.width + pad * 2, spr.height + pad * 2), (0, 0, 0, 0))
        from PIL import ImageDraw
        ImageDraw.Draw(box).rounded_rectangle((0, 0, box.width - 1, box.height - 1), radius=pad,
                                              fill=(tuple(accent) + (230,)) if i == 1 else (0, 0, 0, 160))
        box.alpha_composite(spr, (pad, pad))
        if mode == "stack":
            pos = (m, m) if i == 0 else (m, H // 2 + m)
        elif mode == "side":
            pos = (m, m) if i == 0 else (W // 2 + m, m)
        else:
            pos = (m, m) if i == 0 else (W - m - box.width, m)
        img.alpha_composite(box, pos)
    img.save(dst)
    return dst


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--before", required=True)
    ap.add_argument("--after", required=True)
    ap.add_argument("--mode", choices=MODES, default="wipe")
    ap.add_argument("--labels", default="BEFORE,AFTER", help="'left,right' ('' for none)")
    ap.add_argument("--aspect", default="", help="output aspect (default: the before input's; stack -> 9:16)")
    ap.add_argument("--res", type=int, default=1080)
    ap.add_argument("--fps", type=int, default=30)
    ap.add_argument("--seconds", type=float, default=6.0, help="length for stills (videos use the shorter clip)")
    ap.add_argument("--sweeps", type=float, default=1.5, help="wipe: back-and-forth passes over the length")
    ap.add_argument("--accent", default="#FFD60A")
    ap.add_argument("--audio", choices=["none", "before", "after"], default="none", help="whose sound to keep (videos)")
    K.add_font_args(ap, "TTF/OTF for the labels")
    ap.add_argument("--out", default="kenerate-out/before-after/compare.mp4")
    ap.add_argument("--keep-temp", action="store_true")
    a = K.parse_args(ap)

    st = K.Status(4, "before-after")
    K.require(status=st)
    tmp = K.workdir(a.keep_temp)
    out = Path(a.out)
    out.parent.mkdir(parents=True, exist_ok=True)

    st.step("Fetching both sides")
    pb, pa = K.fetch(a.before, tmp, "before", st), K.fetch(a.after, tmp, "after", st)
    ib, ia = K.probe(pb), K.probe(pa)
    still_b = K.kind_of(str(pb)) == "image" or ib["duration"] <= 0.05
    still_a = K.kind_of(str(pa)) == "image" or ia["duration"] <= 0.05
    durs = [d for d, s in ((ib["duration"], still_b), (ia["duration"], still_a)) if not s]
    D = min(durs) if durs else a.seconds
    aspect = a.aspect or ("9:16" if a.mode == "stack" else f"{ib['width']}x{ib['height']}")
    W, H = K.size_for(aspect)
    sc = a.res / min(W, H)
    W, H = int(W * sc) // 2 * 2, int(H * sc) // 2 * 2
    if ia["width"] and ib["width"] and abs(ia["width"] / ia["height"] - ib["width"] / ib["height"]) > 0.02:
        st.warn("the two sides have different shapes: both are fitted to the same frame (crop) - check the preview")
    if not still_a and not still_b and abs(ia["duration"] - ib["duration"]) > 0.3:
        st.warn(f"clip lengths differ ({ib['duration']:.1f}s vs {ia['duration']:.1f}s): using the shorter")
    labels = [x.strip() for x in a.labels.split(",")] + ["", ""] if a.labels else ["", ""]

    st.step(f"Preparing {W}x{H} ({a.mode})")
    if a.mode in ("side", "stack"):
        cw, ch = (W // 2 // 2 * 2, H) if a.mode == "side" else (W, H // 2 // 2 * 2)
    else:
        cw, ch = W, H
    nb, na = tmp / "b.mp4", tmp / "a.mp4"
    for src, dst, still in ((pb, nb, still_b), (pa, na, still_a)):
        if still:  # no Ken Burns: both sides must line up pixel for pixel
            K.run_ffmpeg(["-loop", "1", "-framerate", str(a.fps), "-t", f"{D:.3f}", "-i", str(src), "-f", "lavfi", "-t", f"{D:.3f}",
                          "-i", "anullsrc=channel_layout=stereo:sample_rate=48000", "-vf", K.fit_filter(cw, ch, "cover") + f",fps={a.fps},format=yuv420p",
                          "-map", "0:v", "-map", "1:a", "-t", f"{D:.3f}"] + K.video_codec_args(16, "veryfast") + K.AUDIO_ARGS + [str(dst)],
                         duration=D, status=st)
        else:
            K.normalize_clip(src, dst, cw, ch, a.fps, 0, D, "cover", status=st)
    lab = label_png(W, H, labels[:2], a.mode, a.font, tmp / "labels.png", K.hex_rgb(a.accent))

    st.step("Rendering")
    acc = K.rgb_hex(K.hex_rgb(a.accent)).lstrip("#")
    lw = max(2, W // 360)
    if a.mode == "wipe":  # divider painted inside the blend expression so it always sits exactly on the edge
        r, g, b = K.hex_rgb(a.accent)
        yy = 16 + (65.481 * r + 128.553 * g + 24.966 * b) / 255
        uu = 128 + (-37.797 * r - 74.203 * g + 112.0 * b) / 255
        vv = 128 + (112.0 * r - 93.786 * g - 18.214 * b) / 255
        pexpr = f"(0.5-0.42*cos(2*PI*{a.sweeps / 2:.4f}*T/{D:.3f}))"
        half = lw / W

        def e(col):
            return f"if(lt(abs(X/W-{pexpr}),{half:.5f}),{col:.1f},if(lt(X/W,{pexpr}),A,B))"
        graph = (f"[0:v]format=yuv444p[b4];[1:v]format=yuv444p[a4];"
                 f"[b4][a4]blend=c0_expr='{e(yy)}':c1_expr='{e(uu)}':c2_expr='{e(vv)}'[bl];[bl][2:v]overlay=0:0[v]")
    elif a.mode == "split":
        graph = f"[0:v][1:v]blend=all_expr='if(lt(X,W/2),A,B)'[bl];[bl]drawbox=x=iw/2-{lw}:y=0:w={lw * 2}:h=ih:color=0x{acc}@1:t=fill[dv];[dv][2:v]overlay=0:0[v]"
    elif a.mode == "side":
        graph = f"[0:v][1:v]hstack=inputs=2,pad={W}:{H}:(ow-iw)/2:0[hs];[hs]drawbox=x=iw/2-{lw}:y=0:w={lw * 2}:h=ih:color=0x{acc}@1:t=fill[dv];[dv][2:v]overlay=0:0[v]"
    elif a.mode == "stack":
        graph = f"[0:v][1:v]vstack=inputs=2,pad={W}:{H}:0:(oh-ih)/2[vs];[vs]drawbox=x=0:y=ih/2-{lw}:w=iw:h={lw * 2}:color=0x{acc}@1:t=fill[dv];[dv][2:v]overlay=0:0[v]"
    else:  # toggle: 1 s each, labels follow
        graph = ("[0:v][1:v]blend=all_expr='if(lt(mod(T,2),1),A,B)'[bl];"
                 f"[2:v]split[l1][l2];[l1]crop=iw/2:ih:0:0[lb];[l2]crop=iw/2:ih:iw/2:0[la];"
                 f"[bl][lb]overlay=0:0:enable='lt(mod(t,2),1)'[o1];[o1][la]overlay=W/2:0:enable='gte(mod(t,2),1)'[v]")
    amap = {"none": [], "before": ["-map", "0:a"], "after": ["-map", "1:a"]}[a.audio]
    K.run_ffmpeg(["-i", str(nb), "-i", str(na), "-loop", "1", "-t", f"{D:.3f}", "-i", str(lab), "-filter_complex",
                  graph + ";[v]format=yuv420p[vo]", "-map", "[vo]"] + amap + ["-t", f"{D:.3f}"] + K.video_codec_args(18)
                 + (K.AUDIO_ARGS if amap else ["-an"]) + [str(out)], duration=D, status=st)

    st.step("Still + preview")
    from PIL import Image
    fb = K.open_image(K.extract_frame(nb, min(D / 2, D - 0.1), tmp / "fb.jpg")).convert("RGB")
    fa = K.open_image(K.extract_frame(na, min(D / 2, D - 0.1), tmp / "fa.jpg")).convert("RGB")
    if a.mode in ("side", "stack"):
        still = Image.new("RGB", (W, H), (0, 0, 0))
        still.paste(fb, (0, 0))
        still.paste(fa, (cw, 0) if a.mode == "side" else (0, ch))
    else:
        still = fb.copy()
        still.paste(fa.crop((W // 2, 0, W, H)), (W // 2, 0))
        from PIL import ImageDraw
        ImageDraw.Draw(still).rectangle((W // 2 - lw, 0, W // 2 + lw, H), fill=K.hex_rgb(a.accent))
    still = still.convert("RGBA")
    still.alpha_composite(K.open_image(lab).convert("RGBA") if a.mode != "toggle" else Image.new("RGBA", (W, H), (0, 0, 0, 0)))
    still_path = out.with_name(out.stem + "_still.png")
    still.convert("RGB").save(still_path)
    times = [D * f for f in (0.1, 0.3, 0.5, 0.7, 0.9)]
    sheet = K.video_preview(out, out.with_name(out.stem + "_preview.jpg"), tmp=tmp / "pv", times=times)
    K.finish({"output": K.file_info(out), "still": str(still_path.resolve()), "preview": str(sheet.resolve()), "mode": a.mode,
              "seconds": round(D, 2), "elapsed_s": st.elapsed()})


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