from __future__ import annotations

import runpy
from pathlib import Path
from typing import Iterable

import imageio.v2 as imageio
import numpy as np
from PIL import Image, ImageDraw


SCRIPT_DIR = Path(__file__).resolve().parent
gif_ns = runpy.run_path(str(SCRIPT_DIR / "50_make_blog_demo_gif.py"))

ASSET_DIR = gif_ns["ASSET_DIR"]
slides: list[Image.Image] = gif_ns["slides"]
video_effects: list[dict] = gif_ns["video_effects"]

MP4 = ASSET_DIR / "hbm-hbn-comsol-demo.mp4"
POSTER = ASSET_DIR / "hbm-hbn-comsol-demo-poster.png"

W, H = 1280, 720
FPS = 24
FADE_SECONDS = 0.35
CYAN = (125, 249, 255)
GREEN = (108, 255, 181)


def ease(t: float) -> float:
    return t * t * (3 - 2 * t)


def as_rgb_array(img: Image.Image) -> np.ndarray:
    return np.asarray(img.convert("RGB"))


def ken_burns(slide: Image.Image, effect: dict, t: float) -> Image.Image:
    e = ease(t)
    start_zoom = effect.get("start_zoom", 1.0)
    end_zoom = effect.get("end_zoom", start_zoom)
    zoom = start_zoom + (end_zoom - start_zoom) * e
    crop_w = max(1, round(W / zoom))
    crop_h = max(1, round(H / zoom))

    start_focus = effect.get("start_focus", (0.5, 0.5))
    end_focus = effect.get("end_focus", start_focus)
    fx = start_focus[0] + (end_focus[0] - start_focus[0]) * e
    fy = start_focus[1] + (end_focus[1] - start_focus[1]) * e

    max_x = max(0, slide.width - crop_w)
    max_y = max(0, slide.height - crop_h)
    x = round(max_x * fx)
    y = round(max_y * fy)
    frame = slide.crop((x, y, x + crop_w, y + crop_h)).resize((W, H), Image.Resampling.LANCZOS)
    return frame


def draw_overlay(frame: Image.Image, effect: dict, t: float) -> None:
    overlay_kind = effect.get("overlay", "")
    layer = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    d = ImageDraw.Draw(layer)

    # A quiet timeline makes the clip feel recorded without becoming a UI tutorial.
    line_w = 220
    line_x = W - line_w - 54
    line_y = H - 35
    d.rounded_rectangle((line_x, line_y, line_x + line_w, line_y + 4), radius=2, fill=(255, 255, 255, 36))
    d.rounded_rectangle((line_x, line_y, line_x + round(line_w * t), line_y + 4), radius=2, fill=(*CYAN, 180))

    if overlay_kind == "terminal":
        cursor_x = 82 + round(300 * t)
        cursor_y = 548
        if int(t * 12) % 2 == 0:
            d.rectangle((cursor_x, cursor_y, cursor_x + 8, cursor_y + 22), fill=(*GREEN, 210))
        d.rounded_rectangle((54, 154, 484, 584), radius=12, outline=(*GREEN, 90), width=3)

    if overlay_kind in {"spotlight", "results"}:
        pulse = 90 + round(75 * (1 - abs(0.5 - t) * 2))
        color = GREEN if overlay_kind == "results" else CYAN
        for rect in effect.get("rects", []):
            d.rounded_rectangle(rect, radius=10, outline=(*color, pulse), width=4)

    if overlay_kind == "paper":
        d.rounded_rectangle((54, 146, 1226, 606), radius=14, outline=(*CYAN, 42), width=2)

    frame.alpha_composite(layer)


def render_frame(slide: Image.Image, effect: dict, t: float) -> Image.Image:
    frame = ken_burns(slide, effect, t).convert("RGBA")
    draw_overlay(frame, effect, t)
    return frame.convert("RGB")


def frame_sequence() -> Iterable[np.ndarray]:
    fade_count = max(1, round(FPS * FADE_SECONDS))
    for idx, (slide, effect) in enumerate(zip(slides, video_effects)):
        hold_count = max(2, round(FPS * effect["duration"]))
        rendered: list[Image.Image] = []
        for step in range(hold_count):
            t = step / (hold_count - 1)
            frame = render_frame(slide, effect, t)
            rendered.append(frame)
            yield as_rgb_array(frame)

        if idx == len(slides) - 1:
            continue

        next_frame = render_frame(slides[idx + 1], video_effects[idx + 1], 0.0)
        last_frame = rendered[-1]
        for step in range(1, fade_count + 1):
            alpha = step / (fade_count + 1)
            yield as_rgb_array(Image.blend(last_frame, next_frame, alpha))


def write_video() -> None:
    render_frame(slides[0], video_effects[0], 0.0).save(POSTER, optimize=True, compress_level=6)
    frame_count = 0
    with imageio.get_writer(
        MP4,
        fps=FPS,
        codec="libx264",
        quality=8,
        ffmpeg_params=["-pix_fmt", "yuv420p", "-movflags", "+faststart"],
        macro_block_size=16,
    ) as writer:
        for frame in frame_sequence():
            writer.append_data(frame)
            frame_count += 1
    duration = frame_count / FPS
    print(f"wrote {MP4} ({MP4.stat().st_size / 1024 / 1024:.2f} MiB, {duration:.1f}s)")
    print(f"wrote {POSTER} ({POSTER.stat().st_size / 1024:.1f} KiB)")


if __name__ == "__main__":
    write_video()