#!/usr/bin/env python3
from __future__ import annotations

import argparse
import subprocess
import sys
from collections import deque
from dataclasses import dataclass
from pathlib import Path

from PIL import Image, ImageDraw


@dataclass(frozen=True)
class IconSpec:
    source: str
    output_name: str
    target_size: tuple[int, int]
    tres_path: str
    icon_height_ratio: float = 0.72


ICON_SPECS: tuple[IconSpec, ...] = (
    IconSpec("ChatGPT Image 2026年5月11日 21_57_01 (1).png", "world", (156, 161), "assets/egret/menu2/menu2_m_word_png.tres", 0.78),
    IconSpec("ChatGPT Image 2026年5月11日 21_57_01 (2).png", "home", (155, 145), "assets/egret/menu2/menu2_m_home_png.tres", 0.76),
    IconSpec("ChatGPT Image 2026年5月11日 23_13_41 (4).png", "store", (86, 89), "assets/egret/menu2/menu2_m_shop_png.tres", 0.72),
    IconSpec("ChatGPT Image 2026年5月11日 23_13_40 (2).png", "storage", (84, 101), "assets/egret/menu2/menu2_m_warehouse_png.tres", 0.72),
    IconSpec("ChatGPT Image 2026年5月11日 23_13_42 (6).png", "friends", (74, 92), "assets/egret/menu2/menu2_m_rank_png.tres", 0.72),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_54 (1).png", "benefits_hall", (85, 89), "assets/egret/menu2/menu2_left_fuli_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_54 (2).png", "news", (86, 78), "assets/egret/menu2/menu2_left_notice_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_54 (3).png", "daily_check_in", (86, 78), "assets/egret/menu2/menu2_left_sign_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_55 (4).png", "online_gift", (85, 76), "assets/egret/menu2/menu2_left_rewardgift_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_55 (5).png", "log", (74, 74), "assets/egret/menu2/menu2_top_log_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_55 (6).png", "land_upgrade", (96, 82), "assets/egret/menu2/menu2_top_landup_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_56 (7).png", "top_up", (56, 80), "assets/egret/menu2/menu2_toppay_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_56 (8).png", "level_pack", (85, 80), "assets/egret/menu2/menu2_top_paygift_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_56 (9).png", "exchange", (93, 83), "assets/egret/menu2/menu2_top_crystal_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 22_15_09.png", "factory", (105, 90), "assets/egret/mofajiagong/mofajiagong_tubiao_png.tres", 0.70),
    IconSpec("ChatGPT Image 2026年5月11日 21_36_57 (10).png", "merge", (78, 85), "assets/egret/menu2/menu2_tophecheng_png.tres"),
    IconSpec("ChatGPT Image 2026年5月11日 21_57_02 (3).png", "lucky_wheel", (86, 83), "assets/egret/menu2/menu2_top_lottery_png.tres"),
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Batch remove green backgrounds, resize farm UI icons, and rewrite Godot .tres references."
    )
    parser.add_argument("--project", default="godot/FarmGodot", help="Godot project directory.")
    parser.add_argument("--input-dir", default=str(Path.home() / "Downloads"), help="Directory containing source PNG files.")
    parser.add_argument("--method", choices=("green", "ai", "auto"), default="green", help="Matting method.")
    parser.add_argument("--install-ai", action="store_true", help="Install rembg/onnxruntime for --method ai.")
    parser.add_argument("--preview", default="/tmp/farm_prepared_icons_preview.png", help="Preview image path.")
    parser.add_argument(
        "--cutout",
        nargs=4,
        action="append",
        metavar=("SOURCE", "OUTPUT", "WIDTH", "HEIGHT"),
        help="Remove green background from SOURCE and fit it into OUTPUT at WIDTHxHEIGHT.",
    )
    return parser.parse_args()


def maybe_install_ai() -> None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--user", "rembg", "onnxruntime"])


def remove_with_ai(image: Image.Image) -> Image.Image:
    try:
        from rembg import remove
    except ImportError as exc:
        raise RuntimeError("rembg is not installed. Run with --install-ai or use --method green.") from exc
    return remove(image.convert("RGBA"))


def is_green_background(pixel: tuple[int, int, int, int]) -> bool:
    r, g, b, a = pixel
    return a > 0 and g >= 120 and g > r * 1.28 and g > b * 1.28


def remove_green_background(image: Image.Image) -> Image.Image:
    image = image.convert("RGBA")
    pix = image.load()
    width, height = image.size
    seen = bytearray(width * height)
    queue: deque[tuple[int, int]] = deque()

    def add_if_background(x: int, y: int) -> None:
        index = y * width + x
        if seen[index] or not is_green_background(pix[x, y]):
            return
        seen[index] = 1
        queue.append((x, y))

    for x in range(width):
        add_if_background(x, 0)
        add_if_background(x, height - 1)
    for y in range(height):
        add_if_background(0, y)
        add_if_background(width - 1, y)

    while queue:
        x, y = queue.popleft()
        if x > 0:
            add_if_background(x - 1, y)
        if x + 1 < width:
            add_if_background(x + 1, y)
        if y > 0:
            add_if_background(x, y - 1)
        if y + 1 < height:
            add_if_background(x, y + 1)

    for y in range(height):
        for x in range(width):
            if seen[y * width + x]:
                r, g, b, _ = pix[x, y]
                pix[x, y] = (r, g, b, 0)
            else:
                r, g, b, a = pix[x, y]
                if a and is_green_background((r, g, b, a)):
                    pix[x, y] = (r, min(g, max(r, b) + 14), b, a)
    return image


def remove_large_green_regions(image: Image.Image, min_area: int = 2048) -> Image.Image:
    image = image.convert("RGBA")
    pix = image.load()
    width, height = image.size
    seen = bytearray(width * height)

    for y in range(height):
        for x in range(width):
            index = y * width + x
            if seen[index] or not is_green_background(pix[x, y]):
                continue
            seen[index] = 1
            queue: deque[tuple[int, int]] = deque([(x, y)])
            component: list[tuple[int, int]] = []
            while queue:
                point_x, point_y = queue.popleft()
                component.append((point_x, point_y))
                for next_x, next_y in (
                    (point_x - 1, point_y),
                    (point_x + 1, point_y),
                    (point_x, point_y - 1),
                    (point_x, point_y + 1),
                ):
                    if next_x < 0 or next_x >= width or next_y < 0 or next_y >= height:
                        continue
                    next_index = next_y * width + next_x
                    if seen[next_index] or not is_green_background(pix[next_x, next_y]):
                        continue
                    seen[next_index] = 1
                    queue.append((next_x, next_y))
            if len(component) >= min_area:
                for point_x, point_y in component:
                    r, g, b, _ = pix[point_x, point_y]
                    pix[point_x, point_y] = (r, g, b, 0)
    return image


def crop_visible(image: Image.Image, padding: int = 16) -> Image.Image:
    bbox = image.getchannel("A").getbbox()
    if bbox is None:
        return image
    left, top, right, bottom = bbox
    return image.crop(
        (
            max(0, left - padding),
            max(0, top - padding),
            min(image.width, right + padding),
            min(image.height, bottom + padding),
        )
    )


def fit_icon(image: Image.Image, target_size: tuple[int, int], icon_height_ratio: float) -> Image.Image:
    cropped = crop_visible(image)
    canvas = Image.new("RGBA", target_size, (0, 0, 0, 0))
    max_width = max(1, int(target_size[0] * 0.96))
    max_height = max(1, int(target_size[1] * icon_height_ratio))
    cropped.thumbnail((max_width, max_height), Image.Resampling.LANCZOS)
    x = (target_size[0] - cropped.width) // 2
    y = max(0, int((target_size[1] * icon_height_ratio - cropped.height) * 0.5))
    canvas.alpha_composite(cropped, (x, y))
    return canvas


def fit_cutout(image: Image.Image, target_size: tuple[int, int]) -> Image.Image:
    cropped = crop_visible(image, 0)
    canvas = Image.new("RGBA", target_size, (0, 0, 0, 0))
    cropped.thumbnail(target_size, Image.Resampling.LANCZOS)
    x = (target_size[0] - cropped.width) // 2
    y = (target_size[1] - cropped.height) // 2
    canvas.alpha_composite(cropped, (x, y))
    return canvas


def remove_background(image: Image.Image, method: str) -> Image.Image:
    if method == "ai":
        return remove_with_ai(image)
    if method == "auto":
        try:
            return remove_with_ai(image)
        except Exception:
            return remove_green_background(image)
    return remove_green_background(image)


def prepare_cutouts(args: argparse.Namespace) -> bool:
    if not args.cutout:
        return False

    preview_items: list[tuple[str, Path]] = []
    for source, output, width_text, height_text in args.cutout:
        source_path = Path(source).expanduser().resolve()
        output_path = Path(output).expanduser().resolve()
        output_path.parent.mkdir(parents=True, exist_ok=True)
        image = Image.open(source_path)
        matted = remove_large_green_regions(image) if args.method == "green" else remove_background(image, args.method)
        prepared = fit_cutout(matted, (int(width_text), int(height_text)))
        prepared.save(output_path)
        preview_items.append((output_path.stem, output_path))

    build_preview(preview_items, Path(args.preview))
    print(f"Prepared {len(preview_items)} cutouts")
    print(f"Preview: {args.preview}")
    return True


def write_tres(project: Path, spec: IconSpec) -> None:
    tres_path = project / spec.tres_path
    resource_path = f"res://assets/replacement/menu2/{spec.output_name}.png"
    tres_path.write_text(
        "[gd_resource type=\"AtlasTexture\" load_steps=2 format=3]\n\n"
        f"[ext_resource type=\"Texture2D\" path=\"{resource_path}\" id=\"1_replacement\"]\n\n"
        "[resource]\n"
        "atlas = ExtResource(\"1_replacement\")\n"
        f"region = Rect2(0, 0, {spec.target_size[0]}, {spec.target_size[1]})\n",
        encoding="utf-8",
    )


def build_preview(items: list[tuple[str, Path]], output_path: Path) -> None:
    thumb = 116
    label_height = 24
    columns = 5
    rows = (len(items) + columns - 1) // columns
    preview = Image.new("RGBA", (columns * thumb, rows * (thumb + label_height)), (255, 255, 255, 255))
    draw = ImageDraw.Draw(preview)
    for index, (name, path) in enumerate(items):
        bg = Image.new("RGBA", (thumb, thumb), (245, 245, 245, 255))
        bg_draw = ImageDraw.Draw(bg)
        for y in range(0, thumb, 10):
            for x in range(0, thumb, 10):
                if (x // 10 + y // 10) % 2:
                    bg_draw.rectangle((x, y, x + 9, y + 9), fill=(225, 225, 225, 255))
        icon = Image.open(path).convert("RGBA")
        icon.thumbnail((thumb - 8, thumb - 8), Image.Resampling.LANCZOS)
        bg.alpha_composite(icon, ((thumb - icon.width) // 2, (thumb - icon.height) // 2))
        x = (index % columns) * thumb
        y = (index // columns) * (thumb + label_height)
        preview.alpha_composite(bg, (x, y))
        draw.text((x + 4, y + thumb + 3), name[:18], fill=(0, 0, 0, 255))
    preview.save(output_path)


def main() -> None:
    args = parse_args()
    if args.install_ai:
        maybe_install_ai()
    if prepare_cutouts(args):
        return

    project = Path(args.project).resolve()
    input_dir = Path(args.input_dir).resolve()
    output_dir = project / "assets" / "replacement" / "menu2"
    output_dir.mkdir(parents=True, exist_ok=True)

    preview_items: list[tuple[str, Path]] = []
    for spec in ICON_SPECS:
        source_path = input_dir / spec.source
        if not source_path.exists():
            raise FileNotFoundError(source_path)
        image = Image.open(source_path)
        matted = remove_background(image, args.method)

        output = fit_icon(matted, spec.target_size, spec.icon_height_ratio)
        output_path = output_dir / f"{spec.output_name}.png"
        output.save(output_path)
        write_tres(project, spec)
        preview_items.append((spec.output_name, output_path))

    build_preview(preview_items, Path(args.preview))
    print(f"Prepared {len(ICON_SPECS)} icons")
    print(f"Preview: {args.preview}")


if __name__ == "__main__":
    main()