#!/usr/bin/env python3
"""Premium App Store marketing screenshot framer.

Upgrades over frame_screenshot.py:
  - Realistic iPhone device bezel (rounded metal edge + Dynamic Island).
  - Background with depth: base vertical gradient + soft radial glow + vignette.
  - SF Pro headline with an optional lighter subtitle line, CJK-capable fonts.
  - Soft, realistic device drop shadow.

Requirements: pip install Pillow

Example:
  python frame_premium.py --input raw.png --output out.png \
    --width 1320 --height 2868 \
    --headline "Watch your wealth grow" --subline "Track net worth over time" \
    --grad-top #3B82F6 --grad-bottom #10B981
"""
from __future__ import annotations

import argparse
import math
from pathlib import Path

from PIL import Image, ImageDraw, ImageFilter, ImageFont


# ---------- fonts ----------
HEADLINE_FONTS = [
    "/System/Library/Fonts/SFNSDisplay.ttf",
    "/System/Library/Fonts/SFNS.ttf",
    "/System/Library/Fonts/Supplemental/HelveticaNeue.ttc",
    "/Library/Fonts/Arial Unicode.ttf",
]
CJK_FONTS = [
    "/System/Library/Fonts/ヒラギノ角ゴシック W6.ttc",
    "/System/Library/Fonts/Hiragino Sans GB.ttc",
    "/System/Library/Fonts/PingFang.ttc",
]


def _has_cjk(s: str) -> bool:
    return any(ord(c) > 0x2E00 for c in s)


def load_font(size: int, text: str, weight_paths=None) -> ImageFont.FreeTypeFont:
    paths = (CJK_FONTS if _has_cjk(text) else []) + (weight_paths or HEADLINE_FONTS)
    for p in paths:
        if Path(p).exists():
            try:
                return ImageFont.truetype(p, size)
            except Exception:
                continue
    return ImageFont.load_default()


def hex_rgb(h: str) -> tuple[int, int, int]:
    h = h.lstrip("#")
    return tuple(int(h[i : i + 2], 16) for i in (0, 2, 4))  # type: ignore


def lerp(a, b, t):
    return tuple(int(a[i] + (b[i] - a[i]) * t) for i in range(3))


# ---------- background ----------
def make_background(w: int, h: int, top: tuple, bottom: tuple) -> Image.Image:
    base = Image.new("RGB", (w, h))
    px = base.load()
    for y in range(h):
        c = lerp(top, bottom, y / max(1, h - 1))
        for x in range(w):
            px[x, y] = c
    # soft radial glow, lighter, upper-center
    glow = Image.new("L", (w, h), 0)
    gd = ImageDraw.Draw(glow)
    cx, cy = w * 0.5, h * 0.28
    r = w * 0.75
    gd.ellipse([cx - r, cy - r, cx + r, cy + r], fill=90)
    glow = glow.filter(ImageFilter.GaussianBlur(w * 0.18))
    light = Image.new("RGB", (w, h), lerp(top, (255, 255, 255), 0.35))
    base = Image.composite(light, base, glow)
    # vignette (darken edges)
    vig = Image.new("L", (w, h), 0)
    vd = ImageDraw.Draw(vig)
    vd.rectangle([0, 0, w, h], fill=70)
    vd.rounded_rectangle([w * 0.06, h * 0.05, w * 0.94, h * 0.95],
                         radius=int(w * 0.1), fill=0)
    vig = vig.filter(ImageFilter.GaussianBlur(w * 0.12))
    dark = Image.new("RGB", (w, h), lerp(bottom, (0, 0, 0), 0.35))
    base = Image.composite(dark, base, vig)
    return base


# ---------- device ----------
def rounded_mask(size, radius) -> Image.Image:
    m = Image.new("L", size, 0)
    ImageDraw.Draw(m).rounded_rectangle([0, 0, size[0], size[1]], radius=radius, fill=255)
    return m


def build_device(screenshot: Image.Image, target_w: int) -> Image.Image:
    """Return an RGBA image of the phone (bezel + screen + island) at target_w width."""
    bezel = max(10, int(target_w * 0.028))          # metal frame thickness
    screen_w = target_w - bezel * 2
    ar = screenshot.height / screenshot.width
    screen_h = int(screen_w * ar)
    shot = screenshot.resize((screen_w, screen_h), Image.LANCZOS).convert("RGB")

    scr_radius = int(target_w * 0.11)
    dev_radius = scr_radius + bezel
    dev_w, dev_h = target_w, screen_h + bezel * 2

    canvas = Image.new("RGBA", (dev_w, dev_h), (0, 0, 0, 0))
    # metal frame (dark) with a subtle lighter outer edge
    frame = Image.new("RGBA", (dev_w, dev_h), (0, 0, 0, 0))
    fd = ImageDraw.Draw(frame)
    fd.rounded_rectangle([0, 0, dev_w, dev_h], radius=dev_radius, fill=(18, 18, 20, 255))
    fd.rounded_rectangle([1, 1, dev_w - 1, dev_h - 1], radius=dev_radius, outline=(60, 62, 68, 255), width=2)
    canvas.alpha_composite(frame)
    # screen
    scr_mask = rounded_mask((screen_w, screen_h), scr_radius)
    canvas.paste(shot, (bezel, bezel), scr_mask)
    # Dynamic Island
    di_w, di_h = int(screen_w * 0.30), int(target_w * 0.032)
    di_x = bezel + (screen_w - di_w) // 2
    di_y = bezel + int(target_w * 0.022)
    ImageDraw.Draw(canvas).rounded_rectangle(
        [di_x, di_y, di_x + di_w, di_y + di_h], radius=di_h // 2, fill=(0, 0, 0, 255)
    )
    return canvas


# ---------- text ----------
def wrap(draw, text, font, max_w):
    words, lines, cur = text.split(), [], ""
    for wd in words:
        t = (cur + " " + wd).strip()
        if draw.textlength(t, font=font) <= max_w or not cur:
            cur = t
        else:
            lines.append(cur)
            cur = wd
    if cur:
        lines.append(cur)
    return lines


def draw_center_text(img, text, font, y, color, line_gap=1.12):
    d = ImageDraw.Draw(img)
    lines = wrap(d, text, font, int(img.width * 0.86))
    asc, desc = font.getmetrics()
    lh = int((asc + desc) * line_gap)
    for i, ln in enumerate(lines):
        w = d.textlength(ln, font=font)
        d.text(((img.width - w) / 2, y + i * lh), ln, font=font, fill=color)
    return y + len(lines) * lh


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--output", required=True)
    ap.add_argument("--width", type=int, required=True)
    ap.add_argument("--height", type=int, required=True)
    ap.add_argument("--headline", required=True)
    ap.add_argument("--subline", default="")
    ap.add_argument("--grad-top", default="#3B82F6")
    ap.add_argument("--grad-bottom", default="#10B981")
    ap.add_argument("--device-scale", type=float, default=0.82)
    a = ap.parse_args()

    W, H = a.width, a.height
    bg = make_background(W, H, hex_rgb(a.grad_top), hex_rgb(a.grad_bottom)).convert("RGBA")

    # headline + subline
    hfont = load_font(int(W * 0.082), a.headline)
    y = int(H * 0.055)
    y = draw_center_text(bg, a.headline, hfont, y, (255, 255, 255, 255))
    if a.subline:
        sfont = load_font(int(W * 0.036), a.subline)
        y = draw_center_text(bg, a.subline, sfont, y + int(H * 0.006),
                             (255, 255, 255, 210))

    # device
    shot = Image.open(a.input).convert("RGB")
    device = build_device(shot, int(W * a.device_scale))
    dx = (W - device.width) // 2
    dy = y + int(H * 0.03)
    # keep device within canvas
    if dy + device.height > H - int(H * 0.03):
        dy = H - int(H * 0.03) - device.height

    # shadow
    shadow = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    sh = Image.new("RGBA", device.size, (0, 0, 0, 150))
    sh.putalpha(device.split()[3].point(lambda p: int(p * 0.55)))
    shadow.alpha_composite(sh, (dx, dy + int(W * 0.02)))
    shadow = shadow.filter(ImageFilter.GaussianBlur(int(W * 0.03)))
    bg.alpha_composite(shadow)
    bg.alpha_composite(device, (dx, dy))

    out = bg.convert("RGB")
    Path(a.output).parent.mkdir(parents=True, exist_ok=True)
    out.save(a.output, "PNG")
    assert out.size == (W, H), f"size {out.size} != {(W, H)}"
    print(f"OK {a.output} {out.size}")


if __name__ == "__main__":
    main()
