frames

#!/usr/bin/env python3
"""frames — Apple Frames CLI

Frame device screenshots with Apple product bezels. Auto-detects devices from
screenshot dimensions, applies colored frames with masks, merges multiple
screenshots side by side. Replicates Apple Frames shortcut functionality for
use in terminal workflows and AI agent pipelines.

Assets default to the Apple Frames Shortcuts location on macOS and standard
app data directories elsewhere.
"""

import argparse
import json
import os
import random
import shutil
import ssl
import stat
import subprocess
import sys
import tempfile
import warnings
from functools import lru_cache
from pathlib import Path, PurePosixPath

try:
    from PIL import Image, UnidentifiedImageError
except ImportError:
    print("Error: Pillow is required. Install with: pip3 install Pillow", file=sys.stderr)
    sys.exit(1)

warnings.simplefilter("error", Image.DecompressionBombWarning)

VERSION = "1.2.7"
APP_NAME = "frames"
HOME = Path.home()
CONFIG_DIR = HOME / ".config" / "frames"
CONFIG_FILE = CONFIG_DIR / "config.json"
ASSETS_URL = "https://cdn.macstories.net/AppleFrames40.zip"
ASSETS_MIN_VERSION = 4
LARGE_IMAGE_WARN_DIMENSION = 20000
LARGE_IMAGE_MAX_DIMENSION = 50000
VALID_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".heic", ".tiff", ".webp"}


def get_shortcuts_assets_dir(home=None):
    """Return the Apple Frames Shortcuts asset location on macOS."""
    home = Path(home) if home is not None else HOME
    return home / "Library/Mobile Documents/iCloud~is~workflow~my~workflows/Documents/Frames"


def get_app_data_assets_dir(home=None, platform=None, environ=None):
    """Return the standard per-platform app data assets directory."""
    home = Path(home) if home is not None else HOME
    platform = platform or sys.platform
    environ = environ or os.environ

    if platform == "darwin":
        return Path(environ.get("FRAMES_DATA_HOME", home / "Library/Application Support")) / APP_NAME / "assets"
    if platform.startswith("win"):
        return Path(environ.get("LOCALAPPDATA", home / "AppData/Local")) / APP_NAME / "assets"
    return Path(environ.get("XDG_DATA_HOME", home / ".local/share")) / APP_NAME / "assets"


def get_default_assets_dir(home=None, platform=None, environ=None):
    """Return the preferred default assets directory for the current platform."""
    platform = platform or sys.platform
    if platform == "darwin":
        return get_shortcuts_assets_dir(home)
    return get_app_data_assets_dir(home, platform, environ)


SHORTCUTS_ASSETS = get_shortcuts_assets_dir()
APPDATA_ASSETS = get_app_data_assets_dir()
DEFAULT_ASSETS = get_default_assets_dir()


def _expand_path(path_value):
    return Path(path_value).expanduser()


def _get_config_assets_dir():
    if CONFIG_FILE.exists():
        try:
            with open(CONFIG_FILE, encoding="utf-8") as f:
                cfg = json.load(f)
            if "assets_path" in cfg:
                return _expand_path(cfg["assets_path"])
        except (json.JSONDecodeError, KeyError, OSError):
            pass
    return None


def read_config_file():
    """Read config file without swallowing errors."""
    if not CONFIG_FILE.exists():
        return None, None
    try:
        with open(CONFIG_FILE, encoding="utf-8") as f:
            return json.load(f), None
    except (json.JSONDecodeError, OSError) as e:
        return None, str(e)


def validate_subfolder_name(name):
    """Allow only a single relative directory component for --subfolder."""
    candidate = str(name).strip()
    if not candidate:
        raise ValueError("Subfolder name cannot be empty.")

    pure = PurePosixPath(candidate.replace("\\", "/"))
    if pure.is_absolute() or len(pure.parts) != 1 or pure.parts[0] in {".", ".."}:
        raise ValueError(
            "Subfolder must be a single directory name without path separators or '..'."
        )

    return pure.parts[0]


def resolve_assets_dir(args_assets=None):
    """Resolve assets directory and report which source won."""
    if args_assets is not None:
        return _expand_path(args_assets), "--assets"
    env = os.environ.get("FRAMES_ASSETS")
    if env:
        return _expand_path(env), "FRAMES_ASSETS"
    configured = _get_config_assets_dir()
    if configured is not None:
        return configured, "config"
    if check_assets_version(DEFAULT_ASSETS)[0]:
        return DEFAULT_ASSETS, "default"
    if DEFAULT_ASSETS != APPDATA_ASSETS and check_assets_version(APPDATA_ASSETS)[0]:
        return APPDATA_ASSETS, "app-data"
    return DEFAULT_ASSETS, "default"


def get_assets_dir(args_assets=None):
    """Resolve assets directory: CLI flag > env var > config file > default."""
    return resolve_assets_dir(args_assets)[0]

def get_subfolder_setting(args, config=None):
    """Get subfolder name if enabled. Returns folder name string or None.
    CLI flag > config > default (None). -f alone = 'framed', -f custom = 'custom'."""
    if hasattr(args, 'subfolder') and args.subfolder is not None:
        return validate_subfolder_name(args.subfolder)
    cfg = config if config is not None else load_config()
    sub = cfg.get("use_subfolder", False)
    if sub is True:
        return "framed"
    elif isinstance(sub, str):
        return validate_subfolder_name(sub)
    return None

# ── Asset Management ────────────────────────────────────────────────────────

def check_assets_version(assets_dir):
    """Check if assets directory has valid Apple Frames 4+ assets.
    Returns (ok, version_or_none)."""
    assets_dir = Path(assets_dir)
    json_path = assets_dir / "NewFrames.json"

    if not assets_dir.exists() or not json_path.exists():
        return False, None

    version_path = assets_dir / "version.txt"
    if version_path.exists():
        try:
            ver = version_path.read_text(encoding="utf-8").strip().split('\n')[0].strip()
            major = int(ver.split('.')[0])
            return major >= ASSETS_MIN_VERSION, major
        except (ValueError, IndexError):
            pass

    # No version.txt — check for v4 indicators in JSON
    try:
        with open(json_path, encoding="utf-8") as f:
            data = json.load(f)
        if "variants" in data:
            return True, 4
        return False, None
    except (json.JSONDecodeError, IOError):
        return False, None


def _zip_member_is_symlink(member):
    unix_mode = member.external_attr >> 16
    return stat.S_ISLNK(unix_mode)


def _iter_safe_zip_members(zip_file):
    members = []
    top_items = set()

    for member in zip_file.infolist():
        filename = member.filename.replace("\\", "/")
        if not filename or filename.startswith("__MACOSX/"):
            continue
        members.append((member, filename))
        top_items.add(filename.split("/", 1)[0])

    prefix = ""
    if len(top_items) == 1:
        prefix = next(iter(top_items)) + "/"

    for member, filename in members:
        relative_name = filename
        if prefix and relative_name.startswith(prefix):
            relative_name = relative_name[len(prefix):]
        if not relative_name:
            continue

        relative_path = PurePosixPath(relative_name)
        if relative_path.is_absolute() or ".." in relative_path.parts:
            raise ValueError(f"Archive contains unsafe path: {filename}")
        if _zip_member_is_symlink(member):
            raise ValueError(f"Archive contains unsupported symlink: {filename}")

        yield member, relative_path


def _extract_zip_to_dir(zip_file, destination):
    for member, relative_path in _iter_safe_zip_members(zip_file):
        target = destination.joinpath(*relative_path.parts)
        if member.is_dir():
            target.mkdir(parents=True, exist_ok=True)
            continue

        target.parent.mkdir(parents=True, exist_ok=True)
        with zip_file.open(member) as src, open(target, "wb") as dst:
            shutil.copyfileobj(src, dst)


def _find_assets_root(path):
    json_path = path / "NewFrames.json"
    if json_path.exists():
        return path
    found = list(path.rglob("NewFrames.json"))
    if found:
        return found[0].parent
    return None

def _is_ssl_download_error(exc):
    reason = getattr(exc, "reason", exc)
    if isinstance(reason, ssl.SSLError):
        return True

    text = str(reason).lower()
    ssl_markers = (
        "certificate verify failed",
        "self-signed certificate",
        "tlsv1 alert",
        "ssl:",
    )
    return any(marker in text for marker in ssl_markers)

def _download_with_urllib(tmp_path, reporthook):
    import urllib.request

    urllib.request.urlretrieve(ASSETS_URL, tmp_path, reporthook=reporthook)

def _download_with_curl(tmp_path, curl_path=None):
    curl_path = curl_path or shutil.which("curl")
    if curl_path is None:
        raise FileNotFoundError("curl is not installed")

    result = subprocess.run(
        [curl_path, "--location", "--fail", "--progress-bar", "--show-error", "-o", tmp_path, ASSETS_URL],
        check=False,
    )
    if result.returncode != 0:
        raise OSError(f"curl download failed with exit code {result.returncode}")

def download_assets(dest_dir):
    """Download and extract Apple Frames 4 assets. Returns True on success."""
    import urllib.request, urllib.error, zipfile

    dest_dir = Path(dest_dir)

    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
        tmp_path = tmp.name

    try:
        def progress(block_num, block_size, total_size):
            if total_size > 0:
                pct = min(100, int(block_num * block_size * 100 / total_size))
                filled = int(25 * pct / 100)
                bar = '\u2588' * filled + '\u2591' * (25 - filled)
                mb = total_size / (1024 * 1024)
                print(f"\r  Downloading [{C.rgb(232,114,42, bar)}] {pct}% of {mb:.0f} MB ", end="", flush=True)

        print(f"  {C.dim('Connecting to cdn.macstories.net...')}", flush=True)
        try:
            _download_with_urllib(tmp_path, progress)
            print()
        except urllib.error.URLError as e:
            if not _is_ssl_download_error(e):
                raise

            curl_path = shutil.which("curl")
            if curl_path is None:
                raise

            print()
            print(f"  {C.dim('Python SSL verification failed; retrying with curl...')}", flush=True)
            _download_with_curl(tmp_path, curl_path=curl_path)
            print()

        print(f"  {C.dim('Extracting...')}", end=" ", flush=True)
        with tempfile.TemporaryDirectory(prefix="frames-assets-") as stage_dir:
            stage_path = Path(stage_dir)
            with zipfile.ZipFile(tmp_path, "r") as zf:
                _extract_zip_to_dir(zf, stage_path)

            source_root = _find_assets_root(stage_path)
            if source_root is None:
                print(C.red('failed'))
                print(f"  {C.red('Error')}: NewFrames.json not found in archive.", file=sys.stderr)
                return False

            dest_dir.mkdir(parents=True, exist_ok=True)
            for item in source_root.iterdir():
                target = dest_dir / item.name
                if item.is_dir():
                    shutil.copytree(item, target, dirs_exist_ok=True)
                else:
                    shutil.copy2(item, target)

        pngs = list(dest_dir.glob("*.png"))
        print(C.green('done!'))
        print(f"  {C.bold(str(len(pngs)))} frame assets installed")
        return True

    except urllib.error.URLError as e:
        print(f"\n  {C.red('Download failed')}: {e.reason}", file=sys.stderr)
        return False
    except zipfile.BadZipFile:
        print(f"\n  {C.red('Error')}: Downloaded file is corrupted.", file=sys.stderr)
        return False
    except ValueError as e:
        print(C.red("failed"))
        print(f"  {C.red('Error')}: {e}", file=sys.stderr)
        return False
    except OSError as e:
        print(f"\n  {C.red('Error')}: {e}", file=sys.stderr)
        return False
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass

def guided_setup(reason=None):
    """Interactive guided setup — download and install Apple Frames 4 assets.
    Returns True on success, False if cancelled or failed."""
    default_path = get_default_assets_dir()
    short_default = str(default_path).replace(str(HOME), "~")

    print()
    if reason:
        print(f"  {C.bold('Apple Frames 4')} assets {reason}.")
    else:
        print(f"  {C.bold('Apple Frames 4')} asset setup")
    print()
    print(f"  Download device frames from {C.cyan('cdn.macstories.net')}? (~40 MB)")
    print()
    print(f"  Install to {C.cyan(short_default)}")
    print()
    print(f"  {C.bold('Enter')} {C.dim('to confirm  |')}  {C.bold('q')} {C.dim('to cancel  |  or type a different path')}")
    print()

    try:
        choice = input(f"  {C.bold('›')} ").strip()
    except (EOFError, KeyboardInterrupt):
        print(f"\n\n  {C.dim('Cancelled.')}\n")
        return False

    if choice.lower() in ('q', 'quit', 'n', 'no'):
        print(f"\n  {C.dim('Cancelled.')}\n")
        return False

    dest = Path(choice).expanduser().resolve() if choice else default_path

    print()
    if not download_assets(dest):
        return False

    config = load_config()
    config["assets_path"] = str(dest)
    save_config(config)

    short_dest = str(dest).replace(str(HOME), "~")
    print()
    print(f"  {C.green('Setup complete!')}")
    print(f"  {C.bold('Assets:')}  {short_dest}")
    print(f"  {C.bold('Config:')}  {CONFIG_FILE}")
    print()
    print(f"  {C.dim('Run')} {C.cyan('frames')} screenshot.png {C.dim('to get started.')}")
    print()
    return True

# ── Color System ─────────────────────────────────────────────────────────────

class C:
    """ANSI color helpers. Respects NO_COLOR env and --no-color flag."""
    enabled = True

    @staticmethod
    def _c(code, text):
        return f"\033[{code}m{text}\033[0m" if C.enabled else str(text)

    @staticmethod
    def red(t):     return C._c("31", t)
    @staticmethod
    def green(t):   return C._c("32", t)
    @staticmethod
    def yellow(t):  return C._c("33", t)
    @staticmethod
    def blue(t):    return C._c("34", t)
    @staticmethod
    def magenta(t): return C._c("35", t)
    @staticmethod
    def cyan(t):    return C._c("36", t)
    @staticmethod
    def dim(t):     return C._c("2", t)
    @staticmethod
    def bold(t):    return C._c("1", t)
    @staticmethod
    def rgb(r, g, b, t):
        return f"\033[38;2;{r};{g};{b}m{t}\033[0m" if C.enabled else str(t)

# ── Splash Screen ────────────────────────────────────────────────────────────

LOGO_LINES = [
    r" █████  ██████  ██████  ██      ███████ ",
    r"██   ██ ██   ██ ██   ██ ██      ██      ",
    r"███████ ██████  ██████  ██      █████   ",
    r"██   ██ ██      ██      ██      ██      ",
    r"██   ██ ██      ██      ███████ ███████ ",
    r"",
    r"███████ ██████   █████  ███    ███ ███████ ███████ ",
    r"██      ██   ██ ██   ██ ████  ████ ██      ██      ",
    r"█████   ██████  ███████ ██ ████ ██ █████   ███████ ",
    r"██      ██   ██ ██   ██ ██  ██  ██ ██           ██ ",
    r"██      ██   ██ ██   ██ ██      ██ ███████ ███████ ",
]

LOGO_COLORS = [
    (232, 114, 42),   # Cosmic Orange
    (235, 125, 50),
    (238, 136, 58),
    (235, 125, 50),
    (232, 114, 42),
    (0, 0, 0),
    (232, 114, 42),
    (235, 125, 50),
    (238, 136, 58),
    (235, 125, 50),
    (232, 114, 42),
]

def show_splash():
    """Display colored ASCII art splash screen."""
    print()
    for line, (r, g, b) in zip(LOGO_LINES, LOGO_COLORS):
        if line:
            print(f"  {C.rgb(r, g, b, line)}")
        else:
            print()

    ver = C.dim(f"v{VERSION}")
    author = C.dim("by Federico Viticci")
    print(f"\n  {C.bold('Apple Frames CLI')} {ver}  {author}")
    print(f"  {C.dim('Frame screenshots with Apple device bezels')}\n")
    print(f"  {C.dim('Quick start:')}")
    print(f"    {C.cyan('frames')} screenshot.png         {C.dim('Frame a screenshot')}")
    print(f"    {C.cyan('frames')} *.png                  {C.dim('Frame all PNGs in directory')}")
    print(f"    {C.cyan('frames')} -m *.png               {C.dim('Merge side by side (proportional)')}")
    print(f"    {C.cyan('frames')} -b 3 *.png             {C.dim('Merge in batches of 3')}")
    print(f"    {C.cyan('frames')} -m --no-scale *.png    {C.dim('Merge without proportional scaling')}")
    print(f"    {C.cyan('frames')} -c random *.png        {C.dim('Random frame colors')}")
    print(f"    {C.cyan('frames')} ~/Screenshots/          {C.dim('Frame all images in a folder')}")
    print(f"    {C.cyan('frames')} -f *.png               {C.dim('Save to /framed/ subfolder')}")
    print(f"    {C.cyan('frames')} --subfolder mockups *.png {C.dim('Save to custom subfolder')}")
    print(f"    {C.cyan('frames')} info shot.png           {C.dim('Identify device, colors, variants')}")
    print(f"    {C.cyan('frames')} info ~/Screenshots/     {C.dim('Scan folder for device matches')}")
    print(f"    {C.cyan('frames')} colors                 {C.dim('Pick default colors per device')}")
    print(f"    {C.cyan('frames')} list                   {C.dim('Show supported devices')}")
    print(f"    {C.cyan('frames')} setup                  {C.dim('Download + configure assets')}")

    # Show asset status
    assets_dir = get_assets_dir()
    ok, ver = check_assets_version(assets_dir)
    if ok:
        short = str(assets_dir).replace(str(HOME), "~")
        pngs = len(list(assets_dir.glob("*.png")))
        print(f"  {C.green('Assets installed')} {C.dim(f'(v{ver}, {pngs} frames)')}")
        print(f"  {C.dim(short)}")
    else:
        print(f"  {C.yellow('Assets not installed.')} Run {C.cyan('frames setup')} to download device frames.")
    print()

# ── Frame Colors Dictionary ──────────────────────────────────────────────────

FRAME_COLORS = {
    "iPhone 17 Pro Portrait":      ["Cosmic Orange", "Deep Blue", "Silver"],
    "iPhone 17 Pro Landscape":     ["Cosmic Orange", "Deep Blue", "Silver"],
    "iPhone 17 Pro Max Portrait":  ["Cosmic Orange", "Deep Blue", "Silver"],
    "iPhone 17 Pro Max Landscape": ["Cosmic Orange", "Deep Blue", "Silver"],
    "iPhone 17 Portrait":          ["Black", "Lavender", "Mist Blue", "Sage", "White"],
    "iPhone 17 Landscape":         ["Black", "Lavender", "Mist Blue", "Sage", "White"],
    "iPhone Air Portrait":         ["Black", "White", "Gold", "Blue"],
    "iPhone Air Landscape":        ["Black", "White", "Gold", "Blue"],
    "MacBook Pro M5 14":           ["Silver", "Space Black"],
    "MacBook Pro M5 16":           ["Silver", "Space Black"],
    "MacBook Neo":                 ["Blush", "Citrus", "Indigo", "Silver"],
    "iPhone 16 Pro Portrait":      ["Natural", "Desert", "Black", "White"],
    "iPhone 16 Pro Landscape":     ["Natural", "Desert", "Black", "White"],
    "iPhone 16 Pro Max Portrait":  ["White", "Desert", "Natural", "Black"],
    "iPhone 16 Pro Max Landscape": ["Black", "Desert", "Natural", "White"],
    "iPhone 16 Portrait":          ["Ultramarine", "Teal", "Pink", "Black", "White"],
    "iPhone 16 Landscape":         ["Ultramarine", "Teal", "Pink", "Black", "White"],
    "iPhone 16 Plus Portrait":     ["Ultramarine", "White", "Teal", "Pink", "Black"],
    "iPhone 16 Plus Landscape":    ["Ultramarine", "Teal", "Pink", "Black", "White"],
    "Watch Ultra 3":               ["Black + Alpine Loop Black", "Black + Alpine Loop Light Blue", "Black + Milanese Loop", "Black + Ocean Band Anchor Blue", "Black + Ocean Band Black", "Black + Trail Loop Black Charcoal", "Natural + Alpine Loop Light Blue", "Natural + Alpine Loop Terra Cotta", "Natural + Milanese Loop", "Natural + Ocean Band Anchor Blue", "Natural + Ocean Band Neon Green", "Natural + Trail Loop Blue Bright Blue", "Natural + Trail Loop Green Neon"],
    "Watch Ultra 2024":            ["Orange Beige Trail Loop", "Blue Alpine Loop", "Orange Ocean Band", "Black + Alpine Loop Dark Green", "Black + Alpine Loop Navy", "Black + Alpine Loop Tan", "Black + Ocean Band Black", "Black + Ocean Band Ice Blue", "Black + Ocean Band Navy", "Black + Titanium Milanese Loop", "Black + Trail Loop Black", "Black + Trail Loop Blue", "Black + Trail Loop Green", "Natural + Alpine Loop Dark Green", "Natural + Alpine Loop Navy", "Natural + Alpine Loop Tan", "Natural + Ocean Band Black", "Natural + Ocean Band Ice Blue", "Natural + Ocean Band Navy", "Natural + Titanium Milanese Loop", "Natural + Trail Loop Black", "Natural + Trail Loop Blue"],
    "MacBook Air M5 13":           ["Midnight", "Silver", "Sky Blue", "Starlight"],
    "MacBook Air M5 15":           ["Midnight", "Silver", "Sky Blue", "Starlight"],
    "iMac M4":                     ["Silver", "Blue", "Green", "Orange", "Pink", "Purple", "Yellow"],
    "Studio Display":              ["Light", "Dark"],
    "Studio Display XDR":          ["Light", "Dark"],
    "Watch Series 11 42":           ["Aluminum Jet Black + Sport Band Black", "Aluminum Jet Black + Sport Loop Dark Gray", "Aluminum Rose Gold + Sport Band Light Blush", "Aluminum Rose Gold + Sport Loop Purple Fog", "Aluminum Silver + Sport Band Neon Yellow", "Aluminum Silver + Sport Band Purple Fog", "Aluminum Silver + Sport Loop Forest", "Aluminum Silver + Sport Loop Neon Yellow", "Aluminum Space Gray + Sport Band Anchor Blue", "Aluminum Space Gray + Sport Band Black", "Aluminum Space Gray + Sport Loop Anchor Blue", "Aluminum Space Gray + Sport Loop Forest", "Titanium Gold + Magnetic Link Sage Gray", "Titanium Gold + Milanese Loop", "Titanium Gold + Sport Band Light Blush", "Titanium Gold + Sport Band Purple Fog", "Titanium Natural + Magnetic Link Caramel", "Titanium Natural + Milanese Loop", "Titanium Natural + Sport Band Stone Gray", "Titanium Slate + Magnetic Link Navy", "Titanium Slate + Milanese Loop", "Titanium Slate + Sport Band Black"],
    "Watch Series 11 46":           ["Aluminum Jet Black + Sport Band Black", "Aluminum Jet Black + Sport Loop Dark Gray", "Aluminum Rose Gold + Sport Band Light Blush", "Aluminum Rose Gold + Sport Loop Purple Fog", "Aluminum Silver + Sport Band Neon Yellow", "Aluminum Silver + Sport Band Purple Fog", "Aluminum Silver + Sport Loop Forest", "Aluminum Silver + Sport Loop Neon Yellow", "Aluminum Space Gray + Sport Band Anchor Blue", "Aluminum Space Gray + Sport Band Black", "Aluminum Space Gray + Sport Loop Anchor Blue", "Aluminum Space Gray + Sport Loop Forest", "Titanium Gold + Magnetic Link Sage Gray", "Titanium Gold + Milanese Loop", "Titanium Gold + Sport Band Light Blush", "Titanium Gold + Sport Band Purple Fog", "Titanium Natural + Magnetic Link Caramel", "Titanium Natural + Milanese Loop", "Titanium Natural + Sport Band Stone Gray", "Titanium Slate + Magnetic Link Navy", "Titanium Slate + Milanese Loop", "Titanium Slate + Sport Band Black"],
}

# Frame Variants: first item = default device
FRAME_VARIANTS = {
    "iPhone 17 Portrait":          ["iPhone 17 Pro Portrait", "iPhone 17 Portrait", "iPhone 16 Pro Portrait"],
    "iPhone 17 Landscape":         ["iPhone 17 Pro Landscape", "iPhone 17 Landscape", "iPhone 16 Pro Landscape"],
    "iPhone 17 Pro Max Portrait":  ["iPhone 17 Pro Max Portrait", "iPhone 16 Pro Max Portrait"],
    "iPhone 17 Pro Max Landscape": ["iPhone 17 Pro Max Landscape", "iPhone 16 Pro Max Landscape"],
    "iPhone 16 Portrait":          ["iPhone 16 Portrait", "iPhone 15 Pro Portrait"],
    "iPhone 16 Landscape":         ["iPhone 16 Landscape", "iPhone 15 Pro Landscape"],
    "iPhone 16 Plus Portrait":     ["iPhone 16 Plus Portrait", "iPhone 15 Pro Max Portrait"],
    "iPhone 16 Plus Landscape":    ["iPhone 16 Plus Landscape", "iPhone 15 Pro Max Landscape"],
    "MacBook Pro M5 14":           ["MacBook Pro M5 14", "MacBook Pro 2021 14"],
    "MacBook Pro M5 16":           ["MacBook Pro M5 16", "MacBook Pro 2021 16"],
    "MacBook Air M5 13":           ["MacBook Air M5 13", "MacBook Air 2022"],
    "Studio Display":              ["Studio Display", "Studio Display XDR"],
    "Watch Series 11 42":          ["Watch Series 11 42", "Watch Series 10 42"],
    "Watch Series 11 46":          ["Watch Series 11 46", "Watch Series 10 46"],
}

# ── Core Engine ──────────────────────────────────────────────────────────────

def load_json(assets_dir):
    """Load the device dictionary from NewFrames.json."""
    json_path = assets_dir / "NewFrames.json"
    if not json_path.exists():
        print(f"{C.red('Error')}: NewFrames.json not found at {json_path}", file=sys.stderr)
        sys.exit(1)
    try:
        with open(json_path, encoding="utf-8") as f:
            return json.load(f)
    except json.JSONDecodeError as e:
        print(f"{C.red('Error')}: Invalid JSON in {json_path}: {e}", file=sys.stderr)
        sys.exit(1)

def detect_device(img, device_dict):
    """Detect device from screenshot dimensions using width + overlap height check."""
    w, h = img.size
    width_key = str(w)

    if width_key not in device_dict:
        return None, None

    entry = device_dict[width_key]

    # Check overlap (height disambiguation)
    if "overlap" in entry:
        height_key = str(h)
        if height_key in entry["overlap"]:
            entry = entry["overlap"][height_key]
        else:
            return None, None

    return entry, entry["name"]

def resolve_variant(model_name, device_dict):
    """Apply variant resolution — first item in variant array is the default."""
    if model_name not in FRAME_VARIANTS:
        return None, model_name

    variants = FRAME_VARIANTS[model_name]
    default_variant = variants[0]

    if default_variant != model_name:
        # Load variant device data from JSON
        if "variants" in device_dict and default_variant in device_dict["variants"]:
            variant_data = device_dict["variants"][default_variant]
            return variant_data, variant_data["name"]

    return None, model_name

def load_config():
    """Load user config from ~/.config/frames/config.json."""
    if CONFIG_FILE.exists():
        try:
            with open(CONFIG_FILE, encoding="utf-8") as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            pass
    return {}

def save_config(config):
    """Save user config."""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", dir=CONFIG_DIR, delete=False, encoding="utf-8") as tmp:
        json.dump(config, tmp, indent=2)
        tmp.write("\n")
        tmp_path = Path(tmp.name)
    os.replace(tmp_path, CONFIG_FILE)


def get_color(model_name, color_arg, user_defaults=None):
    """Resolve frame color for a device. Checks: CLI arg > user config > first in list."""
    if model_name not in FRAME_COLORS:
        return None

    colors = FRAME_COLORS[model_name]
    if not colors:
        return None

    if color_arg == "random":
        return random.choice(colors)
    elif color_arg:
        for c in colors:
            if c.lower() == color_arg.lower():
                return c
        try:
            idx = int(color_arg) - 1
            if 0 <= idx < len(colors):
                return colors[idx]
        except ValueError:
            pass
        print(f"{C.yellow('Warning')}: Color '{color_arg}' not found for {model_name}. Using default.", file=sys.stderr)

    # Check user config for saved default
    user_defaults = user_defaults or {}
    # Strip orientation for lookup (user sets "iPhone 17 Pro" not "iPhone 17 Pro Portrait")
    base = model_name.replace(" Portrait", "").replace(" Landscape", "")
    if base in user_defaults:
        saved = user_defaults[base]
        if saved in colors:
            return saved

    return colors[0]  # Fallback = first color


@lru_cache(maxsize=512)
def _load_rgba_asset_cached(path_str):
    with Image.open(path_str) as image:
        return image.convert("RGBA")

def load_frame(assets_dir, model_name, color):
    """Load the frame PNG for a device + color."""
    if color:
        frame_path = assets_dir / f"{model_name} {color}.png"
        if not frame_path.exists():
            # Fall back to colorless
            frame_path = assets_dir / f"{model_name}.png"
    else:
        frame_path = assets_dir / f"{model_name}.png"

    if not frame_path.exists():
        return None, str(frame_path)

    return _load_rgba_asset_cached(str(frame_path)).copy(), str(frame_path)

def load_mask(assets_dir, model_name):
    """Load mask PNG if it exists."""
    mask_path = assets_dir / f"{model_name}_mask.png"
    if mask_path.exists():
        return _load_rgba_asset_cached(str(mask_path)).copy()
    return None

def apply_mask(screenshot, mask):
    """Apply mask to screenshot — black pixels in mask make screenshot transparent."""
    r, g, b, a = screenshot.split()
    mask_r = mask.split()[0]  # R channel: 0=clip, 255=keep
    return Image.merge("RGBA", (r, g, b, mask_r))

def composite(screenshot, frame, x, y):
    """Alpha-composite: place screenshot on canvas, composite frame on top.
    Produces fully opaque output at semi-transparent bezel inner edges."""
    canvas = Image.new("RGBA", frame.size, (0, 0, 0, 0))
    canvas.paste(screenshot, (x, y))
    return Image.alpha_composite(canvas, frame)

def merge_images(images, spacing=60, physical_heights=None):
    """Merge multiple images horizontally with spacing.

    When physical_heights is provided and values differ, images are scaled
    proportionally (tallest = 1.0) and bottom-aligned. Otherwise, original
    center-aligned behavior is preserved."""
    if len(images) == 1:
        return images[0]

    do_scale = (
        physical_heights is not None
        and len(physical_heights) == len(images)
        and len(set(physical_heights)) > 1
    )

    if do_scale:
        # Normalize all frames to the same pixels-per-mm ratio.
        # The device with the lowest px/mm (coarsest frame) keeps native size;
        # everything else scales down. This avoids upscaling artifacts.
        px_per_mm = [img.height / ph for img, ph in zip(images, physical_heights)]
        ref_ppm = min(px_per_mm)
        scaled = []
        for img, ph in zip(images, physical_heights):
            target_h = round(ph * ref_ppm)
            scale = target_h / img.height
            if scale < 0.999:
                new_w = max(1, round(img.width * scale))
                new_h = max(1, target_h)
                scaled.append(img.resize((new_w, new_h), Image.LANCZOS))
            else:
                scaled.append(img)
        images = scaled

    total_width = sum(img.width for img in images) + spacing * (len(images) - 1)
    max_height = max(img.height for img in images)

    merged = Image.new("RGBA", (total_width, max_height), (0, 0, 0, 0))
    x_offset = 0
    for img in images:
        if do_scale:
            y_offset = max_height - img.height  # Bottom-align
        else:
            y_offset = (max_height - img.height) // 2  # Center-align
        merged.paste(img, (x_offset, y_offset), img)
        x_offset += img.width + spacing

    return merged

def ensure_output_dir(path):
    """Create output directory or raise a clean write error."""
    try:
        path.mkdir(parents=True, exist_ok=True)
    except OSError as e:
        raise OSError(f"Could not create output directory {path}: {e}") from e

def save_png(image, path):
    """Save a PNG or raise a clean write error."""
    try:
        image.save(path, "PNG")
    except OSError as e:
        raise OSError(f"Could not save {path}: {e}") from e

def guard_image_size(width, height, source):
    """Warn on very large images and reject extreme dimensions."""
    largest = max(width, height)
    if largest > LARGE_IMAGE_MAX_DIMENSION:
        raise ValueError(
            f"Refusing to process {source}: image dimensions exceed "
            f"{LARGE_IMAGE_MAX_DIMENSION}px ({width}x{height})."
        )
    if largest > LARGE_IMAGE_WARN_DIMENSION:
        print(
            f"{C.yellow('Warning')}: Large image {source} ({width}x{height}) "
            f"may use significant memory.",
            file=sys.stderr,
        )


def build_device_name_index(device_dict):
    """Build a name→device-entry lookup for exact --device matches."""
    index = {}
    for key, val in device_dict.items():
        if key == "variants":
            continue
        sub_entries = list(val.get("overlap", {}).values()) or [val]
        for entry in sub_entries:
            index[entry["name"]] = entry
    for name, entry in device_dict.get("variants", {}).items():
        index[name] = entry
    return index


def frame_screenshot(
    img_path,
    assets_dir,
    device_dict,
    color_arg=None,
    force_device=None,
    verbose=False,
    user_defaults=None,
    device_index=None,
):
    """Frame a single screenshot. Returns (framed_image, device_info_dict) or raises."""
    with Image.open(img_path) as opened:
        w, h = opened.size
        guard_image_size(w, h, img_path)
        img = opened.convert("RGBA")

    # 1. Device detection
    if force_device:
        entry = (device_index or build_device_name_index(device_dict)).get(force_device)
        if not entry:
            raise ValueError(f"Device '{force_device}' not found in dictionary")
        model_name = entry["name"]
    else:
        entry, model_name = detect_device(img, device_dict)
        if not entry:
            raise ValueError(f"Unknown device for {w}x{h} screenshot. Use 'frames list' to see supported devices.")

    # 2. Variant resolution (skip auto-resolution when --device explicitly selects a frame)
    if not force_device:
        variant_data, model_name = resolve_variant(model_name, device_dict)
        if variant_data:
            entry = variant_data
            if verbose:
                print(f"  {C.dim('Variant resolved:')} {model_name}")

    # 3. Color
    has_colors = entry.get("colors") == "yes"
    color = get_color(model_name, color_arg, user_defaults=user_defaults) if has_colors else None

    # 4. Resize (for More Space resolutions)
    resize_width = entry.get("resizeWidth")
    if resize_width:
        resize_width = int(resize_width)
        ratio = resize_width / img.width
        new_h = int(img.height * ratio)
        img = img.resize((resize_width, new_h), Image.LANCZOS)
        if verbose:
            print(f"  {C.dim('Resized:')} {w}x{h} → {resize_width}x{new_h}")

    # 5. Load frame
    frame, frame_path = load_frame(assets_dir, model_name, color)
    if not frame:
        raise FileNotFoundError(f"Frame asset not found: {frame_path}")

    # 6. Apply mask
    has_mask = entry.get("mask") == "yes"
    mask_missing = False
    if has_mask:
        mask = load_mask(assets_dir, model_name)
        if mask:
            # Ensure mask matches screenshot dimensions
            if mask.size != img.size:
                mask = mask.resize(img.size, Image.LANCZOS)
            img = apply_mask(img, mask)
            if verbose:
                print(f"  {C.dim('Mask applied')}")
        else:
            mask_missing = True
            print(f"  {C.yellow('Warning')}: mask file missing for {model_name} "
                  f"(expected {model_name}_mask.png). Rendering without mask — "
                  f"corners may show screenshot bleed.", file=sys.stderr)

    # 7. Composite
    x, y = int(entry["x"]), int(entry["y"])
    result = composite(img, frame, x, y)

    info = {
        "source": str(img_path),
        "device": model_name,
        "color": color,
        "dimensions": f"{w}x{h}",
        "frame_size": f"{result.width}x{result.height}",
        "resized": bool(resize_width),
        "masked": has_mask and not mask_missing,
        "mask_missing": mask_missing,
        "physicalHeight": float(entry.get("physicalHeight", 0)),
    }

    return result, info

# ── CLI Commands ─────────────────────────────────────────────────────────────

def cmd_frame(args):
    """Frame one or more screenshots."""
    assets_dir = Path(args.assets) if hasattr(args, 'assets') and args.assets else get_assets_dir()
    device_dict = load_json(assets_dir)
    config = load_config()
    user_defaults = config.get("default_colors", {})
    device_index = build_device_name_index(device_dict)

    if not args.files:
        print(f"{C.red('Error')}: No input files specified.", file=sys.stderr)
        sys.exit(1)

    batch_size = getattr(args, 'batch', None)
    if batch_size is not None and batch_size < 2:
        print(f"{C.red('Error')}: --batch must be at least 2.", file=sys.stderr)
        sys.exit(1)
    if args.spacing < 0:
        print(f"{C.red('Error')}: --spacing must be 0 or greater.", file=sys.stderr)
        sys.exit(1)

    # Expand directories into image files
    expanded = _gather_image_paths(args.files)
    if not expanded:
        print(f"{C.red('Error')}: No images found in the specified paths.", file=sys.stderr)
        sys.exit(1)

    try:
        subfolder = get_subfolder_setting(args, config=config)
    except ValueError as e:
        print(f"{C.red('Error')}: {e}", file=sys.stderr)
        sys.exit(1)

    framed = []
    infos = []

    for p in expanded:

        try:
            if not args.json:
                print(f"  {C.dim('Framing')} {p.name}...", end=" ", flush=True)

            result, info = frame_screenshot(
                p, assets_dir, device_dict,
                color_arg=args.color,
                force_device=args.device,
                verbose=args.verbose,
                user_defaults=user_defaults,
                device_index=device_index,
            )
            framed.append((result, p, info))
            infos.append(info)

            if not args.json:
                color_str = f" ({C.rgb(235, 125, 40, info['color'])})" if info["color"] else ""
                print(f"{C.green(info['device'])}{color_str}")

        except (
            ValueError,
            FileNotFoundError,
            OSError,
            UnidentifiedImageError,
            Image.DecompressionBombWarning,
            Image.DecompressionBombError,
        ) as e:
            print(f"{C.red('Error')}: {e}", file=sys.stderr)
            continue

    if not framed:
        print(f"{C.red('No screenshots were framed.')}", file=sys.stderr)
        sys.exit(1)

    # Merge if requested (--merge or --batch implies merge)
    do_merge = args.merge or batch_size

    if do_merge and len(framed) > 1:
        # Split into batches
        if batch_size and batch_size > 0:
            batches = [framed[i:i+batch_size] for i in range(0, len(framed), batch_size)]
        else:
            batches = [framed]  # Single batch = merge all

        if args.output:
            out_dir = Path(args.output)
        elif subfolder:
            out_dir = framed[0][1].parent / subfolder
        else:
            out_dir = framed[0][1].parent
        try:
            ensure_output_dir(out_dir)
        except OSError as e:
            print(f"{C.red('Error')}: {e}", file=sys.stderr)
            sys.exit(1)

        merged_outputs = []
        for batch_idx, batch in enumerate(batches):
            images = [f[0] for f in batch]

            # Collect physical heights for proportional scaling
            no_scale = getattr(args, 'no_scale', False)
            if not no_scale:
                phs = [f[2].get("physicalHeight", 0) for f in batch]
                physical_heights = phs if all(ph > 0 for ph in phs) else None
            else:
                physical_heights = None

            if len(images) == 1:
                merged = images[0]
            else:
                merged = merge_images(images, spacing=args.spacing, physical_heights=physical_heights)

            if len(batches) == 1:
                out_name = "merged_framed.png"
            else:
                out_name = f"merged_{batch_idx + 1}_framed.png"

            out_path = out_dir / out_name
            try:
                save_png(merged, out_path)
            except OSError as e:
                print(f"{C.red('Error')}: {e}", file=sys.stderr)
                sys.exit(1)
            merged_outputs.append(str(out_path))

            if not args.json:
                batch_count = len(batch)
                scaled_msg = ""
                if physical_heights and len(set(physical_heights)) > 1:
                    max_ph = max(physical_heights)
                    scaled_msg = f" {C.dim('(proportional)')}"
                print(f"\n  {C.green('Merged')} {batch_count} frames → {C.bold(str(out_path))}{scaled_msg}")

        if args.copy and len(merged_outputs) == 1:
            copy_to_clipboard(Path(merged_outputs[0]), announce=not args.json)

        if args.json:
            # Add scale factors when proportional scaling was applied
            all_phs = [info.get("physicalHeight", 0) for info in infos]
            if not no_scale and all(ph > 0 for ph in all_phs) and len(set(all_phs)) > 1:
                # Mirror the actual merge algorithm: normalize to lowest px/mm
                frame_heights = [int(info["frame_size"].split("x")[1]) for info in infos]
                px_per_mm = [fh / ph for fh, ph in zip(frame_heights, all_phs)]
                ref_ppm = min(px_per_mm)
                for info, fh, ph in zip(infos, frame_heights, all_phs):
                    target_h = round(ph * ref_ppm)
                    info["scale_factor"] = round(target_h / fh, 4)

            if len(batches) == 1:
                result_json = {"merged": merged_outputs[0], "count": len(framed), "frames": infos}
            else:
                result_json = {
                    "batches": [{"merged": p, "count": len(b)} for p, b in zip(merged_outputs, batches)],
                    "batch_size": batch_size,
                    "total": len(framed),
                    "frames": infos,
                }
            print(json.dumps(result_json, indent=2))

    else:
        # Save individual framed images
        outputs = []
        for result_img, src_path, info in framed:
            if args.output:
                out_dir = Path(args.output)
            elif subfolder:
                out_dir = src_path.parent / subfolder
            else:
                out_dir = src_path.parent
            try:
                ensure_output_dir(out_dir)
                out_path = out_dir / f"{src_path.stem}_framed.png"
                save_png(result_img, out_path)
            except OSError as e:
                print(f"{C.red('Error')}: {e}", file=sys.stderr)
                sys.exit(1)

            info["output"] = str(out_path)
            outputs.append(str(out_path))

            if not args.json:
                print(f"  {C.dim('→')} {C.bold(str(out_path))}")

        if args.copy and len(outputs) == 1:
            copy_to_clipboard(Path(outputs[0]), announce=not args.json)

        if args.json:
            if len(infos) == 1:
                print(json.dumps(infos[0], indent=2))
            else:
                print(json.dumps({"frames": infos}, indent=2))

    if not args.json:
        print()

def cmd_list(args):
    """List all supported devices."""
    assets_dir = Path(args.assets) if hasattr(args, 'assets') and args.assets else get_assets_dir()
    device_dict = load_json(assets_dir)

    print(f"\n  {C.bold('Supported Devices')}\n")

    # Group by category
    categories = {
        "iPhone 17": [], "iPhone Air": [], "iPhone 16": [],
        "iPhone 12-13": [], "iPhone 8": [],
        "iPad": [], "MacBook": [], "iMac": [], "Studio Display": [], "Watch": [],
    }

    for key in sorted(device_dict.keys(), key=lambda k: int(k) if k != "variants" else 0):
        if key == "variants":
            continue
        entry = device_dict[key]

        # Expand overlap entries (detect_device only checks overlap sub-entries)
        if "overlap" in entry:
            sub_entries = list(entry["overlap"].values())
        else:
            sub_entries = [entry]

        for sub in sub_entries:
            name = sub["name"]
            has_colors = sub.get("colors") == "yes"
            resize = sub.get("resizeWidth")

            for cat in categories:
                if name.startswith(cat) or (cat == "MacBook" and "MacBook" in name) or (cat == "Watch" and "Watch" in name) or (cat == "Studio Display" and "Studio Display" in name):
                    flags = []
                    if has_colors:
                        colors = FRAME_COLORS.get(name, [])
                        if colors:
                            flags.append(f"{len(colors)} colors")
                    if resize:
                        flags.append("scaled")
                    flag_str = C.dim(f" ({', '.join(flags)})") if flags else ""
                    categories[cat].append(f"    {C.cyan(f'{key:>5}px')} → {name}{flag_str}")
                    break

    for cat, entries in categories.items():
        if entries:
            print(f"  {C.bold(cat)}")
            for e in entries:
                print(e)
            print()

    # Show variant defaults
    print(f"  {C.bold('Default Variants')} {C.dim('(newest device used automatically)')}")
    for primary, variants in FRAME_VARIANTS.items():
        default = variants[0]
        if default != primary:
            print(f"    {C.dim(primary)} → {C.green(default)}")
    print()

def cmd_list_colors(args):
    """List colors for a specific device (partial match)."""
    device = args.device_name
    matches = [(k, v) for k, v in FRAME_COLORS.items() if device.lower() in k.lower()]

    if not matches:
        print(f"{C.red('Error')}: No device matching '{device}'", file=sys.stderr)
        sys.exit(1)

    print()
    for name, colors in matches:
        print(f"  {C.bold(name)}")
        for i, color in enumerate(colors):
            default = C.dim(" (default)") if i == 0 else ""
            print(f"    {C.cyan(str(i+1))}. {color}{default}")
        print()

def _build_device_list():
    """Build deduplicated device→colors list (no portrait/landscape dupes)."""
    devices = []
    seen = set()
    for name, colors in FRAME_COLORS.items():
        base = name.replace(" Portrait", "").replace(" Landscape", "")
        if base not in seen:
            seen.add(base)
            devices.append((base, list(colors)))
    return devices

def _print_colors_fallback():
    """Print device colors without curses."""
    O = lambda t: C.rgb(232, 114, 42, t)
    devices = _build_device_list()
    cfg = load_config()
    user_defaults = cfg.get("default_colors", {})
    print(f"\n  {C.bold('All Device Colors')}\n")
    for base, colors in devices:
        print(f"  {C.bold(base)}")
        for i, color in enumerate(colors):
            is_default = (base in user_defaults and user_defaults[base] == color) or (base not in user_defaults and i == 0)
            marker = O("●") if is_default else C.dim("○")
            print(f"    {marker} {color}")
        print()

def _colors_curses(stdscr):
    """Curses-based interactive color picker."""
    import curses

    curses.curs_set(0)
    curses.use_default_colors()
    curses.start_color()
    curses.init_pair(1, 209, -1)   # Orange
    curses.init_pair(2, 250, -1)   # Dim

    devices = _build_device_list()
    cfg = load_config()
    user_defaults = cfg.get("default_colors", {})
    selected = {}
    for base, colors in devices:
        if base in user_defaults and user_defaults[base] in colors:
            selected[base] = colors.index(user_defaults[base])
        else:
            selected[base] = 0

    # Flat list of navigable items
    items = []
    for base, colors in devices:
        for ci in range(len(colors)):
            items.append((base, ci))

    cursor = 0
    changed = False

    while True:
        stdscr.clear()
        h, w = stdscr.getmaxyx()

        # Header
        try:
            stdscr.addstr(0, 2, "Default Colors", curses.A_BOLD)
            stdscr.addstr(0, 18, "  ↑↓ navigate  space select  enter save  q cancel", curses.color_pair(2))
        except curses.error:
            pass

        row = 2
        idx = 0
        # Calculate scroll offset so cursor is always visible
        # Each device takes len(colors) + 2 rows (header + colors + blank)
        # Simple approach: center the view around cursor
        total_rows = sum(len(c) + 2 for _, c in devices)
        scroll_offset = 0
        if total_rows > h - 4:
            # Find which row the cursor is on
            cursor_row = 2
            ci = 0
            for base, colors in devices:
                cursor_row += 1  # device header
                for c_idx in range(len(colors)):
                    if ci == cursor:
                        break
                    cursor_row += 1
                    ci += 1
                else:
                    cursor_row += 1  # blank line
                    continue
                break
            scroll_offset = max(0, cursor_row - h // 2)

        row = 2 - scroll_offset
        idx = 0
        for base, colors in devices:
            if 0 <= row < h:
                try:
                    stdscr.addstr(row, 2, base, curses.A_BOLD)
                except curses.error:
                    pass
            row += 1

            for ci, color in enumerate(colors):
                if 0 <= row < h - 1:
                    is_cursor = (idx == cursor)
                    is_selected = (selected[base] == ci)

                    try:
                        if is_cursor:
                            stdscr.addstr(row, 2, "▸ ", curses.A_BOLD)
                        else:
                            stdscr.addstr(row, 2, "  ")

                        if is_selected:
                            stdscr.addstr(row, 4, "● ", curses.color_pair(1))
                        else:
                            stdscr.addstr(row, 4, "○ ", curses.color_pair(2))

                        attr = curses.A_BOLD if is_cursor else 0
                        stdscr.addstr(row, 6, color, attr)
                    except curses.error:
                        pass
                row += 1
                idx += 1
            row += 1  # blank between devices

        stdscr.refresh()

        key = stdscr.getch()
        if key == curses.KEY_DOWN:
            cursor = min(cursor + 1, len(items) - 1)
        elif key == curses.KEY_UP:
            cursor = max(cursor - 1, 0)
        elif key == ord(' '):
            base, ci = items[cursor]
            selected[base] = ci
            changed = True
        elif key in (ord('\n'), ord('\r')):
            if changed:
                defaults = {}
                for base, colors in devices:
                    defaults[base] = colors[selected[base]]
                cfg["default_colors"] = defaults
                save_config(cfg)
            return changed, selected, devices
        elif key in (ord('q'), 27):  # q or Esc
            return False, selected, devices

def cmd_colors(args):
    """Interactive color picker — arrow keys to navigate, space to select, enter to save."""
    O = lambda t: C.rgb(232, 114, 42, t)

    if not sys.stdin.isatty():
        _print_colors_fallback()
        return

    try:
        import curses
    except ImportError:
        print(f"{C.yellow('Warning')}: curses is unavailable. Showing non-interactive color list instead.", file=sys.stderr)
        _print_colors_fallback()
        return

    changed, selected, devices = curses.wrapper(_colors_curses)

    if changed:
        print(f"\n  {C.green('Defaults saved!')}\n")
        for base, colors in devices:
            color = colors[selected[base]]
            print(f"    {O('●')} {C.bold(base)}: {color}")
        print()
    else:
        print(f"\n  {C.dim('No changes saved.')}\n")

def cmd_setup(args):
    """Configure the assets directory path, or download assets interactively."""
    if args.path is None:
        if not sys.stdin.isatty():
            print(f"{C.red('Error')}: Run 'frames setup' in an interactive terminal, or provide a path.", file=sys.stderr)
            sys.exit(1)
        if guided_setup():
            return
        sys.exit(1)

    assets_path = Path(args.path).resolve()

    # Validate the path
    if not assets_path.exists():
        print(f"{C.red('Error')}: Path does not exist: {assets_path}", file=sys.stderr)
        sys.exit(1)

    json_file = assets_path / "NewFrames.json"
    if not json_file.exists():
        print(f"{C.red('Error')}: NewFrames.json not found in {assets_path}", file=sys.stderr)
        print(f"  {C.dim('The path should point to the folder containing NewFrames.json and frame PNGs.')}", file=sys.stderr)
        sys.exit(1)

    ok, ver = check_assets_version(assets_path)
    if not ok:
        reason = "missing or invalid Apple Frames 4 assets" if ver is None else f"outdated asset version {ver}"
        print(f"{C.red('Error')}: {reason} in {assets_path}", file=sys.stderr)
        sys.exit(1)

    # Count assets
    pngs = list(assets_path.glob("*.png"))
    with open(json_file, encoding="utf-8") as f:
        device_dict = json.load(f)
    device_count = len([k for k in device_dict if k != "variants"])

    # Save config
    config = load_config()
    config["assets_path"] = str(assets_path)
    if hasattr(args, 'subfolder') and args.subfolder is not None:
        config["use_subfolder"] = args.subfolder
    save_config(config)

    subfolder_status = config.get("use_subfolder", False)

    print()
    print(f"  {C.green('Setup complete!')}")
    print(f"  {C.bold('Assets path:')} {assets_path}")
    print(f"  {C.bold('Frame PNGs:')}  {len(pngs)}")
    print(f"  {C.bold('Devices:')}     {device_count}")
    print(f"  {C.bold('Subfolder:')}   {'on' if subfolder_status else 'off'} {C.dim('(save to /framed/ subfolder)')}")
    print(f"  {C.bold('Config:')}      {CONFIG_FILE}")
    print()
    print(f"  {C.dim('You can now run')} {C.cyan('frames screenshot.png')} {C.dim('from anywhere.')}")
    print()

def _gather_image_paths(paths):
    """Expand a mix of files and directories into image file paths."""
    result = []
    for p in paths:
        p = Path(p)
        if p.is_dir():
            for f in sorted(p.iterdir()):
                if f.suffix.lower() in VALID_IMAGE_EXTS and not f.name.startswith("."):
                    result.append(f)
        elif p.is_file():
            result.append(p)
        else:
            print(f"{C.red('Error')}: {p} not found", file=sys.stderr)
    return result

def _info_for_image(p, device_dict):
    """Analyze a single image. Returns an info dict."""
    with Image.open(p) as img:
        w, h = img.size
        guard_image_size(w, h, p)
        entry, model_name = detect_device(img, device_dict)
    info = {"file": str(p), "name": p.name, "width": w, "height": h}

    if not entry:
        info["device"] = None
        info["error"] = f"Unknown device (width {w}px)"
        return info

    # Resolve variant
    variant_entry, resolved_name = resolve_variant(model_name, device_dict)
    if variant_entry:
        entry = variant_entry

    colors = FRAME_COLORS.get(resolved_name, [])
    variants = FRAME_VARIANTS.get(model_name, [])

    info.update({
        "device": resolved_name,
        "primary_match": model_name,
        "x": entry["x"], "y": entry["y"],
        "colors": colors,
        "color_count": len(colors),
        "has_mask": entry.get("mask") == "yes",
        "resize_width": entry.get("resizeWidth"),
        "variants": variants,
        "is_variant": resolved_name != model_name,
    })
    return info

def cmd_info(args):
    """Detect device for screenshots or scan a folder."""
    assets_dir = Path(args.assets) if hasattr(args, 'assets') and args.assets else get_assets_dir()
    device_dict = load_json(assets_dir)

    image_paths = _gather_image_paths(args.files)
    if not image_paths:
        print(f"{C.red('No images found.')}", file=sys.stderr)
        return

    results = []
    for p in image_paths:
        try:
            info = _info_for_image(p, device_dict)
            results.append(info)
        except (
            ValueError,
            FileNotFoundError,
            OSError,
            UnidentifiedImageError,
            Image.DecompressionBombWarning,
            Image.DecompressionBombError,
        ) as e:
            results.append({"file": str(p), "name": p.name, "error": str(e)})

    if args.json:
        print(json.dumps(results if len(results) > 1 else results[0], indent=2))
        return

    # Pretty print
    matched = [r for r in results if r.get("device")]
    unmatched = [r for r in results if not r.get("device")]

    if len(image_paths) > 1:
        print(f"\n  {C.bold('Scanned')} {len(image_paths)} images: "
              f"{C.green(str(len(matched)) + ' matched')}, "
              f"{C.red(str(len(unmatched)) + ' unknown') if unmatched else C.dim('0 unknown')}")

    for r in results:
        print()
        print(f"  {C.bold(r['name'])}")
        if r.get("error"):
            print(f"  {C.red('✗')} {r['error']}")
            continue

        dims = f'{r["width"]} × {r["height"]}'
        print(f"  {C.dim(dims)}")

        device = r["device"]
        colors = r.get("colors", [])
        variants = r.get("variants", [])

        print(f"  {C.green('✓')} {C.bold(device)}")

        if r.get("is_variant"):
            print(f"    {C.dim('Primary match:')} {r['primary_match']} {C.dim('→ variant resolved')}")

        if colors:
            print(f"    {C.cyan(str(len(colors)))} colors: {C.dim(', '.join(colors[:5]))}"
                  + (f" {C.dim(f'(+{len(colors)-5} more)')}" if len(colors) > 5 else ""))
        else:
            print(f"    {C.dim('No color variants')}")

        if len(variants) > 1:
            alt = [v for v in variants if v != device]
            print(f"    {C.dim('Also available:')} {', '.join(alt)}")

        if r.get("resize_width"):
            print(f"    {C.dim('Scaled from')} {r['width']}px {C.dim('→')} {r['resize_width']}px")
        if r.get("has_mask"):
            print(f"    {C.dim('Mask applied')}")

    # Summary by device when scanning folders with many files
    if len(matched) > 3:
        from collections import Counter
        device_counts = Counter(r["device"] for r in matched)
        print(f"\n  {C.bold('Summary')}")
        for device, count in device_counts.most_common():
            colors = FRAME_COLORS.get(device, [])
            color_tag = f" ({len(colors)} colors)" if colors else ""
            print(f"    {count:3d}× {device}{C.dim(color_tag)}")
        if unmatched:
            print(f"    {C.red(f'{len(unmatched)}× unknown')}")
    print()

def copy_to_clipboard(path, announce=True):
    """Copy image to macOS clipboard."""
    if sys.platform != "darwin":
        print(f"  {C.yellow('Warning')}: Clipboard copy is only supported on macOS", file=sys.stderr)
        return
    try:
        script = (
            'set the clipboard to '
            '(read (POSIX file (system attribute "FRAMES_CLIPBOARD_PATH")) as «class PNGf»)'
        )
        env = os.environ.copy()
        env["FRAMES_CLIPBOARD_PATH"] = str(path)
        subprocess.run([
            "osascript", "-e",
            script,
        ], check=True, capture_output=True, env=env)
        if announce:
            print(f"  {C.green('Copied to clipboard')}", file=sys.stderr)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print(f"  {C.yellow('Warning')}: Could not copy to clipboard", file=sys.stderr)


def cmd_doctor(args):
    """Diagnose common asset and config problems."""
    assets_dir, assets_source = resolve_assets_dir(getattr(args, "assets", None))
    config, config_error = read_config_file()
    ok, ver = check_assets_version(assets_dir)

    json_path = assets_dir / "NewFrames.json"
    png_count = len(list(assets_dir.glob("*.png"))) if assets_dir.exists() else 0
    issues = []
    notes = []

    if config_error:
        issues.append(f"Config file is unreadable: {config_error}")
    elif config is None:
        notes.append("Config file not found. This is fine if you use defaults or FRAMES_ASSETS.")

    if not assets_dir.exists():
        issues.append(f"Assets directory does not exist: {assets_dir}")
    elif not json_path.exists():
        issues.append(f"NewFrames.json is missing from assets directory: {assets_dir}")
    elif not ok:
        if ver is None:
            issues.append(f"Assets at {assets_dir} do not look like Apple Frames 4.")
        else:
            issues.append(f"Assets are outdated at {assets_dir} (version {ver}).")

    if assets_dir.exists() and ok and png_count == 0:
        issues.append(f"No frame PNGs were found in {assets_dir}.")

    bad_default_colors = []
    if config and isinstance(config.get("default_colors"), dict):
        for device, color in config["default_colors"].items():
            matches = [
                name for name, colors in FRAME_COLORS.items()
                if name.replace(" Portrait", "").replace(" Landscape", "") == device
            ]
            if matches and all(color not in FRAME_COLORS[name] for name in matches):
                bad_default_colors.append(f"{device}: {color}")
    if bad_default_colors:
        issues.append("Some saved default colors no longer match supported device colors.")

    if assets_source == "FRAMES_ASSETS":
        notes.append("Using FRAMES_ASSETS from the environment.")
    elif assets_source == "--assets":
        notes.append("Using assets directory from the CLI flag.")
    elif assets_source == "default" and sys.platform == "darwin":
        notes.append("Using the Apple Frames Shortcuts/iCloud assets directory.")
    elif assets_source == "app-data":
        notes.append("Using the fallback Application Support assets directory.")

    next_steps = []
    if any("Config file is unreadable" in issue for issue in issues):
        next_steps.append(f"Fix or remove {CONFIG_FILE}.")
    if any("Assets" in issue or "NewFrames.json" in issue or "frame PNGs" in issue for issue in issues):
        next_steps.append("Run `frames setup` to re-download or re-point your assets.")
    if bad_default_colors:
        next_steps.append("Run `frames colors` to save fresh defaults, or remove stale entries from config.")
    if not next_steps:
        next_steps.append("No action needed.")

    result = {
        "ok": not issues,
        "assets_path": str(assets_dir),
        "assets_source": assets_source,
        "assets_version": ver,
        "frame_pngs": png_count,
        "config_path": str(CONFIG_FILE),
        "config_present": CONFIG_FILE.exists(),
        "issues": issues,
        "notes": notes,
        "next_steps": next_steps,
    }
    if bad_default_colors:
        result["bad_default_colors"] = bad_default_colors

    if args.json:
        print(json.dumps(result, indent=2))
        return

    print()
    status = C.green("OK") if result["ok"] else C.red("Needs attention")
    print(f"  {C.bold('Frames Doctor')} {status}")
    print(f"  {C.bold('Assets:')} {assets_dir}")
    print(f"  {C.bold('Source:')} {assets_source}")
    version_text = f"v{ver}" if ver is not None else "unknown"
    print(f"  {C.bold('Version:')} {version_text}")
    print(f"  {C.bold('Frame PNGs:')} {png_count}")
    print(f"  {C.bold('Config:')} {CONFIG_FILE}")

    if notes:
        print()
        print(f"  {C.bold('Notes')}")
        for note in notes:
            print(f"    {C.dim('-')} {note}")

    if issues:
        print()
        print(f"  {C.bold('Issues')}")
        for issue in issues:
            print(f"    {C.red('•')} {issue}")

    print()
    print(f"  {C.bold('Next steps')}")
    for step in next_steps:
        print(f"    {C.cyan('•')} {step}")
    print()

# ── Argument Parsing ─────────────────────────────────────────────────────────

def build_parser():
    epilog = (
        "Default framing flags also work without the `frame` subcommand.\n"
        "Examples:\n"
        "  frames -m shot1.png shot2.png\n"
        "  frames -m --no-scale shot1.png shot2.png\n"
        "  frames frame --help\n"
        "  frames doctor"
    )
    p = argparse.ArgumentParser(
        prog="frames",
        description="Frame screenshots with Apple device bezels",
        add_help=True,
        epilog=epilog,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--version", action="version", version=f"frames v{VERSION}")
    p.add_argument("--no-color", action="store_true", help="Disable colored output")
    p.add_argument("--json", action="store_true", help="Output JSON (for AI agents)")
    p.add_argument("--assets", default=None, help="Assets directory path")
    p.add_argument("-v", "--verbose", action="store_true", help="Verbose output")

    sub = p.add_subparsers(dest="cmd")

    # Default frame command (also works without subcommand)
    frame_p = sub.add_parser("frame", help="Frame screenshots (default command)")
    frame_p.add_argument("files", nargs="*", help="Screenshot files to frame")
    frame_p.add_argument("--color", "-c", help="Frame color name, index, or 'random'")
    frame_p.add_argument("--merge", "-m", action="store_true", help="Merge framed images side by side")
    frame_p.add_argument("--batch", "-b", type=int, metavar="N", help="Merge in batches of N (e.g. -b 3 for groups of 3)")
    frame_p.add_argument("--spacing", "-s", type=int, default=60, help="Merge spacing in pixels (default: 60)")
    frame_p.add_argument("--no-scale", action="store_true", help="Disable proportional scaling when merging different devices")
    frame_p.add_argument("--output", "-o", help="Output directory")
    frame_p.add_argument("--copy", action="store_true", help="Copy result to clipboard")
    frame_p.add_argument("--device", "-d", help="Force specific device frame")
    frame_p.add_argument("-f", action="store_const", const="framed", dest="subfolder", help="Save to /framed/ subfolder")
    frame_p.add_argument("--subfolder", dest="subfolder", metavar="NAME", help="Save to custom subfolder (e.g. --subfolder mockups)")
    frame_p.add_argument("--json", action="store_true", dest="frame_json", help=argparse.SUPPRESS)

    # List command
    sub.add_parser("list", help="List supported devices")

    # List colors command (partial match)
    colors_p = sub.add_parser("list-colors", help="Show colors for a specific device")
    colors_p.add_argument("device_name", help="Device name (partial match)")

    # Colors command (browse all)
    sub.add_parser("colors", help="Browse all device colors")

    # Info command
    info_p = sub.add_parser("info", help="Detect device for screenshots or folders")
    info_p.add_argument("files", nargs="+", help="Screenshot files or directories to scan")

    # Setup command
    setup_p = sub.add_parser("setup", help="Download assets or configure path")
    setup_p.add_argument("path", nargs='?', default=None, help="Path to assets folder (omit to download)")
    setup_p.add_argument("--subfolder", action="store_true", default=None, help="Save framed images to /framed/ subfolder")
    setup_p.add_argument("--no-subfolder", action="store_false", dest="subfolder", help="Save framed images next to originals (default)")

    # Doctor command
    sub.add_parser("doctor", help="Diagnose asset and config problems")

    return p

def main():
    p = build_parser()

    # Pre-process argv: if the first non-flag arg isn't a known subcommand,
    # assume it's a file and inject "frame" as the subcommand
    known_cmds = {"frame", "list", "list-colors", "colors", "info", "setup", "doctor", "-h", "--help", "--version"}
    argv = sys.argv[1:]
    if argv:
        global_flags = {"-h", "--help", "--version", "--no-color", "--json", "-v", "--verbose"}
        insert_at = 0
        idx = 0
        while idx < len(argv):
            arg = argv[idx]
            if arg in known_cmds:
                break
            if arg in global_flags:
                idx += 1
                insert_at = idx
                continue
            if arg == "--assets":
                idx += 2
                insert_at = min(idx, len(argv))
                continue
            if arg.startswith("--assets="):
                idx += 1
                insert_at = idx
                continue
            argv = argv[:insert_at] + ["frame"] + argv[insert_at:]
            break

    args = p.parse_args(argv)

    # Disable colors if requested
    if args.no_color or os.environ.get("NO_COLOR") or not sys.stdout.isatty():
        C.enabled = False

    # No command → show splash
    if args.cmd is None:
        if not args.json:
            show_splash()
        return

    # Propagate global flags
    if not hasattr(args, "verbose"):
        args.verbose = False
    if not hasattr(args, "json"):
        args.json = False
    if hasattr(args, "frame_json") and args.frame_json:
        args.json = True

    # Resolve assets directory (setup doesn't need it)
    if args.cmd not in ("setup", "doctor"):
        args.assets = str(get_assets_dir(args.assets))

    # Route commands
    cmds = {
        "frame": cmd_frame,
        "list": cmd_list,
        "list-colors": cmd_list_colors,
        "colors": cmd_colors,
        "info": cmd_info,
        "setup": cmd_setup,
        "doctor": cmd_doctor,
    }
    handler = cmds.get(args.cmd)
    if handler:
        # Check assets exist and are v4+ (except for setup, colors, list-colors)
        if args.cmd not in ("setup", "colors", "list-colors", "doctor"):
            ok, ver = check_assets_version(Path(args.assets))
            if not ok:
                if not sys.stdin.isatty() or args.json:
                    reason = "not found" if ver is None else f"outdated (version {ver})"
                    if args.json:
                        print(json.dumps({"error": f"Assets {reason}", "setup_required": True}))
                    else:
                        print(f"{C.red('Error')}: Assets {reason}. Run 'frames setup' interactively.", file=sys.stderr)
                    sys.exit(1)
                reason = "not found" if ver is None else f"outdated (version {ver})"
                if guided_setup(reason=reason):
                    args.assets = str(get_assets_dir())
                else:
                    sys.exit(1)
        handler(args)

if __name__ == "__main__":
    main()