#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""把「Python 爱心弹窗」效果渲染成一条 1080p 的片子（纯离屏绘制，不碰桌面）。

    python tools/render_reel.py                 # 渲染 out/爱心弹窗-复刻.mp4
    python tools/render_reel.py --preview 7,9   # 只导出这几个时间点的预览图
"""
from __future__ import annotations

import argparse
import io
import math
import os
import random
import re
import subprocess
import tokenize
from collections import OrderedDict

import numpy as np
from PIL import Image, ImageColor, ImageDraw, ImageFilter, ImageFont

W, H = 1920, 1080
FPS = 30
DUR = 22.0

FFMPEG = os.environ.get("FFMPEG") or "ffmpeg"
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)

# ---------------- 时间轴 ----------------
T_BLACK_END = 0.30
T_PLATE_END = 1.30
T_WIPE_END = 1.75
T_HEART_START = 4.80
HEART_N = 72
HEART_STEP = 0.06
HEART_S = 33

HEART_DY = 120
T_DESTROY = 10.30
T_FLOOD_START = 10.35
FLOOD_N = 150
FLOOD_STEP = 0.04
T_TITLE_IN = 21.10

MESSAGES = ["老吴多喝水", "老吴好好爱自己", "老吴好好吃饭", "老吴保持好心情",
            "老吴我想你了", "老吴顺顺利利", "老吴别熬夜", "老吴天凉了多穿衣服",
            "老吴谢谢你，把冬天过得像春天", "老吴有你真好，谢谢", "老吴你让我很安心",
            "老吴那颗糖甜到心里", "老吴你信我这一次", "老吴有你在我踏实"]
COLORS = ["pink", "lightblue", "lightgreen",
          "lemonchiffon", "hotpink", "skyblue"]

FONT_DIR = r"C:\Windows\Fonts"


def font(name, size):
    return ImageFont.truetype(os.path.join(FONT_DIR, name), size)


def darken(rgb, k=0.08):
    return tuple(max(0, int(c * (1 - k))) for c in rgb)


def is_wide(ch: str) -> bool:
    o = ord(ch)
    return (0x2E80 <= o <= 0x9FFF) or (0xFF00 <= o <= 0xFF60) or (0x3000 <= o <= 0x303F)


# ---------------- 配色 ----------------
BG0, BG1 = (22, 24, 29), (34, 37, 44)
EDITOR_BG = (30, 30, 30)
GUTTER_FG = (133, 133, 133)
LINE_FG = (212, 212, 212)
SYN = {
    "kw": (86, 156, 214),
    "flow": (197, 134, 192),
    "str": (206, 145, 120),
    "num": (181, 206, 168),
    "com": (106, 153, 85),
    "fn": (220, 220, 170),
    "builtin": (78, 201, 176),
    "var": (156, 220, 254),
}

FLOW = {"for", "in", "if", "else", "elif", "while", "lambda", "and", "or",
        "not", "with", "as", "return", "yield", "try", "except", "finally",
        "break", "continue", "pass", "raise", "global", "assert", "del"}
KWS = {"def", "class", "import", "from", "is", "None", "True", "False"}
BUILTINS = {"range", "print", "int", "len", "max", "min", "enumerate",
            "random", "math", "time", "sys", "exit", "isinstance"}


def highlight(src: str):
    """返回 {行号: [(起始列, 结束列, 颜色), ...]}，行号 0 基。"""
    spans = {}
    try:
        toks = list(tokenize.generate_tokens(io.StringIO(src).readline))
    except (tokenize.TokenError, IndentationError):
        return spans
    for tok in toks:
        if tok.type in (tokenize.NEWLINE, tokenize.NL, tokenize.INDENT,
                        tokenize.DEDENT, tokenize.ENCODING, tokenize.ENDMARKER):
            continue
        srow, scol = tok.start
        erow, ecol = tok.end
        if srow != erow:
            continue
        color = None
        if tok.type == tokenize.COMMENT:
            color = SYN["com"]
        elif tok.type == tokenize.STRING:
            color = SYN["str"]
        elif tok.type == tokenize.NUMBER:
            color = SYN["num"]
        elif tok.type == tokenize.NAME:
            if tok.string in FLOW:
                color = SYN["flow"]
            elif tok.string in KWS:
                color = SYN["kw"]
            elif tok.string in BUILTINS:
                color = SYN["builtin"]
            elif re.match(r"^[a-z_][a-z0-9_]{1,}$", tok.string):
                color = SYN["fn"] if tok.string not in (
                    "th", "sx", "sy", "sw", "sh", "tip", "n", "i", "x", "y") else SYN["var"]
        if color:
            spans.setdefault(srow - 1, []).append((scol, ecol, color))
    return spans


# ---------------- 代码编辑器底图 ----------------
def build_editor():
    snippet = os.path.join(ROOT, "assets", "code_snippet.py")
    src = open(snippet, encoding="utf-8").read().rstrip("\n")
    lines = src.split("\n")
    spans = highlight(src)

    f_code = font("consola.ttf", 24)
    f_cjk = font("msyh.ttc", 24)
    f_ln = font("consola.ttf", 18)
    f_ui = font("msyh.ttc", 13)

    img = Image.new("RGB", (W, H), EDITOR_BG)
    d = ImageDraw.Draw(img)

    # 露在外面的桌面
    for y in range(0, H, 2):
        k = y / H
        d.rectangle([0, y, W, y + 2],
                    fill=tuple(int(BG0[i] + (BG1[i] - BG0[i]) * k) for i in range(3)))

    # 标题栏
    d.rectangle([0, 0, W, 34], fill=(50, 50, 55))
    d.rectangle([0, 0, 300, 34], fill=(40, 40, 45))
    d.text((18, 8), "heart_popup.py", font=f_ui, fill=(220, 220, 220))
    d.text((205, 9), "code_snippet.py ×", font=f_ui, fill=(150, 150, 150))
    for i, g in enumerate(["—", "❐", "✕"]):
        d.text((W - 92 + i * 28, 8), g, font=font("msyh.ttc", 13), fill=(200, 200, 200))

    # 编辑区
    d.rectangle([0, 34, W, 872], fill=EDITOR_BG)
    d.rectangle([0, 34, 72, 872], fill=(37, 37, 38))

    cw = f_code.getlength("0")
    line_h = 33
    first = 3
    vis = (872 - 48) // line_h
    y = 48
    for idx in range(first, min(len(lines), first + vis)):
        ln = idx + 1
        d.text((72 - 12, y + 4), str(ln), font=f_ln, fill=GUTTER_FG, anchor="ra")
        text = lines[idx].replace("\t", "    ")
        cmap = {}
        for (a, b, color) in spans.get(idx, []):
            for c in range(a, b):
                cmap[c] = color
        x = 96.0
        for col, ch in enumerate(text):
            if ch == " ":
                x += cw
                continue
            wide = is_wide(ch)
            d.text((x, y), ch, font=(f_cjk if wide else f_code),
                   fill=cmap.get(col, LINE_FG))
            x += cw * (2 if wide else 1)
        y += line_h

    # 终端
    d.rectangle([0, 872, W, 1012], fill=(24, 24, 26))
    d.rectangle([0, 872, W, 874], fill=(60, 60, 65))
    d.text((24, 886), "终端", font=f_ui, fill=(180, 180, 180))
    d.text((92, 887), "powershell", font=f_ui, fill=(120, 120, 120))
    d.text((24, 920), r"C:\Users\Administrator\Desktop> python heart_popup.py",
           font=font("consola.ttf", 17), fill=(205, 212, 205))
    d.text((24, 952), ">>> 正在创建置顶窗口 …",
           font=font("msyh.ttc", 16), fill=(140, 200, 140))

    # 状态栏
    d.rectangle([0, 1012, W, 1036], fill=(0, 122, 204))
    d.text((16, 1015), "main*   heart_popup.py", font=f_ui, fill=(240, 245, 250))
    d.text((W - 270, 1015), "UTF-8   LF   Python 3.12.14",
           font=f_ui, fill=(240, 245, 250))

    # 任务栏
    d.rectangle([0, 1036, W, H], fill=(28, 28, 32))
    d.rectangle([W // 2 - 20, 1046, W // 2 + 20, 1070], fill=(70, 120, 200))
    for i in range(5):
        x = W // 2 + 46 + i * 46
        d.rectangle([x, 1048, x + 24, 1072], fill=(60, 60, 66))
    d.text((W - 150, 1046), "20:31", font=f_ui, fill=(230, 230, 230))
    d.text((W - 150, 1062), "2026/11/7", font=font("msyh.ttc", 12), fill=(190, 190, 190))
    return img


# ---------------- 虚焦开场空镜 ----------------
def build_plate():
    rng = np.random.default_rng(7)
    base = np.zeros((H // 4, W // 4, 3), dtype=np.float32)
    base[:, :] = (238, 236, 232)
    for _ in range(70):
        cx, cy = rng.integers(0, W // 4), rng.integers(0, H // 4)
        rw, rh = rng.integers(20, 170), rng.integers(10, 70)
        col = rng.choice([(200, 200, 205), (175, 180, 186),
                          (250, 248, 244), (150, 160, 170)])
        base[cy:cy + rh, cx:cx + rw] = col
    img = Image.fromarray(base.clip(0, 255).astype(np.uint8)).resize((W, H), Image.BICUBIC)
    img = img.filter(ImageFilter.GaussianBlur(40))
    arr = np.asarray(img).astype(np.float32) * 1.15 + 22
    return Image.fromarray(arr.clip(0, 255).astype(np.uint8))


# ---------------- 弹窗贴图 ----------------
TILE_W, TILE_H = 210, 84
PAD = 4
_tiles: "OrderedDict[tuple, Image.Image]" = OrderedDict()


def popup_tile(bg_name: str, text: str) -> Image.Image:
    key = (bg_name, text)
    if key in _tiles:
        return _tiles[key]
    bg = ImageColor.getrgb(bg_name)
    tile = Image.new("RGBA", (TILE_W + PAD * 2, TILE_H + PAD * 2), (0, 0, 0, 0))
    sh = Image.new("RGBA", tile.size, (0, 0, 0, 0))
    ImageDraw.Draw(sh).rectangle(
        [PAD + 2, PAD + 4, PAD + TILE_W, PAD + TILE_H + 1], fill=(0, 0, 0, 80))
    tile.alpha_composite(sh.filter(ImageFilter.GaussianBlur(2.2)))

    d = ImageDraw.Draw(tile)
    d.rectangle([PAD, PAD, PAD + TILE_W - 1, PAD + TILE_H - 1], fill=bg)
    d.rectangle([PAD, PAD, PAD + TILE_W - 1, PAD + 24], fill=(242, 243, 245))
    d.line([PAD, PAD + 25, PAD + TILE_W - 1, PAD + 25], fill=(215, 216, 220))
    d.rectangle([PAD, PAD, PAD + TILE_W - 1, PAD + TILE_H - 1], outline=(188, 189, 194))

    d.ellipse([PAD + 7, PAD + 8, PAD + 25, PAD + 26], fill=(110, 150, 200))
    d.ellipse([PAD + 11, PAD + 12, PAD + 21, PAD + 22], fill=(242, 243, 245))

    f_btn = font("msyh.ttc", 11)
    for i, glyph in enumerate(["—", "❐", "✕"]):
        d.text((PAD + TILE_W - 80 + i * 25, PAD + 7), glyph,
               font=f_btn, fill=(110, 110, 115))

    size, lines = fit_text(text, TILE_W - 18, 27)
    if len(lines) == 1:
        d.text((PAD + TILE_W / 2, PAD + 54), lines[0],
               font=font("msyh.ttc", size), fill=(28, 28, 28), anchor="mm")
    else:
        step = int(size * 1.24)
        y0 = PAD + 54 - step / 2
        for i, ln in enumerate(lines):
            d.text((PAD + TILE_W / 2, y0 + i * step), ln,
                   font=font("msyh.ttc", size), fill=(28, 28, 28), anchor="mm")
    _tiles[key] = tile
    return tile


def fit_text(text, max_w, base):
    """单行放不下就缩字号，再放不下就折两行。"""
    if font("msyh.ttc", base).getlength(text) <= max_w:
        return base, [text]
    size = base
    while size > int(base * 0.72):
        size -= 1
        if font("msyh.ttc", size).getlength(text) <= max_w:
            return size, [text]
    half = (len(text) + 1) // 2
    lines = [text[:half], text[half:]]
    size = int(base * 0.92)
    while size > 11:
        f = font("msyh.ttc", size)
        if all(f.getlength(x) <= max_w for x in lines):
            return size, lines
        size -= 1
    return 11, lines


# ---------------- 布局 ----------------
def heart_points(n, w, h):
    """心形参数方程，按弧长等距取 n 个点（等 θ 会挤在一侧）。"""
    dense = 4000
    raw = []
    for k in range(dense + 1):
        th = k / dense * 2 * math.pi
        x = 16 * math.sin(th) ** 3
        y = (13 * math.cos(th) - 5 * math.cos(2 * th)
             - 2 * math.cos(3 * th) - math.cos(4 * th))
        raw.append((x, y))
    cum = [0.0]
    for i in range(1, len(raw)):
        dx = raw[i][0] - raw[i - 1][0]
        dy = raw[i][1] - raw[i - 1][1]
        cum.append(cum[-1] + math.hypot(dx, dy))
    total = cum[-1]

    pts = []
    j = 0
    for i in range(n):
        target = total * i / n
        while j < len(cum) - 2 and cum[j + 1] < target:
            j += 1
        span = cum[j + 1] - cum[j] or 1.0
        t = (target - cum[j]) / span
        x = raw[j][0] + (raw[j + 1][0] - raw[j][0]) * t
        y = raw[j][1] + (raw[j + 1][1] - raw[j][1]) * t
        cx = w / 2 + x * HEART_S
        cy = h / 2 - y * HEART_S - HEART_DY
        sx = int(cx - TILE_W / 2)
        sy = int(cy - TILE_H / 2)
        pts.append((max(0, min(sx, w - TILE_W)), max(0, min(sy, h - TILE_H))))
    return pts


def build_schedule():
    rng = random.Random(20261107)
    heart = []
    for i, (x, y) in enumerate(heart_points(HEART_N, W, H)):
        heart.append((T_HEART_START + i * HEART_STEP, x, y,
                      rng.choice(COLORS), rng.choice(MESSAGES)))
    flood = []
    for i in range(FLOOD_N):
        flood.append((T_FLOOD_START + i * FLOOD_STEP,
                      rng.randint(0, W - TILE_W), rng.randint(0, H - TILE_H),
                      rng.choice(COLORS), rng.choice(MESSAGES)))
    return heart, flood


# ---------------- 后期 ----------------
def build_vignette():
    yy, xx = np.mgrid[0:H, 0:W]
    cx, cy = W / 2, H / 2
    r = np.sqrt(((xx - cx) / cx) ** 2 + ((yy - cy) / cy) ** 2)
    a = (np.clip((r - 0.55) / 0.75, 0, 1) ** 1.5 * 95).astype(np.uint8)
    v = np.zeros((H, W, 4), dtype=np.uint8)
    v[..., 3] = a
    return Image.fromarray(v, "RGBA")


def build_title():
    f = font("msyh.ttc", 52)
    txt = "心之所向   爱之所往"
    layer = Image.new("RGBA", (W, 240), (0, 0, 0, 0))
    d = ImageDraw.Draw(layer)
    d.text((W / 2 + 2, 122), txt, font=f, fill=(0, 0, 0, 170), anchor="mm")
    d.text((W / 2, 120), txt, font=f, fill=(255, 255, 255, 255), anchor="mm")
    return layer


# ---------------- 主渲染 ----------------
def compose(t, editor, plate, vig, title, heart, flood):
    if t < T_BLACK_END:
        return Image.new("RGB", (W, H), (0, 0, 0))

    if t < T_PLATE_END:
        k = (t - T_BLACK_END) / (T_PLATE_END - T_BLACK_END)
        z = 1.0 + 0.06 * k
        cw_, ch_ = int(W / z), int(H / z)
        x0, y0 = (W - cw_) // 2, (H - ch_) // 2
        return plate.crop((x0, y0, x0 + cw_, y0 + ch_)).resize((W, H), Image.BICUBIC)

    if t < T_WIPE_END:
        k = (t - T_PLATE_END) / (T_WIPE_END - T_PLATE_END)
        cut = int(k * H)
        mosaic = editor.resize((W // 26, H // 26), Image.NEAREST).resize((W, H), Image.NEAREST)
        img = plate.copy()
        if cut > 0:
            img.paste(mosaic.crop((0, 0, W, cut)), (0, 0))
        if 0 < cut < H:
            ImageDraw.Draw(img).line([0, cut, W, cut], fill=(255, 255, 255), width=3)
        return img

    img = editor.copy()

    def draw_windows(items):
        for (_st, x, y, col, txt) in items:
            tile = popup_tile(col, txt)
            img.paste(tile, (x - PAD, y - PAD), tile)

    if T_HEART_START <= t < T_DESTROY:
        k = min(len(heart), int((t - T_HEART_START) / HEART_STEP) + 1)
        draw_windows(heart[:k])
    elif t >= T_FLOOD_START:
        k = min(len(flood), int((t - T_FLOOD_START) / FLOOD_STEP) + 1)
        draw_windows(flood[:k])

    img = img.convert("RGBA")
    img.alpha_composite(vig)
    if t >= T_TITLE_IN:
        k = min(1.0, (t - T_TITLE_IN) / 0.70)
        layer = title.copy()
        alpha = layer.split()[3].point(lambda v: int(v * k))
        layer.putalpha(alpha)
        img.alpha_composite(layer, (0, int(H * 0.40) - 120))
    return img.convert("RGB")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=os.path.join(ROOT, "out", "爱心弹窗-复刻.mp4"))
    ap.add_argument("--preview", default=None)
    ap.add_argument("--audio", default=os.path.join(ROOT, "assets", "bed.wav"))
    args = ap.parse_args()

    print("build editor…", flush=True)
    editor = build_editor()
    plate = build_plate()
    vig = build_vignette()
    title = build_title()
    heart, flood = build_schedule()
    print(f"windows: heart={len(heart)} flood={len(flood)}", flush=True)

    if args.preview:
        outdir = os.path.join(ROOT, "preview")
        os.makedirs(outdir, exist_ok=True)
        for s in args.preview.split(","):
            t = float(s)
            img = compose(t, editor, plate, vig, title, heart, flood)
            p = os.path.join(outdir, f"t{t:05.2f}.png")
            img.save(p)
            print("preview", p)
        return 0

    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    cmd = [FFMPEG, "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
           "-s", f"{W}x{H}", "-r", str(FPS), "-i", "-"]
    has_audio = os.path.exists(args.audio)
    if has_audio:
        cmd += ["-i", args.audio]
    cmd += ["-c:v", "libx264", "-preset", "medium", "-crf", "19",
            "-pix_fmt", "yuv420p", "-movflags", "+faststart"]
    if has_audio:
        cmd += ["-c:a", "aac", "-b:a", "192k", "-shortest"]
    cmd += [args.out]

    proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    n = int(DUR * FPS)
    for i in range(n):
        t = i / FPS
        proc.stdin.write(compose(t, editor, plate, vig, title, heart, flood).tobytes())
        if i % 90 == 0:
            print(f"  frame {i}/{n}  t={t:5.2f}", flush=True)
    proc.stdin.close()
    proc.wait()
    print("OK", args.out, os.path.getsize(args.out) // 1024, "KB")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())