"""
encode.py — turn a directory of PNG frames into a Discord-friendly GIF or MP4.

Use it INSIDE a single code_exec call. Render your frames to a scratch dir as
frame_00000.png, frame_00001.png, … (fixed 5-digit width, ONE continuous
counter across all scenes), then:

    from encode import encode_animation
    art = encode_animation("/tmp/frames", fps=15, seconds_hint=8)
    # art == {"path": "/workspace/final.mp4", "mime": "video/mp4", "bytes": 1234567}

It writes the finished artifact to the TOP of /workspace (so code_exec returns a
file_id for it) and picks the format for you:

  * short & light            -> animated GIF (fun, loops, autoplays)
  * long (many seconds+),
    OR a GIF too big to post  -> H.264 MP4: yuv420p, even dimensions, +faststart
                                 — the combination Discord needs to PLAY it inline
                                 (a video whose dimensions Discord can't read is
                                 shown as a dead download, not a player).

Requires ffmpeg (always present in the sandbox). Uses gifski + gifsicle for a
cleaner GIF when they're installed, falling back to ffmpeg's palette pipeline.
This helper only ENCODES — you draw the frames.
"""

import glob
import os
import shutil
import subprocess

# A GIF longer than this balloons past what's worth shipping as a GIF — encode
# straight to MP4 instead. Tune down for chattier channels.
GIF_MAX_SECONDS = 20
# Keep a GIF comfortably under Discord's ~10 MiB inline limit; over this we
# re-encode to an MP4 and drop the GIF.
GIF_MAX_BYTES = 9 * 1024 * 1024


def _which(prog):
    return shutil.which(prog) is not None


def _run(cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        raise RuntimeError(f"{cmd[0]} failed:\n{p.stderr[-2000:]}")


def _frames(frames_dir):
    fs = sorted(glob.glob(os.path.join(frames_dir, "frame_*.png")))
    if not fs:
        raise FileNotFoundError(
            f"no frame_*.png in {frames_dir} — render frames as "
            f"{frames_dir}/frame_00000.png, frame_00001.png, … first"
        )
    return fs


def _encode_gif(frames_dir, fps, out, long_edge):
    fs = _frames(frames_dir)
    if _which("gifski"):
        # gifski takes an explicit, sorted file list — no numbering-gap surprises.
        _run(["gifski", "-o", out, "--fps", str(fps), "--quality", "90"] + fs)
    else:
        vf = (f"fps={fps},scale='min({long_edge},iw)':-1:flags=lanczos,"
              f"split[s0][s1];[s0]palettegen=stats_mode=diff[p];"
              f"[s1][p]paletteuse=dither=bayer:bayer_scale=3")
        _run(["ffmpeg", "-y", "-pattern_type", "glob",
              "-i", os.path.join(frames_dir, "frame_*.png"), "-vf", vf, out])
    if _which("gifsicle"):
        _run(["gifsicle", "-O3", "--lossy=60", out, "-o", out])


def _encode_mp4(frames_dir, fps, out, long_edge):
    _frames(frames_dir)  # validate presence
    # -2 forces EVEN width/height (H.264 requires it); yuv420p + faststart make
    # Discord/browsers play it inline instead of offering a download.
    vf = f"scale='min({long_edge},iw)':-2:flags=lanczos"
    _run(["ffmpeg", "-y", "-framerate", str(fps), "-pattern_type", "glob",
          "-i", os.path.join(frames_dir, "frame_*.png"),
          "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "23",
          "-vf", vf, "-movflags", "+faststart", out])


def encode_animation(frames_dir, fps=15, seconds_hint=0, long_edge=640):
    """Encode frames_dir → /workspace/final.{gif,mp4}. Returns {path, mime, bytes}.

    seconds_hint: your intended runtime in seconds (falls back to frame_count/fps
    when 0). Pass it for long pieces so the format decision doesn't depend on how
    many frames you happened to render.
    long_edge: cap for the output's long side in px (keeps size/encode sane).
    """
    n = len(_frames(frames_dir))
    seconds = seconds_hint or (n / max(fps, 1))
    gif, mp4 = "/workspace/final.gif", "/workspace/final.mp4"

    if seconds > GIF_MAX_SECONDS:
        _encode_mp4(frames_dir, fps, mp4, long_edge)
        return _art(mp4, "video/mp4")

    _encode_gif(frames_dir, fps, gif, long_edge)
    if os.path.getsize(gif) <= GIF_MAX_BYTES:
        return _art(gif, "image/gif")

    # Too heavy as a GIF — hand Discord an MP4 it can actually render.
    _encode_mp4(frames_dir, fps, mp4, long_edge)
    os.remove(gif)
    return _art(mp4, "video/mp4")


def _art(path, mime):
    b = os.path.getsize(path)
    print(f"[encode] {path}  {mime}  {b / 1024 / 1024:.2f} MiB")
    return {"path": path, "mime": mime, "bytes": b}