#!/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")


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",
)


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) -> ImageFont.ImageFont:
    for font_path in FONT_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 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,
) -> tuple[str, ImageFont.ImageFont, int, int]:
    _, height = canvas_size
    max_size = max(34, round(height * 0.052))
    min_size = 28

    for size in range(max_size, min_size - 1, -2):
        font = load_font(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 = load_font(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,
) -> 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

    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) -> Image.Image:
    target_width = round(canvas_width * 0.78)
    target_height = round(canvas_height * 0.72)
    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,
) -> None:
    del device  # Reserved for future per-device tuning.
    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 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", required=True, help="Top gradient color as #RRGGBB.")
    parser.add_argument("--grad-bottom", required=True, help="Bottom gradient color as #RRGGBB.")
    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)

    parse_hex_color(args.grad_top)
    parse_hex_color(args.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=args.grad_top,
        grad_bottom=args.grad_bottom,
        device=args.device,
    )
    print(f"Wrote {args.output} ({args.width}x{args.height})")
    return 0


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