#!/usr/bin/env python3
"""
Created: 2026-06-02
Updated: 2026-06-02
Name: remove_icon_padding.py
Desc: Remove transparent padding from icon images, fill content to full size,
      and apply iOS-style rounded corners (superellipse mask).
      Original files are backed up to a 'backup' subfolder.
"""

import os
import shutil
import numpy as np
from PIL import Image, ImageDraw


INPUT_DIR = '/Users/wushu/Documents/trae_projects/project/xianyan/assets/templates/resized'
BACKUP_DIR = os.path.join(INPUT_DIR, 'backup_original')

ALPHA_THRESHOLD = 10
CORNER_RATIO = 0.2237


def find_content_bounds(arr):
    """Find bounding box of non-transparent content."""
    alpha = arr[:, :, 3]
    rows = np.any(alpha > ALPHA_THRESHOLD, axis=1)
    cols = np.any(alpha > ALPHA_THRESHOLD, axis=0)
    if not rows.any():
        return None
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    return (cmin, rmin, cmax + 1, rmax + 1)


def create_superellipse_mask(size, n=5):
    """
    Create iOS-style superellipse (squircle) mask.
    The equation: |x|^n + |y|^n = 1 with n=5 approximates Apple's icon shape.
    """
    w, h = size
    mask = Image.new('L', (w, h), 0)
    draw = ImageDraw.Draw(mask)

    cx, cy = w / 2.0, h / 2.0
    a = w / 2.0
    b = h / 2.0

    points = []
    num_points = 360
    for i in range(num_points + 1):
        t = 2.0 * np.pi * i / num_points
        cos_t = np.cos(t)
        sin_t = np.sin(t)

        x = np.sign(cos_t) * abs(cos_t) ** (2.0 / n) * a
        y = np.sign(sin_t) * abs(sin_t) ** (2.0 / n) * b

        points.append((cx + x, cy + y))

    draw.polygon(points, fill=255)
    return mask


def create_rounded_rect_mask(size, radius):
    """Create a simple rounded rectangle mask as fallback."""
    w, h = size
    mask = Image.new('L', (w, h), 0)
    draw = ImageDraw.Draw(mask)
    draw.rounded_rectangle([(0, 0), (w - 1, h - 1)], radius=radius, fill=255)
    return mask


def process_icon(filepath, backup_dir):
    """Process a single icon: backup, crop padding, resize, apply rounded mask."""
    filename = os.path.basename(filepath)
    img = Image.open(filepath)
    original_size = img.size
    arr = np.array(img)

    if img.mode != 'RGBA':
        print(f"  SKIP {filename}: not RGBA mode ({img.mode})")
        return False

    bounds = find_content_bounds(arr)
    if bounds is None:
        print(f"  SKIP {filename}: no visible content found")
        return False

    cmin, rmin, cmax, rmax = bounds
    h, w = arr.shape[:2]
    pad_top = rmin
    pad_bottom = h - rmax
    pad_left = cmin
    pad_right = w - cmax

    if pad_top < 2 and pad_bottom < 2 and pad_left < 2 and pad_right < 2:
        print(f"  SKIP {filename}: no significant padding detected")
        return False

    print(f"  Padding: top={pad_top} bottom={pad_bottom} left={pad_left} right={pad_right}")

    # Backup original
    backup_path = os.path.join(backup_dir, filename)
    if not os.path.exists(backup_path):
        shutil.copy2(filepath, backup_path)
        print(f"  Backed up to {backup_path}")

    # Crop to content
    cropped = img.crop(bounds)

    # Resize back to original dimensions (high quality)
    resized = cropped.resize(original_size, Image.LANCZOS)

    # Apply iOS superellipse mask for rounded corners
    mask = create_superellipse_mask(original_size, n=5)

    # Composite: apply mask as alpha channel
    result = resized.copy()
    result.putalpha(mask)

    # Save result
    result.save(filepath, 'PNG')
    print(f"  DONE: {filename} -> cropped, resized to {original_size}, superellipse mask applied")
    return True


def main():
    os.makedirs(BACKUP_DIR, exist_ok=True)
    print(f"Input dir: {INPUT_DIR}")
    print(f"Backup dir: {BACKUP_DIR}")
    print()

    files = sorted([
        f for f in os.listdir(INPUT_DIR)
        if f.lower().endswith('.png')
    ])

    processed = 0
    skipped = 0

    for f in files:
        filepath = os.path.join(INPUT_DIR, f)
        if not os.path.isfile(filepath):
            continue
        print(f"[{processed + skipped + 1}/{len(files)}] {f}")
        if process_icon(filepath, BACKUP_DIR):
            processed += 1
        else:
            skipped += 1

    print(f"\n{'='*50}")
    print(f"Total: {len(files)} files")
    print(f"Processed: {processed}")
    print(f"Skipped: {skipped}")
    print(f"Backup location: {BACKUP_DIR}")


if __name__ == '__main__':
    main()