"""
鸿蒙(HarmonyOS) 分层图标生成脚本
从源图标提取 foreground + background，符合华为审核规范 1024x1024
"""
import os
import json
import math

from PIL import Image, ImageDraw

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
ASSETS_ICONS = os.path.join(SCRIPT_DIR, "..", "assets", "icons")
OHOS_MEDIA = os.path.join(SCRIPT_DIR, "..", "ohos", "AppScope", "resources", "base", "media")
OHOS_ENTRY_MEDIA = os.path.join(SCRIPT_DIR, "..", "ohos", "entry", "src", "main", "resources", "base", "media")

SOURCE_ICON = os.path.join(ASSETS_ICONS, "icon_1024x1024.png")
OUTPUT_SIZE = 1024


def get_dominant_bg_color(img):
    samples = [
        (80, 80), (img.width - 81, 80),
        (80, img.height - 81), (img.width - 81, img.height - 81),
        (img.width // 2, 60), (img.width // 2, img.height - 61),
        (60, img.height // 2), (img.width - 61, img.height // 2),
        (150, 150), (img.width - 151, img.height - 151),
    ]
    colors = []
    for pos in samples:
        px = img.getpixel(pos)
        if len(px) >= 3:
            colors.append((px[0], px[1], px[2]))
    r_sum = sum(c[0] for c in colors)
    g_sum = sum(c[1] for c in colors)
    b_sum = sum(c[2] for c in colors)
    n = len(colors)
    return (r_sum // n, g_sum // n, b_sum // n)


def color_distance(c1, c2):
    return math.sqrt(
        (c1[0] - c2[0]) ** 2 +
        (c1[1] - c2[1]) ** 2 +
        (c1[2] - c2[2]) ** 2
    )


def create_rounded_mask(size, radius):
    mask = Image.new("L", (size, size), 0)
    draw = ImageDraw.Draw(mask)
    draw.rounded_rectangle([0, 0, size - 1, size - 1], radius=radius, fill=255)
    return mask


def generate_background(img, output_path, bg_color):
    bg = Image.new("RGB", (OUTPUT_SIZE, OUTPUT_SIZE), bg_color)
    bg.save(output_path, "PNG")
    print(f"  background: {output_path} ({os.path.getsize(output_path) / 1024:.1f}KB)")
    print(f"    颜色: RGB{bg_color}")


def generate_foreground(img, output_path, bg_color, tolerance=50):
    from PIL import ImageChops

    fg = img.convert("RGBA").resize((OUTPUT_SIZE, OUTPUT_SIZE), Image.LANCZOS)

    pixels = fg.load()
    width, height = fg.size

    for y in range(height):
        for x in range(width):
            r, g, b, a = pixels[x, y]
            dist = color_distance((r, g, b), bg_color)
            if dist < tolerance and a > 200:
                alpha = int(255 * (dist / tolerance))
                pixels[x, y] = (r, g, b, alpha)
            elif dist < tolerance * 0.5:
                pixels[x, y] = (r, g, b, 0)

    mask = create_rounded_mask(OUTPUT_SIZE, int(OUTPUT_SIZE * 0.22))

    fg.putalpha(ImageChops.multiply(fg.split()[3], mask))

    fg.save(output_path, "PNG")
    print(f"  foreground: {output_path} ({os.path.getsize(output_path) / 1024:.1f}KB)")


def generate_layered_json(output_dir, fg_name="foreground_icon.png", bg_name="background_icon.png"):
    config = {
        "layered-image": {
            "foreground": f"$media:{fg_name.replace('.png', '')}",
            "background": f"$media:{bg_name.replace('.png', '')}"
        }
    }
    json_path = os.path.join(output_dir, "layered_image.json")
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, ensure_ascii=False)
    print(f"  layered_image.json: {json_path}")
    return json_path


def main():
    if not os.path.exists(SOURCE_ICON):
        print(f"错误: 源文件不存在 {SOURCE_ICON}")
        return

    print(f"源图标: {SOURCE_ICON}")
    img = Image.open(SOURCE_ICON).convert("RGBA")
    print(f"  尺寸: {img.size}")

    bg_color = get_dominant_bg_color(img)
    print(f"  背景色采样: RGB{bg_color}")

    os.makedirs(OHOS_MEDIA, exist_ok=True)
    os.makedirs(OHOS_ENTRY_MEDIA, exist_ok=True)

    print("\n--- AppScope/resources/base/media ---")

    bg_path = os.path.join(OHOS_MEDIA, "background_icon.png")
    fg_path = os.path.join(OHOS_MEDIA, "foreground_icon.png")

    generate_background(img, bg_path, bg_color)
    generate_foreground(img, fg_path, bg_color)
    generate_layered_json(OHOS_MEDIA)

    print("\n--- entry/src/main/resources/base/media ---")

    bg_entry = os.path.join(OHOS_ENTRY_MEDIA, "background_icon.png")
    fg_entry = os.path.join(OHOS_ENTRY_MEDIA, "foreground_icon.png")

    generate_background(img, bg_entry, bg_color)
    generate_foreground(img, fg_entry, bg_color)
    generate_layered_json(OHOS_ENTRY_MEDIA)

    print("\n完成!")


if __name__ == "__main__":
    main()