#!/usr/bin/env python3
"""GIF maker: a clip (or a set of images) -> an optimized GIF with a 2-pass palette, optional caption, boomerang,
and an automatic size cap (shrinks width / fps until it fits). Also WebP or a silent looping MP4.

Examples
  python gif.py --input https://.../clip.mp4 --start 1.2 --end 4 --width 480 --max-mb 8 --text "when the coffee hits"
  python gif.py --input clip.mp4 --boomerang --format webp
  python gif.py --images a.png b.png c.png --frame-seconds 0.5 --width 600
"""
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


def caption_png(text: str, W: int, H: int, dst: Path, pos: str, font) -> Path:
    from PIL import Image
    img = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    bh = int(H * 0.24)
    y0 = int(H * 0.04) if pos == "top" else H - bh - int(H * 0.04)
    K.draw_block(img, text, (int(W * 0.05), y0, int(W * 0.95), y0 + bh), fill=(255, 255, 255), max_size=int(min(W, H) * 0.14),
                 min_size=12, kind="impact", font_override=font, stroke_ratio=0.08, upper=True, max_lines=2)
    img.save(dst)
    return dst


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--input", help="video URL or file")
    ap.add_argument("--images", nargs="*", default=[], help="or images, in order")
    ap.add_argument("--frame-seconds", type=float, default=0.4, help="images: seconds per image")
    ap.add_argument("--start", type=float, default=0.0)
    ap.add_argument("--end", type=float, default=0.0, help="0 = start + 6 s (or the clip's end)")
    ap.add_argument("--width", type=int, default=480)
    ap.add_argument("--fps", type=int, default=15)
    ap.add_argument("--max-mb", type=float, default=8.0, help="size cap (Slack / Discord / email friendly)")
    ap.add_argument("--text", default="")
    ap.add_argument("--text-pos", choices=["top", "bottom"], default="bottom")
    ap.add_argument("--boomerang", action="store_true", help="play forward then backward")
    ap.add_argument("--format", choices=["gif", "webp", "mp4"], default="gif")
    K.add_font_args(ap)
    ap.add_argument("--out", default="kenerate-out/gif-maker/clip.gif")
    ap.add_argument("--keep-temp", action="store_true")
    a = K.parse_args(ap)
    if not a.input and not a.images:
        K.fail("Give --input VIDEO or --images ...")

    st = K.Status(4, "gif-maker")
    K.require(status=st)
    tmp = K.workdir(a.keep_temp)
    out = Path(a.out).with_suffix("." + a.format)
    out.parent.mkdir(parents=True, exist_ok=True)

    st.step("Fetching input")
    if a.images:
        imgs = K.fetch_all(a.images, tmp, st, prefix="img")
        first = K.open_image(imgs[0])
        Wsrc, Hsrc = first.size
        lst = K.write_concat([(p, a.frame_seconds) for p in imgs], tmp / "imgs.ffconcat")
        src_args = ["-f", "concat", "-safe", "0", "-i", str(lst)]
        D = a.frame_seconds * len(imgs)
        pre = f"scale={Wsrc}:{Hsrc}:force_original_aspect_ratio=decrease,pad={Wsrc}:{Hsrc}:(ow-iw)/2:(oh-ih)/2,"
    else:
        src = K.fetch(a.input, tmp, "input", st)
        info = K.probe(src)
        end = a.end or min(info["duration"], a.start + 6)
        end = min(end, info["duration"]) if info["duration"] else end
        D = max(0.2, end - a.start)
        Wsrc, Hsrc = info["width"], info["height"]
        src_args = ["-ss", f"{a.start:.3f}", "-t", f"{D:.3f}", "-i", str(src)]
        pre = ""
    st.info(f"{D:.2f} s of {Wsrc}x{Hsrc}")

    st.step("Caption")
    width = min(a.width, Wsrc) // 2 * 2
    fps = a.fps
    cap = None
    if a.text:
        h = int(round(Hsrc * width / Wsrc / 2)) * 2
        cap = caption_png(a.text, width, h, tmp / "cap.png", a.text_pos, a.font)

    st.step(f"Encoding {a.format} (cap {a.max_mb} MB)")
    tries = 0
    while True:
        tries += 1
        h = int(round(Hsrc * width / Wsrc / 2)) * 2
        base = f"[0:v]{pre}fps={fps},scale={width}:{h}:flags=lanczos,setsar=1"
        if a.boomerang:
            base += ",split[f][r];[r]reverse[rv];[f][rv]concat=n=2:v=1:a=0"
        base += "[b]"
        inputs = list(src_args)
        if cap:
            if cap and h != K.open_image(cap).height:
                cap = caption_png(a.text, width, h, tmp / f"cap{tries}.png", a.text_pos, a.font)
            inputs += ["-i", str(cap)]
            base += ";[b][1:v]overlay=0:0[c]"
            last = "[c]"
        else:
            last = "[b]"
        DD = D * (2 if a.boomerang else 1)
        if a.format == "gif":
            graph = base + f";{last}split[x][y];[x]palettegen=max_colors=192:stats_mode=diff[p];[y][p]paletteuse=dither=sierra2_4a:diff_mode=rectangle[o]"
            K.run_ffmpeg(inputs + ["-filter_complex", graph, "-map", "[o]", "-loop", "0", str(out)], duration=DD, status=st)
        elif a.format == "webp":
            K.run_ffmpeg(inputs + ["-filter_complex", base + f";{last}null[o]", "-map", "[o]", "-c:v", "libwebp", "-quality", "75",
                                   "-loop", "0", "-an", str(out)], duration=DD, status=st)
        else:
            K.run_ffmpeg(inputs + ["-filter_complex", base + f";{last}format=yuv420p[o]", "-map", "[o]", "-an"]
                         + K.video_codec_args(23) + [str(out)], duration=DD, status=st)
        mb = out.stat().st_size / 1e6
        if mb <= a.max_mb or tries >= 6 or a.format == "mp4":
            break
        st.info(f"{mb:.1f} MB > {a.max_mb} MB: retrying smaller")
        if fps > 10:
            fps = max(10, int(fps * 0.75))
        width = max(160, int(width * 0.8) // 2 * 2)

    st.step("Done")
    mb = out.stat().st_size / 1e6
    if mb > a.max_mb:
        st.warn(f"still {mb:.1f} MB: shorten it (--end) or use --format mp4")
    K.finish({"output": str(out.resolve()), "mb": round(mb, 2), "width": width, "fps": fps, "seconds": round(D, 2),
              "boomerang": a.boomerang, "attempts": tries, "elapsed_s": st.elapsed()})


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