from collections import deque
import math
from pathlib import Path

from PIL import Image, ImageChops, ImageDraw, ImageFont


REPO_ROOT = Path(__file__).resolve().parents[1]
REFERENCE_PATH = (
    REPO_ROOT
    / "public"
    / "branding"
    / "taonier-logo-punch-hole-concepts"
    / "taonier-punch-color-inlay.png"
)
OUTPUT_DIR = (
    REPO_ROOT
    / "public"
    / "branding"
    / "taonier-logo-ref04-locked-color-concepts"
)


def is_red(pixel):
    r, g, b = pixel
    return r > 160 and g < 155 and b < 145 and r - g > 45


def is_cyan(pixel):
    r, g, b = pixel
    return r < 120 and g > 125 and b > 125 and b - r > 65


def is_open_light(pixel):
    r, g, b = pixel
    lum = (r + g + b) // 3
    return lum > 174 and max(pixel) - min(pixel) < 105 and not is_red(pixel) and not is_cyan(pixel)


def colorize(pixel, target, category):
    r, g, b = pixel
    lum = (0.2126 * r + 0.7152 * g + 0.0722 * b) / 255
    if category == "dark":
        factor = 0.88 + min(lum, 0.42) * 0.44
    else:
        factor = 0.9 + lum * 0.2
    return tuple(max(0, min(255, round(channel * factor))) for channel in target)


def build_masks(image):
    width, height = image.size
    pixels = image.load()
    open_mask = bytearray(width * height)

    for y in range(height):
        for x in range(width):
            if is_open_light(pixels[x, y]):
                open_mask[y * width + x] = 1

    # 从画布边缘连通的浅色区域是外部背景；剩下的浅色闭合区域就是中孔。
    external_mask = bytearray(width * height)
    queue = deque()
    for x in range(width):
        for y in (0, height - 1):
            index = y * width + x
            if open_mask[index] and not external_mask[index]:
                external_mask[index] = 1
                queue.append((x, y))
    for y in range(height):
        for x in (0, width - 1):
            index = y * width + x
            if open_mask[index] and not external_mask[index]:
                external_mask[index] = 1
                queue.append((x, y))

    while queue:
        x, y = queue.popleft()
        for next_x, next_y in ((x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)):
            if 0 <= next_x < width and 0 <= next_y < height:
                index = next_y * width + next_x
                if open_mask[index] and not external_mask[index]:
                    external_mask[index] = 1
                    queue.append((next_x, next_y))

    hole_mask = bytearray(width * height)
    red_mask = bytearray(width * height)
    cyan_mask = bytearray(width * height)
    dark_mask = bytearray(width * height)
    for y in range(height):
        for x in range(width):
            index = y * width + x
            pixel = pixels[x, y]
            if open_mask[index] and not external_mask[index]:
                hole_mask[index] = 1
            elif not external_mask[index]:
                if is_red(pixel):
                    red_mask[index] = 1
                elif is_cyan(pixel):
                    cyan_mask[index] = 1
                elif not open_mask[index]:
                    dark_mask[index] = 1

    return {
        "dark": dark_mask,
        "red": red_mask,
        "cyan": cyan_mask,
        "hole": hole_mask,
    }


def mask_bounds(mask, width, height):
    xs = []
    ys = []
    for index, value in enumerate(mask):
        if value:
            xs.append(index % width)
            ys.append(index // width)
    return min(xs), min(ys), max(xs), max(ys)


def draw_center_content(size, hole_mask, variant):
    width, height = size
    hole_bounds = mask_bounds(hole_mask, width, height)
    left, top, right, bottom = hole_bounds
    center_x = (left + right) / 2
    center_y = (top + bottom) / 2
    hole_w = right - left
    hole_h = bottom - top

    scale = 4
    layer = Image.new("RGBA", (width * scale, height * scale), (0, 0, 0, 0))
    draw = ImageDraw.Draw(layer)

    def box(cx, cy, w, h):
        return [
            int((cx - w / 2) * scale),
            int((cy - h / 2) * scale),
            int((cx + w / 2) * scale),
            int((cy + h / 2) * scale),
        ]

    def rounded(cx, cy, w, h, radius, fill):
        draw.rounded_rectangle(box(cx, cy, w, h), radius=int(radius * scale), fill=fill)

    def ellipse(cx, cy, w, h, fill):
        draw.ellipse(box(cx, cy, w, h), fill=fill)

    def star(cx, cy, outer, inner, fill):
        points = []
        for index in range(10):
            angle = -90 + index * 36
            radius = outer if index % 2 == 0 else inner
            points.append(
                (
                    int((cx + radius * math.cos(math.radians(angle))) * scale),
                    int((cy + radius * math.sin(math.radians(angle))) * scale),
                )
            )
        draw.polygon(points, fill=fill)

    def sparkle(cx, cy, radius, fill, with_rays=False):
        points = []
        point_count = 128
        for index in range(point_count):
            theta = -math.pi / 2 + index * math.tau / point_count
            pulse = abs(math.cos(2 * theta)) ** 4.2
            current_radius = radius * (0.12 + 0.88 * pulse)
            points.append(
                (
                    int((cx + math.cos(theta) * current_radius * 0.78) * scale),
                    int((cy + math.sin(theta) * current_radius * 1.12) * scale),
                )
            )
        draw.polygon(points, fill=fill)
        if not with_rays:
            return

        ray_color = fill
        line_width = max(4, int(radius * 0.11 * scale))
        cap = line_width // 2

        def rounded_line(start, end):
            draw.line(
                (
                    int(start[0] * scale),
                    int(start[1] * scale),
                    int(end[0] * scale),
                    int(end[1] * scale),
                ),
                fill=ray_color,
                width=line_width,
            )
            for point in (start, end):
                draw.ellipse(
                    (
                        int(point[0] * scale) - cap,
                        int(point[1] * scale) - cap,
                        int(point[0] * scale) + cap,
                        int(point[1] * scale) + cap,
                    ),
                    fill=ray_color,
                )

        rounded_line((cx - radius * 1.48, cy - radius * 0.15), (cx - radius * 1.18, cy - radius * 0.08))
        rounded_line((cx - radius * 1.35, cy + radius * 0.5), (cx - radius * 1.1, cy + radius * 0.3))
        rounded_line((cx + radius * 1.12, cy - radius * 0.42), (cx + radius * 1.33, cy - radius * 0.64))
        rounded_line((cx + radius * 1.25, cy + radius * 0.08), (cx + radius * 1.52, cy + radius * 0.15))

    if variant == "cream_seed":
        rounded(center_x, center_y, hole_w * 0.42, hole_h * 0.34, 44, (244, 216, 166, 255))
    elif variant == "soft_dot":
        rounded(center_x, center_y, hole_w * 0.32, hole_h * 0.28, 36, (250, 219, 157, 255))
    elif variant == "double_piece":
        rounded(center_x - hole_w * 0.08, center_y + hole_h * 0.01, hole_w * 0.24, hole_h * 0.22, 30, (249, 202, 174, 255))
        rounded(center_x + hole_w * 0.13, center_y - hole_h * 0.02, hole_w * 0.22, hole_h * 0.21, 28, (143, 207, 205, 255))
    elif variant == "tiny_kernel":
        ellipse(center_x, center_y, hole_w * 0.26, hole_h * 0.24, (252, 223, 157, 255))
    elif variant == "filled_core":
        rounded(center_x, center_y, hole_w * 0.58, hole_h * 0.5, 58, (248, 231, 196, 255))
        ellipse(center_x - hole_w * 0.09, center_y + hole_h * 0.02, hole_w * 0.13, hole_h * 0.12, (240, 93, 82, 255))
        ellipse(center_x + hole_w * 0.1, center_y - hole_h * 0.02, hole_w * 0.13, hole_h * 0.12, (20, 183, 196, 255))
    elif variant == "clay_pearl":
        rounded(center_x, center_y, hole_w * 0.36, hole_h * 0.3, 40, (255, 226, 177, 255))
        ellipse(center_x + hole_w * 0.06, center_y - hole_h * 0.05, hole_w * 0.08, hole_h * 0.07, (255, 246, 220, 180))
    elif variant == "cream_star":
        star(center_x, center_y, min(hole_w, hole_h) * 0.17, min(hole_w, hole_h) * 0.075, (255, 223, 154, 255))
    elif variant == "small_star":
        star(center_x, center_y, min(hole_w, hole_h) * 0.135, min(hole_w, hole_h) * 0.06, (255, 231, 177, 255))
    elif variant == "soft_star_badge":
        rounded(center_x, center_y, hole_w * 0.38, hole_h * 0.34, 42, (255, 239, 207, 255))
        star(center_x, center_y, min(hole_w, hole_h) * 0.115, min(hole_w, hole_h) * 0.052, (238, 129, 80, 255))
    elif variant == "coral_star":
        star(center_x, center_y, min(hole_w, hole_h) * 0.15, min(hole_w, hole_h) * 0.067, (241, 108, 82, 255))
    elif variant == "mint_star":
        star(center_x, center_y, min(hole_w, hole_h) * 0.15, min(hole_w, hole_h) * 0.067, (78, 198, 183, 255))
    elif variant == "soft_sparkle":
        sparkle(center_x, center_y, min(hole_w, hole_h) * 0.18, (255, 205, 61, 255), True)
    elif variant == "small_sparkle":
        sparkle(center_x, center_y, min(hole_w, hole_h) * 0.145, (255, 214, 91, 255), True)
    elif variant == "bright_sparkle":
        sparkle(center_x, center_y, min(hole_w, hole_h) * 0.17, (255, 197, 43, 255), True)
    elif variant == "quiet_sparkle":
        sparkle(center_x, center_y, min(hole_w, hole_h) * 0.155, (255, 224, 139, 255), False)

    layer = layer.resize(size, Image.Resampling.LANCZOS)
    alpha = Image.frombytes("L", size, bytes(255 if value else 0 for value in hole_mask))
    layer_alpha = layer.getchannel("A")
    layer.putalpha(ImageChops.multiply(layer_alpha, alpha))
    return layer


VARIANTS = [
    {
        "id": "taonier-ref04-locked-warm-ink",
        "label": "01 warm",
        "dark": (63, 58, 53),
        "red": (243, 82, 69),
        "cyan": (14, 183, 198),
        "content": "cream_seed",
    },
    {
        "id": "taonier-ref04-locked-blue-ink",
        "label": "02 blue",
        "dark": (30, 39, 72),
        "red": (255, 89, 84),
        "cyan": (28, 181, 207),
        "content": "soft_dot",
    },
    {
        "id": "taonier-ref04-locked-plum-ink",
        "label": "03 plum",
        "dark": (69, 53, 72),
        "red": (255, 98, 86),
        "cyan": (34, 188, 198),
        "content": "double_piece",
    },
    {
        "id": "taonier-ref04-locked-green-ink",
        "label": "04 green",
        "dark": (11, 83, 78),
        "red": (255, 107, 88),
        "cyan": (68, 209, 192),
        "content": "tiny_kernel",
    },
    {
        "id": "taonier-ref04-locked-shrink-core",
        "label": "05 fill",
        "dark": (43, 43, 47),
        "red": (239, 84, 75),
        "cyan": (17, 178, 193),
        "content": "filled_core",
    },
    {
        "id": "taonier-ref04-locked-soft-charcoal",
        "label": "06 soft",
        "dark": (82, 76, 68),
        "red": (242, 105, 90),
        "cyan": (37, 188, 195),
        "content": "clay_pearl",
    },
]

STAR_OUTPUT_DIR = (
    REPO_ROOT
    / "public"
    / "branding"
    / "taonier-logo-ref04-warm-star-concepts"
)

STAR_VARIANTS = [
    {
        "id": "taonier-ref04-warm-star-terracotta",
        "label": "01 clay",
        "dark": (121, 76, 54),
        "red": (244, 86, 70),
        "cyan": (15, 184, 198),
        "content": "cream_star",
    },
    {
        "id": "taonier-ref04-warm-star-caramel",
        "label": "02 caramel",
        "dark": (142, 94, 51),
        "red": (247, 91, 73),
        "cyan": (13, 185, 196),
        "content": "small_star",
    },
    {
        "id": "taonier-ref04-warm-star-cocoa",
        "label": "03 cocoa",
        "dark": (89, 64, 47),
        "red": (240, 88, 72),
        "cyan": (17, 181, 194),
        "content": "soft_star_badge",
    },
    {
        "id": "taonier-ref04-warm-star-rust",
        "label": "04 rust",
        "dark": (111, 62, 54),
        "red": (249, 93, 75),
        "cyan": (15, 184, 198),
        "content": "cream_star",
    },
    {
        "id": "taonier-ref04-warm-star-olive",
        "label": "05 olive",
        "dark": (92, 81, 48),
        "red": (245, 94, 76),
        "cyan": (25, 185, 187),
        "content": "coral_star",
    },
    {
        "id": "taonier-ref04-warm-star-plum",
        "label": "06 plum",
        "dark": (95, 57, 66),
        "red": (250, 94, 81),
        "cyan": (26, 185, 197),
        "content": "mint_star",
    },
]

SPARKLE_OUTPUT_DIR = (
    REPO_ROOT
    / "public"
    / "branding"
    / "taonier-logo-ref04-warm-sparkle-concepts"
)

SPARKLE_V2_OUTPUT_DIR = (
    REPO_ROOT
    / "public"
    / "branding"
    / "taonier-logo-ref04-warm-sparkle-v2-concepts"
)

PALETTE_TRANSFER_OUTPUT_DIR = (
    REPO_ROOT
    / "public"
    / "branding"
    / "taonier-logo-ref04-palette-transfer"
)

SPARKLE_VARIANTS = [
    {
        "id": "taonier-ref04-warm-sparkle-terracotta",
        "label": "01 clay",
        "dark": (121, 76, 54),
        "red": (244, 86, 70),
        "cyan": (15, 184, 198),
        "content": "soft_sparkle",
    },
    {
        "id": "taonier-ref04-warm-sparkle-rust",
        "label": "02 rust",
        "dark": (111, 62, 54),
        "red": (249, 93, 75),
        "cyan": (15, 184, 198),
        "content": "soft_sparkle",
    },
    {
        "id": "taonier-ref04-warm-sparkle-caramel",
        "label": "03 caramel",
        "dark": (142, 94, 51),
        "red": (247, 91, 73),
        "cyan": (13, 185, 196),
        "content": "small_sparkle",
    },
    {
        "id": "taonier-ref04-warm-sparkle-cocoa",
        "label": "04 cocoa",
        "dark": (89, 64, 47),
        "red": (240, 88, 72),
        "cyan": (17, 181, 194),
        "content": "bright_sparkle",
    },
    {
        "id": "taonier-ref04-warm-sparkle-clay-quiet",
        "label": "05 quiet",
        "dark": (121, 76, 54),
        "red": (244, 86, 70),
        "cyan": (15, 184, 198),
        "content": "quiet_sparkle",
    },
    {
        "id": "taonier-ref04-warm-sparkle-plum",
        "label": "06 plum",
        "dark": (95, 57, 66),
        "red": (250, 94, 81),
        "cyan": (26, 185, 197),
        "content": "soft_sparkle",
    },
]

PALETTE_TRANSFER_VARIANTS = [
    {
        "id": "taonier-ref04-palette-transfer-warm-yellow-sparkle",
        "label": "transfer",
        "dark": (224, 162, 58),
        "red": (255, 113, 132),
        "cyan": (91, 213, 192),
        "content": "soft_sparkle",
    },
]


def apply_variant(reference, masks, variant):
    image = reference.copy().convert("RGBA")
    source = reference.convert("RGB")
    width, height = source.size
    source_pixels = source.load()
    result_pixels = image.load()

    for y in range(height):
        for x in range(width):
            index = y * width + x
            pixel = source_pixels[x, y]
            if masks["dark"][index]:
                result_pixels[x, y] = (*colorize(pixel, variant["dark"], "dark"), 255)
            elif masks["red"][index]:
                result_pixels[x, y] = (*colorize(pixel, variant["red"], "accent"), 255)
            elif masks["cyan"][index]:
                result_pixels[x, y] = (*colorize(pixel, variant["cyan"], "accent"), 255)

    content = draw_center_content(source.size, masks["hole"], variant["content"])
    return Image.alpha_composite(image, content).convert("RGB")


def build_contact_sheet(items, output_path):
    thumb = 260
    label_h = 34
    pad = 18
    cols = 4
    rows = (len(items) + cols - 1) // cols
    sheet_w = cols * thumb + (cols + 1) * pad
    sheet_h = rows * (thumb + label_h) + (rows + 1) * pad
    sheet = Image.new("RGB", (sheet_w, sheet_h), "#f7f3ea")
    draw = ImageDraw.Draw(sheet)
    try:
        font = ImageFont.truetype("arial.ttf", 18)
    except OSError:
        font = ImageFont.load_default()

    for index, (label, path) in enumerate(items):
        image = Image.open(path).convert("RGB")
        image.thumbnail((thumb, thumb), Image.Resampling.LANCZOS)
        row, col = divmod(index, cols)
        x = pad + col * (thumb + pad)
        y = pad + row * (thumb + label_h + pad)
        bg = Image.new("RGB", (thumb, thumb), "#fffaf1")
        bg.paste(image, ((thumb - image.width) // 2, (thumb - image.height) // 2))
        sheet.paste(bg, (x, y))
        draw.rectangle((x, y, x + thumb - 1, y + thumb - 1), outline="#ded5c6", width=1)
        bbox = draw.textbbox((0, 0), label, font=font)
        draw.text((x + (thumb - (bbox[2] - bbox[0])) // 2, y + thumb + 8), label, fill="#211f1c", font=font)

    sheet.save(output_path)


def generate_set(output_dir, variants, contact_name):
    output_dir.mkdir(parents=True, exist_ok=True)
    reference = Image.open(REFERENCE_PATH).convert("RGB")
    masks = build_masks(reference)
    contact_items = [("REF-04", REFERENCE_PATH)]

    for variant in variants:
        output_path = output_dir / f"{variant['id']}.png"
        apply_variant(reference, masks, variant).save(output_path)
        contact_items.append((variant["label"], output_path))

    build_contact_sheet(
        contact_items,
        output_dir / contact_name,
    )


def main():
    generate_set(
        OUTPUT_DIR,
        VARIANTS,
        "taonier-logo-ref04-locked-color-contact-sheet.png",
    )
    generate_set(
        STAR_OUTPUT_DIR,
        STAR_VARIANTS,
        "taonier-logo-ref04-warm-star-contact-sheet.png",
    )
    generate_set(
        SPARKLE_OUTPUT_DIR,
        SPARKLE_VARIANTS,
        "taonier-logo-ref04-warm-sparkle-contact-sheet.png",
    )
    generate_set(
        SPARKLE_V2_OUTPUT_DIR,
        SPARKLE_VARIANTS,
        "taonier-logo-ref04-warm-sparkle-v2-contact-sheet.png",
    )
    generate_set(
        PALETTE_TRANSFER_OUTPUT_DIR,
        PALETTE_TRANSFER_VARIANTS,
        "taonier-logo-ref04-palette-transfer-contact-sheet.png",
    )


if __name__ == "__main__":
    main()