from __future__ import annotations

import colorsys
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont


REPO_ROOT = Path(__file__).resolve().parents[1]
SOURCE_PATH = (
    REPO_ROOT
    / "public"
    / "branding"
    / "taonier-logo-magic-dot-concepts"
    / "taonier-magic-dot-squish.png"
)
OUTPUT_DIR = (
    REPO_ROOT / "public" / "branding" / "taonier-logo-squish-recolor-variants"
)


VARIANTS = [
    {
        "id": "taonier-squish-recolor-original-plus",
        "title": "原版提亮",
        "top": ("#ff3f74", "#ff5a8c"),
        "bottom": ("#10c6b1", "#19d5b8"),
        "star": ("#ffd249", "#ffc13c"),
        "background": "#fffdf8",
        "saturation": 1.04,
        "value": 1.02,
    },
    {
        "id": "taonier-squish-recolor-candy-mint",
        "title": "糖果薄荷",
        "top": ("#ff5aa0", "#ff7786"),
        "bottom": ("#1fd3c2", "#46e0cc"),
        "star": ("#ffe071", "#ffc64b"),
        "background": "#fffafd",
        "saturation": 1.02,
        "value": 1.05,
    },
    {
        "id": "taonier-squish-recolor-peach-jelly",
        "title": "桃桃果冻",
        "top": ("#ff6b8d", "#ff8b72"),
        "bottom": ("#30cfb7", "#72dec5"),
        "star": ("#ffe586", "#ffc75c"),
        "background": "#fffaf2",
        "saturation": 0.96,
        "value": 1.07,
    },
    {
        "id": "taonier-squish-recolor-pop-bright",
        "title": "亮彩出圈",
        "top": ("#ff2f82", "#ff4faf"),
        "bottom": ("#00c5b9", "#00d8e8"),
        "star": ("#fff15a", "#ffc735"),
        "background": "#fbfbff",
        "saturation": 1.10,
        "value": 1.03,
    },
    {
        "id": "taonier-squish-recolor-coral-soda",
        "title": "珊瑚苏打",
        "top": ("#ff674d", "#ff8372"),
        "bottom": ("#17c5a9", "#55dabc"),
        "star": ("#ffe26f", "#ffbe43"),
        "background": "#fff9ee",
        "saturation": 0.98,
        "value": 1.05,
    },
    {
        "id": "taonier-squish-recolor-bubble-q",
        "title": "泡泡Q感",
        "top": ("#ff68ba", "#ff77a0"),
        "bottom": ("#35d8c9", "#73e7d8"),
        "star": ("#fff08c", "#ffd35c"),
        "background": "#fffaff",
        "saturation": 0.92,
        "value": 1.09,
    },
]


def hex_to_rgb(value: str) -> tuple[int, int, int]:
    value = value.removeprefix("#")
    return tuple(int(value[index : index + 2], 16) for index in (0, 2, 4))


def blend_rgb(
    first: tuple[int, int, int], second: tuple[int, int, int], amount: float
) -> tuple[int, int, int]:
    amount = max(0.0, min(1.0, amount))
    return tuple(round(first[index] * (1 - amount) + second[index] * amount) for index in range(3))


def classify_pixel(red: int, green: int, blue: int) -> str | None:
    hue, saturation, value = colorsys.rgb_to_hsv(red / 255, green / 255, blue / 255)

    if value < 0.42 or saturation < 0.09:
        return None
    if red > 145 and green > 105 and blue < 170 and red > blue + 20 and green > blue + 16:
        return "star"
    if red > 140 and red > green + 18 and red > blue + 12 and (hue < 0.08 or hue > 0.9):
        return "top"
    if green > 105 and blue > 88 and green > red + 15 and blue > red + 6 and 0.36 <= hue <= 0.58:
        return "bottom"
    return None


def remap_color(
    red: int,
    green: int,
    blue: int,
    x: int,
    y: int,
    width: int,
    height: int,
    group: str,
    variant: dict[str, object],
) -> tuple[int, int, int]:
    hue, saturation, value = colorsys.rgb_to_hsv(red / 255, green / 255, blue / 255)
    palette = variant[group]
    assert isinstance(palette, tuple)
    start = hex_to_rgb(palette[0])
    end = hex_to_rgb(palette[1])

    if group == "top":
        gradient_position = 0.68 * (x / width) + 0.32 * (y / height)
    elif group == "bottom":
        gradient_position = 0.42 * (x / width) + 0.58 * (y / height)
    else:
        gradient_position = 0.18 * (x / width) + 0.82 * (y / height)

    target = blend_rgb(start, end, gradient_position)
    target_hue, target_saturation, target_value = colorsys.rgb_to_hsv(
        target[0] / 255, target[1] / 255, target[2] / 255
    )

    saturation_boost = float(variant["saturation"])
    value_boost = float(variant["value"])
    new_saturation = max(0.0, min(1.0, target_saturation * saturation_boost))
    # 原图本身有轻微明暗变化，这里保留它，让换色后仍然像同一个软泥形体。
    new_value = max(0.0, min(1.0, target_value * (0.78 + value * 0.22) * value_boost))
    recolored = colorsys.hsv_to_rgb(target_hue, new_saturation, new_value)

    background = hex_to_rgb(str(variant["background"]))
    edge_coverage = max(0.0, min(1.0, (saturation - 0.08) / 0.55))
    if group == "star":
        edge_coverage = max(0.0, min(1.0, (saturation - 0.06) / 0.5))

    foreground = tuple(round(channel * 255) for channel in recolored)
    return blend_rgb(background, foreground, edge_coverage)


def recolor_variant(source: Image.Image, variant: dict[str, object]) -> Image.Image:
    image = source.convert("RGBA")
    width, height = image.size
    background = hex_to_rgb(str(variant["background"]))
    result = Image.new("RGBA", image.size, (*background, 255))

    source_pixels = image.load()
    result_pixels = result.load()
    for y in range(height):
        for x in range(width):
            red, green, blue, alpha = source_pixels[x, y]
            if alpha == 0:
                result_pixels[x, y] = (*background, 0)
                continue

            group = classify_pixel(red, green, blue)
            if group is None:
                if red > 238 and green > 238 and blue > 238:
                    result_pixels[x, y] = (*background, alpha)
                else:
                    result_pixels[x, y] = (red, green, blue, alpha)
                continue

            new_red, new_green, new_blue = remap_color(
                red, green, blue, x, y, width, height, group, variant
            )
            result_pixels[x, y] = (new_red, new_green, new_blue, alpha)

    return result


def load_font(size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    candidates = [
        Path("C:/Windows/Fonts/msyh.ttc"),
        Path("C:/Windows/Fonts/simhei.ttf"),
        Path("C:/Windows/Fonts/simsun.ttc"),
    ]
    for candidate in candidates:
        if candidate.exists():
            return ImageFont.truetype(str(candidate), size)
    return ImageFont.load_default()


def paste_cell(
    sheet: Image.Image,
    image: Image.Image,
    label: str,
    index: int,
    font: ImageFont.FreeTypeFont | ImageFont.ImageFont,
) -> None:
    cell = 310
    label_height = 54
    gap = 28
    columns = 4
    row = index // columns
    column = index % columns
    x = gap + column * (cell + gap)
    y = gap + row * (cell + label_height + gap)

    draw = ImageDraw.Draw(sheet)
    thumb = image.resize((cell, cell), Image.Resampling.LANCZOS)
    sheet.alpha_composite(thumb, (x, y))
    draw.rounded_rectangle(
        (x, y + cell, x + cell, y + cell + label_height),
        radius=10,
        fill=(255, 253, 248, 255),
    )
    text = f"{index:02d} {label}"
    text_box = draw.textbbox((0, 0), text, font=font)
    text_x = x + (cell - (text_box[2] - text_box[0])) / 2
    text_y = y + cell + (label_height - (text_box[3] - text_box[1])) / 2 - 2
    draw.text((text_x, text_y), text, fill=(50, 42, 36, 255), font=font)


def build_contact_sheet(original: Image.Image, outputs: list[tuple[dict[str, object], Image.Image]]) -> Image.Image:
    cell = 310
    label_height = 54
    gap = 28
    columns = 4
    rows = 2
    width = columns * cell + (columns + 1) * gap
    height = rows * (cell + label_height) + (rows + 1) * gap
    sheet = Image.new("RGBA", (width, height), (246, 242, 235, 255))
    font = load_font(22)

    paste_cell(sheet, original.convert("RGBA"), "原参考", 0, font)
    for index, (variant, image) in enumerate(outputs, start=1):
        paste_cell(sheet, image, str(variant["title"]), index, font)
    return sheet


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    source = Image.open(SOURCE_PATH)
    outputs: list[tuple[dict[str, object], Image.Image]] = []

    for variant in VARIANTS:
        image = recolor_variant(source, variant)
        image.save(OUTPUT_DIR / f"{variant['id']}.png")
        outputs.append((variant, image))

    contact_sheet = build_contact_sheet(source, outputs)
    contact_sheet.convert("RGB").save(OUTPUT_DIR / "taonier-squish-recolor-contact-sheet.png")

    print(
        {
            "ok": True,
            "output_dir": str(OUTPUT_DIR),
            "files": [f"{variant['id']}.png" for variant in VARIANTS]
            + ["taonier-squish-recolor-contact-sheet.png"],
        }
    )


if __name__ == "__main__":
    main()