#!/usr/bin/env python3
"""Logo sting / reveal: a 2-5 s animated logo for intros, outros and ad end cards. Six styles, any background
(colour, gradient, image or a Kenerate clip), optional tagline, and an optional transparent .mov to drop into any
editor. Silent unless you pass --audio (ask the user first).

Styles: pop (spring + ring burst), sweep (fade up + light sweep across the logo), glitch (RGB split, slices,
settles), reveal (accent bar wipes the logo on), zoom-blur (rushes in from a blur), assemble (tiles fly together).

Examples
  python sting.py --logo logo.png --style sweep --tagline "brewandbloom.com" --bg "gradient:#0F172A,#1E293B"
  python sting.py --logo logo.png --style glitch --bg https://.../loop.mp4 --aspect 9:16 --alpha
  python sting.py --logo logo.png --showcase              # all six styles in one contact sheet (pick one)
"""
from __future__ import annotations

import argparse
import math
import random
import sys
from pathlib import Path

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

STYLES = ("pop", "sweep", "glitch", "reveal", "zoom-blur", "assemble")


class Sting:
    def __init__(self, logo, W: int, H: int, D: float, style: str, accent, tagline: str, font, scale: float):
        from PIL import Image
        self.W, self.H, self.D, self.style, self.accent = W, H, D, style, accent
        s = min(W, H)
        box = int(s * 0.42 * scale)
        lg = logo.convert("RGBA")
        lg.thumbnail((box, box), Image.LANCZOS)
        self.logo = lg
        self.tag = K.text_img(tagline, int(s * 0.045), (255, 255, 255, 255), "regular" if not font else "bold", font) if tagline else None
        self.cx, self.cy = W / 2, H / 2 - (s * 0.04 if tagline else 0)
        rnd = random.Random(7)
        self.tiles = []
        if style == "assemble":
            n = 6
            tw, th = math.ceil(lg.width / n), math.ceil(lg.height / n)
            for i in range(n):
                for j in range(n):
                    t = lg.crop((i * tw, j * th, min((i + 1) * tw, lg.width), min((j + 1) * th, lg.height)))
                    if t.getbbox() is None:
                        continue
                    ang = rnd.uniform(0, 2 * math.pi)
                    dist = rnd.uniform(0.4, 0.9) * s
                    self.tiles.append((t, i * tw + t.width / 2 - lg.width / 2, j * th + t.height / 2 - lg.height / 2,
                                       math.cos(ang) * dist, math.sin(ang) * dist, rnd.uniform(-160, 160), rnd.uniform(0, 0.35)))
        self.glitch_seed = [(rnd.random(), rnd.random(), rnd.random()) for _ in range(64)]

    def frame(self, t: float, base=None):
        from PIL import Image, ImageChops, ImageDraw, ImageFilter
        W, H = self.W, self.H
        canvas = base.copy().convert("RGBA") if base is not None else Image.new("RGBA", (W, H), (0, 0, 0, 0))
        lg, cx, cy, st = self.logo, self.cx, self.cy, self.style
        drift = 1 + 0.03 * K.clamp((t - 1.0) / max(self.D - 1.0, 0.1))  # slow push while holding
        if st == "pop":
            p = (t - 0.1) / 0.55
            if p > 0:
                rp = K.clamp((t - 0.1) / 0.8)
                if rp < 1:
                    r = int(max(lg.size) * (0.35 + 0.9 * K.ease_out_cubic(rp)))
                    ring = Image.new("RGBA", (r * 2 + 8, r * 2 + 8), (0, 0, 0, 0))
                    ImageDraw.Draw(ring).ellipse((4, 4, r * 2 + 4, r * 2 + 4), outline=tuple(self.accent) + (int(255 * (1 - rp)),),
                                                 width=max(3, min(W, H) // 120))
                    K.place(canvas, ring, cx, cy)
                K.place(canvas, lg, cx, cy, K.ease_out_back(p, 1.7) * drift, K.clamp(p * 2.5))
        elif st == "sweep":
            p = K.ease_out_cubic(t / 0.6)
            frame = lg
            sp = (t - 0.55) / 0.75
            if 0 < sp < 1:
                band = Image.new("L", lg.size, 0)
                d = ImageDraw.Draw(band)
                x = -lg.width * 0.6 + sp * lg.width * 2.2
                bw = lg.width * 0.22
                d.polygon([(x, 0), (x + bw, 0), (x + bw - lg.height * 0.5, lg.height), (x - lg.height * 0.5, lg.height)], fill=210)
                band = band.filter(ImageFilter.GaussianBlur(max(2, lg.width // 40)))
                band = ImageChops.multiply(band, lg.getchannel("A"))
                shine = Image.new("RGBA", lg.size, (255, 255, 255, 0))
                shine.putalpha(band)
                frame = lg.copy()
                frame.alpha_composite(shine)
            K.place(canvas, frame, cx, cy + (1 - p) * min(W, H) * 0.03, (0.94 + 0.06 * p) * drift, p)
        elif st == "glitch":
            g = 1 - K.clamp((t - 0.15) / 0.85)
            if t < 0.05:
                pass
            elif g <= 0.01:
                K.place(canvas, lg, cx, cy, drift)
            else:
                k = int(t * 24)
                rnd = self.glitch_seed[k % len(self.glitch_seed)]
                r, gch, b, a = lg.split()
                off = int(min(W, H) * 0.03 * g)
                red = Image.merge("RGBA", (r, Image.new("L", lg.size, 0), Image.new("L", lg.size, 0), a))
                cyan = Image.merge("RGBA", (Image.new("L", lg.size, 0), gch, b, a))
                lay = Image.new("RGBA", (lg.width + off * 2 + 2, lg.height), (0, 0, 0, 0))
                lay.alpha_composite(K.with_opacity(red, 0.9), (0, 0))
                lay.alpha_composite(K.with_opacity(cyan, 0.9), (off * 2, 0))
                # horizontal slice displacement
                ms = int(lay.width * 0.13) + 1
                out = Image.new("RGBA", (lay.width + ms * 2, lay.height), (0, 0, 0, 0))
                y, i = 0, 0
                while y < lay.height:
                    hh = max(2, int(lay.height * (0.04 + 0.12 * self.glitch_seed[(k + i) % 64][0])))
                    dx = int((self.glitch_seed[(k * 3 + i) % 64][1] - 0.5) * 2 * ms * g)
                    out.alpha_composite(lay.crop((0, y, lay.width, min(lay.height, y + hh))), (ms + dx, y))
                    y += hh
                    i += 1
                flick = 0.35 + 0.65 * (rnd[2] > 0.25 * g)
                K.place(canvas, out, cx + (rnd[0] - 0.5) * off * 2, cy, drift, flick)
                if rnd[1] < 0.5 * g:  # stray scan bar
                    d = ImageDraw.Draw(canvas)
                    yy = int(cy - lg.height / 2 + rnd[0] * lg.height)
                    d.rectangle((0, yy, W, yy + max(2, H // 160)), fill=tuple(self.accent) + (int(160 * g),))
        elif st == "reveal":
            p = K.ease_in_out((t - 0.15) / 0.75)
            x0 = cx - lg.width / 2 - lg.width * 0.08
            span = lg.width * 1.16
            bar_x = x0 + span * p
            if p > 0:
                vis = lg.crop((0, 0, max(1, int(min(lg.width, bar_x - (cx - lg.width / 2)))), lg.height)) if bar_x > cx - lg.width / 2 else None
                if vis is not None and vis.width > 1:
                    canvas_l = Image.new("RGBA", lg.size, (0, 0, 0, 0))
                    canvas_l.alpha_composite(vis, (0, 0))
                    K.place(canvas, canvas_l, cx, cy, drift if p >= 1 else 1.0)
            if 0 < p < 1 or (1 <= p and t < 1.15):
                fade = 1 - K.clamp((t - 0.9) / 0.25)
                d = ImageDraw.Draw(canvas)
                bw = max(4, int(lg.width * 0.035))
                d.rectangle((bar_x - bw, cy - lg.height * 0.62, bar_x, cy + lg.height * 0.62), fill=tuple(self.accent) + (int(255 * fade),))
        elif st == "zoom-blur":
            p = K.ease_out_cubic(t / 0.7)
            if p > 0:
                rad = (1 - p) * min(W, H) * 0.03
                im = lg.filter(ImageFilter.GaussianBlur(rad)) if rad > 0.5 else lg
                K.place(canvas, im, cx, cy, (2.2 - 1.2 * p) * drift, K.clamp(p * 1.6))
        elif st == "assemble" and t > 0.35 + 0.75 + 0.02:  # landed: draw the whole logo (no tile seams)
            K.place(canvas, lg, cx, cy, drift)
        elif st == "assemble":
            for (tile, ox, oy, fx, fy, rot, delay) in self.tiles:
                p = K.ease_out_cubic((t - delay) / 0.75)
                if p <= 0:
                    continue
                K.place(canvas, tile, cx + ox * drift + fx * (1 - p), cy + oy * drift + fy * (1 - p), drift, K.clamp(p * 2), rot * (1 - p))
        if self.tag is not None:
            q = K.ease_out_cubic((t - 1.0) / 0.45)
            K.place(canvas, self.tag, W / 2, cy + lg.height / 2 + min(W, H) * 0.06 + (1 - q) * 18, 1.0, q)
        return canvas


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--logo", required=True, help="logo PNG (transparent works best)")
    ap.add_argument("--style", choices=STYLES, default="pop")
    ap.add_argument("--tagline", default="")
    ap.add_argument("--bg", default="gradient:#0B0F19,#1F2937", help="'#hex', 'gradient:#a,#b', image or video")
    ap.add_argument("--accent", default="#FFD60A")
    ap.add_argument("--scale", type=float, default=1.0, help="logo size multiplier")
    ap.add_argument("--duration", type=float, default=3.0)
    ap.add_argument("--outro", choices=["none", "fade"], default="none", help="fade to black at the end")
    ap.add_argument("--aspect", default="16:9")
    ap.add_argument("--res", type=int, default=1080)
    ap.add_argument("--fps", type=int, default=30)
    ap.add_argument("--audio", help="sound for the sting (only one the user supplied or approved)")
    ap.add_argument("--alpha", action="store_true", help="also write a transparent .mov (PNG codec) for editors")
    ap.add_argument("--showcase", action="store_true", help="render a key-frame sheet of all six styles and stop")
    K.add_font_args(ap, "TTF/OTF for the tagline")
    ap.add_argument("--out", default="kenerate-out/logo-sting/sting.mp4")
    ap.add_argument("--keep-temp", action="store_true")
    a = K.parse_args(ap)

    st = K.Status(4, "logo-sting")
    K.require(status=st)
    W, H = K.size_for(a.aspect, a.res)
    tmp = K.workdir(a.keep_temp)
    out = Path(a.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    D = max(1.5, a.duration)

    st.step("Fetching logo + background")
    logo = K.open_image(K.fetch(a.logo, tmp, "logo", st)).convert("RGBA")
    bbox = logo.getbbox()
    if bbox:
        logo = logo.crop(bbox)
    if logo.getchannel("A").getextrema()[0] == 255:
        st.warn("the logo has no transparency: it will sit on its own background box (remove_background first for a clean sting)")
    base = K.parse_bg(a.bg, W, H)
    bg_video = None
    if base is None:
        p = K.fetch(a.bg, tmp, "bg", st)
        if K.kind_of(str(p)) == "image":
            base = K.cover(K.open_image(p).convert("RGB"), W, H)
        else:
            bg_video = p
    audio = K.fetch(a.audio, tmp, "audio", st) if a.audio else None
    accent = K.hex_rgb(a.accent)

    if a.showcase:
        st.step("Rendering all styles")
        from PIL import Image
        cells, labels = [], []
        bg_img = base if base is not None else K.cover(K.open_image(K.extract_frame(bg_video, 0.5, tmp / "bg.jpg")).convert("RGB"), W, H)
        for sname in STYLES:
            sc = Sting(logo, W, H, D, sname, accent, a.tagline, a.font, a.scale)
            for t in (0.25, 0.55, D - 0.2):
                p = tmp / f"{sname}_{t:.2f}.jpg"
                sc.frame(t, bg_img).convert("RGB").save(p, quality=85)
                cells.append(p)
                labels.append(f"{sname} {t:.2f}s")
        sheet = K.contact_sheet(cells, out.with_name("styles_showcase.jpg"), cols=6, cell=300, labels=labels)
        K.finish({"showcase": str(sheet.resolve()), "styles": list(STYLES)})
        return

    st.step(f"Laying out '{a.style}' ({W}x{H}, {D:.1f} s)")
    sc = Sting(logo, W, H, D, a.style, accent, a.tagline, a.font, a.scale)
    n = int(round(D * a.fps))
    ex = f"fade=t=out:st={D - 0.35:.3f}:d=0.35" if a.outro == "fade" else ""

    st.step(f"Rendering {n} frames")
    if bg_video is not None:
        K.FramePipe(out, W, H, a.fps, n, background=["-stream_loop", "-1", "-i", str(bg_video)], bg_filter=K.fit_filter(W, H, "cover"),
                    audio=audio, keep_bg_audio=False, extra_filter=ex, status=st).run(sc.frame(i / a.fps) for i in range(n))
    else:
        K.FramePipe(out, W, H, a.fps, n, rgba=False, audio=audio, extra_filter=ex, status=st).run(
            sc.frame(i / a.fps, base).convert("RGB") for i in range(n))
    alpha = None
    if a.alpha:
        alpha = out.with_suffix(".mov")
        st.info("transparent .mov")
        K.encode_alpha_mov((sc.frame(i / a.fps) for i in range(n)), alpha, W, H, a.fps, n, st)

    st.step("Preview")
    sheet = K.video_preview(out, out.with_name(out.stem + "_preview.jpg"), tmp=tmp / "pv",
                            times=[0.15, 0.35, 0.6, 0.9, 1.3, D - 0.1])
    K.finish({"output": K.file_info(out), "alpha_mov": str(alpha.resolve()) if alpha else None, "preview": str(sheet.resolve()),
              "style": a.style, "seconds": D, "audio": bool(audio), "elapsed_s": st.elapsed()})


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