#!/usr/bin/env python3
"""Frame raw app screenshots for App Store marketing artwork.

Requirements:
    pip install Pillow
"""

from __future__ import annotations

import argparse
import json
import re
import sys
import textwrap
import warnings
from pathlib import Path
from typing import Iterable, Sequence

from PIL import Image, ImageDraw, ImageFilter, ImageFont


DEFAULT_GRAD_TOP = "#0B5FFF"
DEFAULT_GRAD_BOTTOM = "#00C2A8"
HEX_RE = re.compile(r"#[0-9A-Fa-f]{6}\b")

# --- Dark premium / fintech theme -----------------------------------------
DARK_GRAD_TOP = "#0F172A"       # slate-900
DARK_GRAD_BOTTOM = "#020617"    # near black
HEADLINE_COLOR = (248, 250, 252)   # #F8FAFC near-white
ACCENT_COLOR = (52, 211, 153)      # #34D399 vibrant green
MUTED_COLOR = (148, 163, 184)      # #94A3B8 muted slate


FONT_CANDIDATES = (
    # CJK-capable macOS fonts first.
    "/System/Library/Fonts/ヒラギノ角ゴシック W6.ttc",
    "/System/Library/Fonts/ヒラギノ角ゴシック W5.ttc",
    "/System/Library/Fonts/Hiragino Sans GB.ttc",
    "/System/Library/Fonts/PingFang.ttc",
    "/System/Library/Fonts/Supplemental/Arial Unicode.ttf",
    "/Library/Fonts/Arial Unicode.ttf",
    # Strong Latin fallbacks common on macOS.
    "/System/Library/Fonts/Supplemental/Arial Bold.ttf",
    "/System/Library/Fonts/Supplemental/Arial.ttf",
    "/System/Library/Fonts/SFNS.ttf",
)

# Bold-weight candidates for headline / stat callout.
BOLD_FONT_CANDIDATES = (
    "/System/Library/Fonts/ヒラギノ角ゴシック W7.ttc",
    "/System/Library/Fonts/ヒラギノ角ゴシック W6.ttc",
    "/System/Library/Fonts/Supplemental/Arial Bold.ttf",
    "/System/Library/Fonts/SFNS.ttf",
    "/System/Library/Fonts/Supplemental/Arial.ttf",
)

# Regular-weight candidates for secondary / substat text.
REGULAR_FONT_CANDIDATES = (
    "/System/Library/Fonts/ヒラギノ角ゴシック W4.ttc",
    "/System/Library/Fonts/ヒラギノ角ゴシック W5.ttc",
    "/System/Library/Fonts/Supplemental/Arial.ttf",
    "/System/Library/Fonts/SFNS.ttf",
    "/System/Library/Fonts/Supplemental/Arial Unicode.ttf",
)


def parse_hex_color(value: str) -> tuple[int, int, int]:
    """Return an RGB tuple from #RRGGBB."""
    if not HEX_RE.fullmatch(value):
        raise argparse.ArgumentTypeError(f"Expected #RRGGBB hex color, got {value!r}")
    value = value.lstrip("#")
    return tuple(int(value[i : i + 2], 16) for i in (0, 2, 4))


def find_repo_root(start: Path) -> Path:
    for path in (start, *start.parents):
        if (path / ".git").exists():
            return path
    return start


def find_brand_colors(repo_root: Path) -> tuple[str | None, str | None]:
    """Look for SwiftUI Color.appPrimary/appSecondary hex values, then assets."""
    primary = secondary = None

    for swift_file in repo_root.rglob("*.swift"):
        try:
            text = swift_file.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            text = swift_file.read_text(errors="ignore")

        primary_match = re.search(r"appPrimary\s*=\s*Color\(hex:\s*\"(#[0-9A-Fa-f]{6})\"", text)
        secondary_match = re.search(r"appSecondary\s*=\s*Color\(hex:\s*\"(#[0-9A-Fa-f]{6})\"", text)
        if primary_match:
            primary = primary_match.group(1).upper()
        if secondary_match:
            secondary = secondary_match.group(1).upper()
        if primary and secondary:
            return primary, secondary

    for contents_file in repo_root.glob("**/*.colorset/Contents.json"):
        name = contents_file.parent.name.lower()
        if "primary" not in name and "secondary" not in name:
            continue
        color = read_xcasset_hex(contents_file)
        if not color:
            continue
        if "primary" in name and not primary:
            primary = color
        elif "secondary" in name and not secondary:
            secondary = color

    return primary, secondary


def read_xcasset_hex(path: Path) -> str | None:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None

    for color_entry in data.get("colors", []):
        components = color_entry.get("color", {}).get("components", {})
        red, green, blue = components.get("red"), components.get("green"), components.get("blue")
        if red is None or green is None or blue is None:
            continue
        try:
            values = [component_to_byte(channel) for channel in (red, green, blue)]
        except ValueError:
            continue
        return "#{:02X}{:02X}{:02X}".format(*values)
    return None


def component_to_byte(value: str | float | int) -> int:
    if isinstance(value, str) and value.startswith("0x"):
        return int(value, 16)
    number = float(value)
    if number <= 1:
        number *= 255
    return max(0, min(255, round(number)))


def print_brand_colors(repo_root: Path) -> tuple[str, str]:
    primary, secondary = find_brand_colors(repo_root)
    top = primary or DEFAULT_GRAD_TOP
    bottom = secondary or DEFAULT_GRAD_BOTTOM
    source = "found" if primary and secondary else "default"
    print(f"Brand colors ({source}): appPrimary={top} appSecondary={bottom}")
    return top, bottom


def make_vertical_gradient(size: tuple[int, int], top: str, bottom: str) -> Image.Image:
    width, height = size
    top_rgb = parse_hex_color(top)
    bottom_rgb = parse_hex_color(bottom)
    gradient = Image.new("RGB", size)
    draw = ImageDraw.Draw(gradient)

    for y in range(height):
        ratio = y / max(1, height - 1)
        color = tuple(round(top_rgb[i] + (bottom_rgb[i] - top_rgb[i]) * ratio) for i in range(3))
        draw.line([(0, y), (width, y)], fill=color)

    return gradient.convert("RGBA")


def load_font(
    size: int,
    candidates: Sequence[str] = FONT_CANDIDATES,
) -> ImageFont.ImageFont:
    for font_path in candidates:
        path = Path(font_path)
        if not path.exists():
            continue
        try:
            return ImageFont.truetype(str(path), size=size)
        except OSError:
            continue

    warnings.warn(
        "No CJK-capable macOS font found; falling back to PIL default. "
        "Japanese or other non-Latin text may not render correctly.",
        RuntimeWarning,
    )
    return ImageFont.load_default()


def load_bold_font(size: int) -> ImageFont.ImageFont:
    return load_font(size, BOLD_FONT_CANDIDATES)


def load_regular_font(size: int) -> ImageFont.ImageFont:
    return load_font(size, REGULAR_FONT_CANDIDATES)


def text_bbox(draw: ImageDraw.ImageDraw, text: str, font: ImageFont.ImageFont) -> tuple[int, int]:
    if not text:
        return 0, 0
    left, top, right, bottom = draw.multiline_textbbox((0, 0), text, font=font, spacing=8, align="center")
    return right - left, bottom - top


def wrap_headline(
    text: str,
    font: ImageFont.ImageFont,
    max_width: int,
    draw: ImageDraw.ImageDraw,
) -> str:
    if contains_cjk(text):
        return wrap_cjk_text(text, font, max_width, draw)
    return wrap_word_text(text, font, max_width, draw)


def contains_cjk(text: str) -> bool:
    return any(
        "\u3040" <= char <= "\u30ff"
        or "\u3400" <= char <= "\u4dbf"
        or "\u4e00" <= char <= "\u9fff"
        or "\uf900" <= char <= "\ufaff"
        for char in text
    )


def wrap_word_text(
    text: str,
    font: ImageFont.ImageFont,
    max_width: int,
    draw: ImageDraw.ImageDraw,
) -> str:
    words = text.split()
    if not words:
        return text

    lines: list[str] = []
    current = words[0]
    for word in words[1:]:
        candidate = f"{current} {word}"
        if text_bbox(draw, candidate, font)[0] <= max_width:
            current = candidate
        else:
            lines.append(current)
            current = word
    lines.append(current)

    return "\n".join(lines)


def wrap_cjk_text(
    text: str,
    font: ImageFont.ImageFont,
    max_width: int,
    draw: ImageDraw.ImageDraw,
) -> str:
    lines: list[str] = []
    current = ""

    for char in text:
        if char.isspace():
            if current:
                lines.append(current)
                current = ""
            continue

        candidate = f"{current}{char}"
        if current and text_bbox(draw, candidate, font)[0] > max_width:
            lines.append(current)
            current = char
        else:
            current = candidate

    if current:
        lines.append(current)
    return "\n".join(lines)


def fit_headline(
    headline: str,
    canvas_size: tuple[int, int],
    max_width: int,
    max_height: int,
    draw: ImageDraw.ImageDraw,
    max_scale: float = 0.052,
    font_loader=load_font,
) -> tuple[str, ImageFont.ImageFont, int, int]:
    _, height = canvas_size
    max_size = max(34, round(height * max_scale))
    min_size = 28

    for size in range(max_size, min_size - 1, -2):
        font = font_loader(size)
        wrapped = wrap_headline(headline, font, max_width, draw)
        text_width, text_height = text_bbox(draw, wrapped, font)
        if text_width <= max_width and text_height <= max_height:
            return wrapped, font, text_width, text_height

    font = font_loader(min_size)
    wrapped = wrap_headline(headline, font, max_width, draw)
    text_width, text_height = text_bbox(draw, wrapped, font)
    return wrapped, font, text_width, text_height


def rounded_image(image: Image.Image, radius: int) -> Image.Image:
    image = image.convert("RGBA")
    mask = Image.new("L", image.size, 0)
    draw = ImageDraw.Draw(mask)
    draw.rounded_rectangle((0, 0, image.width, image.height), radius=radius, fill=255)
    image.putalpha(mask)
    return image


def paste_with_shadow(
    canvas: Image.Image,
    screenshot: Image.Image,
    position: tuple[int, int],
    radius: int,
    glow: bool = False,
) -> None:
    x, y = position
    shadow_offset = max(10, round(canvas.width * 0.014))
    shadow_blur = max(18, round(canvas.width * 0.025))
    shadow_opacity = 105

    if glow:
        # Soft light halo so the device lifts off the dark background.
        halo_blur = max(28, round(canvas.width * 0.05))
        halo_pad = max(24, round(canvas.width * 0.03))
        halo_size = (screenshot.width + halo_pad * 2, screenshot.height + halo_pad * 2)
        halo = Image.new("RGBA", halo_size, (0, 0, 0, 0))
        halo_mask = Image.new("L", halo_size, 0)
        ImageDraw.Draw(halo_mask).rounded_rectangle(
            (halo_pad, halo_pad, halo_pad + screenshot.width, halo_pad + screenshot.height),
            radius=radius,
            fill=70,
        )
        halo.putalpha(halo_mask)
        # Tint the halo toward a cool light so it reads as premium elevation.
        halo_rgb = Image.new("RGBA", halo_size, (148, 197, 255, 0))
        halo_rgb.putalpha(halo.getchannel("A"))
        halo_rgb = halo_rgb.filter(ImageFilter.GaussianBlur(halo_blur))
        canvas.alpha_composite(halo_rgb, (x - halo_pad, y - halo_pad))

    shadow = Image.new("RGBA", screenshot.size, (0, 0, 0, 0))
    mask = Image.new("L", screenshot.size, 0)
    ImageDraw.Draw(mask).rounded_rectangle((0, 0, screenshot.width, screenshot.height), radius=radius, fill=shadow_opacity)
    shadow.putalpha(mask)
    shadow = shadow.filter(ImageFilter.GaussianBlur(shadow_blur))

    canvas.alpha_composite(shadow, (x, y + shadow_offset))
    canvas.alpha_composite(rounded_image(screenshot, radius), position)


def resize_screenshot(
    image: Image.Image,
    canvas_width: int,
    canvas_height: int,
    width_frac: float = 0.78,
    max_height_px: int | None = None,
) -> Image.Image:
    target_width = round(canvas_width * width_frac)
    target_height = round(canvas_height * 0.72)
    if max_height_px is not None:
        target_height = min(target_height, max_height_px)
    scale = min(target_width / image.width, target_height / image.height, 1.0)
    new_size = (max(1, round(image.width * scale)), max(1, round(image.height * scale)))
    if new_size == image.size:
        return image.convert("RGBA")
    return image.convert("RGBA").resize(new_size, Image.Resampling.LANCZOS)


def frame_screenshot(
    input_path: Path,
    output_path: Path,
    width: int,
    height: int,
    headline: str,
    grad_top: str,
    grad_bottom: str,
    device: str | None = None,
    theme: str = "light",
    stat: str | None = None,
    substat: str | None = None,
) -> None:
    del device  # Reserved for future per-device tuning.
    if theme == "dark":
        frame_screenshot_dark(
            input_path, output_path, width, height,
            headline, grad_top, grad_bottom, stat, substat,
        )
        return

    canvas = make_vertical_gradient((width, height), grad_top, grad_bottom)
    draw = ImageDraw.Draw(canvas)

    horizontal_margin = round(width * 0.08)
    top_margin = round(height * 0.052)
    headline_area_height = round(height * 0.14)
    max_text_width = width - horizontal_margin * 2
    wrapped, font, _, text_height = fit_headline(
        headline,
        (width, height),
        max_text_width,
        headline_area_height,
        draw,
    )
    headline_y = top_margin + max(0, (headline_area_height - text_height) // 2)
    draw.multiline_text(
        (width // 2, headline_y),
        wrapped,
        font=font,
        fill=(255, 255, 255, 255),
        anchor="ma",
        align="center",
        spacing=max(8, round(height * 0.006)),
    )

    screenshot = resize_screenshot(Image.open(input_path), width, height)
    screenshot_x = (width - screenshot.width) // 2
    min_screenshot_y = round(height * 0.22)
    screenshot_y = max(top_margin + headline_area_height + round(height * 0.035), min_screenshot_y)
    available_bottom = height - round(height * 0.045)
    if screenshot_y + screenshot.height > available_bottom:
        screenshot_y = max(round(height * 0.19), available_bottom - screenshot.height)

    radius = max(24, round(width * 0.03))
    paste_with_shadow(canvas, screenshot, (screenshot_x, screenshot_y), radius)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    canvas.convert("RGB").save(output_path, "PNG")


def frame_screenshot_dark(
    input_path: Path,
    output_path: Path,
    width: int,
    height: int,
    headline: str,
    grad_top: str,
    grad_bottom: str,
    stat: str | None,
    substat: str | None,
) -> None:
    """Dark premium / fintech framing: near-black gradient, near-white bold
    headline, optional big green stat callout, and a device shot lifted off the
    background with a subtle light glow."""
    canvas = make_vertical_gradient((width, height), grad_top, grad_bottom)
    draw = ImageDraw.Draw(canvas)

    horizontal_margin = round(width * 0.08)
    top_margin = round(height * 0.055)
    max_text_width = width - horizontal_margin * 2
    line_spacing = max(8, round(height * 0.006))

    # --- Headline: bold, near-white, slightly tighter/smaller than the light
    # theme so it reads premium rather than shouty.
    headline_area_height = round(height * 0.11)
    wrapped, headline_font, _, headline_h = fit_headline(
        headline,
        (width, height),
        max_text_width,
        headline_area_height,
        draw,
        max_scale=0.044,
        font_loader=load_bold_font,
    )

    cursor_y = top_margin
    draw.multiline_text(
        (width // 2, cursor_y),
        wrapped,
        font=headline_font,
        fill=HEADLINE_COLOR + (255,),
        anchor="ma",
        align="center",
        spacing=line_spacing,
    )
    cursor_y += headline_h

    # --- Optional stat callout: big vibrant-green number + muted secondary.
    if stat:
        cursor_y += round(height * 0.022)
        stat_size = max(48, round(height * 0.062))
        stat_font = load_bold_font(stat_size)
        # Shrink to fit width if the number is long.
        while stat_size > 40:
            if text_bbox(draw, stat, stat_font)[0] <= max_text_width:
                break
            stat_size -= 2
            stat_font = load_bold_font(stat_size)
        _, stat_h = text_bbox(draw, stat, stat_font)
        draw.text(
            (width // 2, cursor_y),
            stat,
            font=stat_font,
            fill=ACCENT_COLOR + (255,),
            anchor="ma",
        )
        cursor_y += stat_h

        if substat:
            cursor_y += round(height * 0.010)
            sub_size = max(26, round(height * 0.026))
            sub_font = load_regular_font(sub_size)
            sub_wrapped = wrap_headline(substat, sub_font, max_text_width, draw)
            _, sub_h = text_bbox(draw, sub_wrapped, sub_font)
            # A leading "▲" (or up-arrow) is a positive signal -> green.
            sub_color = ACCENT_COLOR if substat.strip().startswith(("▲", "↑", "+")) else MUTED_COLOR
            draw.multiline_text(
                (width // 2, cursor_y),
                sub_wrapped,
                font=sub_font,
                fill=sub_color + (255,),
                anchor="ma",
                align="center",
                spacing=line_spacing,
            )
            cursor_y += sub_h
    elif substat:
        # substat without a big stat -> supporting line. Positive-signal
        # prefixes (▲ ↑ +) get the green accent; otherwise muted slate.
        cursor_y += round(height * 0.014)
        sub_size = max(26, round(height * 0.030))
        sub_font = load_bold_font(sub_size)
        sub_wrapped = wrap_headline(substat, sub_font, max_text_width, draw)
        _, sub_h = text_bbox(draw, sub_wrapped, sub_font)
        sub_color = ACCENT_COLOR if substat.strip().startswith(("▲", "↑", "+")) else MUTED_COLOR
        draw.multiline_text(
            (width // 2, cursor_y),
            sub_wrapped,
            font=sub_font,
            fill=sub_color + (255,),
            anchor="ma",
            align="center",
            spacing=line_spacing,
        )
        cursor_y += sub_h

    # --- Device shot below the text block, with glow + shadow.
    # The device top is pinned just under the text block; it is scaled down so
    # it always fits the remaining space and never overlaps the copy.
    screenshot_y = cursor_y + round(height * 0.05)
    available_bottom = height - round(height * 0.045)
    max_device_height = available_bottom - screenshot_y
    screenshot = resize_screenshot(
        Image.open(input_path), width, height, max_height_px=max_device_height
    )
    screenshot_x = (width - screenshot.width) // 2
    # Center the device within the space left below the text for balance.
    slack = max(0, (available_bottom - screenshot_y) - screenshot.height)
    screenshot_y += slack // 2

    radius = max(24, round(width * 0.03))
    paste_with_shadow(canvas, screenshot, (screenshot_x, screenshot_y), radius, glow=True)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    canvas.convert("RGB").save(output_path, "PNG")


def positive_int(value: str) -> int:
    number = int(value)
    if number <= 0:
        raise argparse.ArgumentTypeError("Value must be greater than zero")
    return number


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Create gradient App Store marketing screenshots from raw PNG captures.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent(
            """\
            Example:
              python Scripts/aso/frame_screenshot.py \\
                --input raw.png --output framed.png --width 1320 --height 2868 \\
                --headline "Watch your wealth grow" --grad-top "#3B82F6" --grad-bottom "#10B981"
            """
        ),
    )
    parser.add_argument("--input", required=True, type=Path, help="Raw app screenshot PNG.")
    parser.add_argument("--output", required=True, type=Path, help="Framed output PNG.")
    parser.add_argument("--width", required=True, type=positive_int, help="Final canvas width.")
    parser.add_argument("--height", required=True, type=positive_int, help="Final canvas height.")
    parser.add_argument("--headline", required=True, help="Localized headline text.")
    parser.add_argument(
        "--grad-top",
        help="Top gradient color as #RRGGBB (defaults per theme).",
    )
    parser.add_argument(
        "--grad-bottom",
        help="Bottom gradient color as #RRGGBB (defaults per theme).",
    )
    parser.add_argument(
        "--theme",
        choices=("light", "dark"),
        default="light",
        help="Framing theme. 'dark' = dark premium / fintech look.",
    )
    parser.add_argument(
        "--stat",
        help="Optional big accent stat callout, e.g. '€717,382' (dark theme).",
    )
    parser.add_argument(
        "--substat",
        help="Optional smaller secondary line, e.g. '▲ +9.4% since last check-in'.",
    )
    parser.add_argument("--device", choices=("iphone", "ipad"), help="Optional device family hint.")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.theme == "dark":
        grad_top = args.grad_top or DARK_GRAD_TOP
        grad_bottom = args.grad_bottom or DARK_GRAD_BOTTOM
    else:
        grad_top = args.grad_top or DEFAULT_GRAD_TOP
        grad_bottom = args.grad_bottom or DEFAULT_GRAD_BOTTOM

    parse_hex_color(grad_top)
    parse_hex_color(grad_bottom)
    print_brand_colors(find_repo_root(Path(__file__).resolve()))

    frame_screenshot(
        input_path=args.input,
        output_path=args.output,
        width=args.width,
        height=args.height,
        headline=args.headline,
        grad_top=grad_top,
        grad_bottom=grad_bottom,
        device=args.device,
        theme=args.theme,
        stat=args.stat,
        substat=args.substat,
    )
    print(f"Wrote {args.output} ({args.width}x{args.height})")
    return 0


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