from pathlib import Path

from PIL import Image, ImageDraw, ImageFilter, ImageFont


ROOT = Path(__file__).resolve().parent
SOURCE = ROOT / "smart-canteen-unified-3d-overview-4k-16x9.png"
OUTPUT = ROOT / "smart-canteen-unified-3d-overview-labeled-4k-16x9.png"
FONT_PATH = "/System/Library/Fonts/STHeiti Medium.ttc"

GREEN = (17, 124, 68, 242)
WHITE = (255, 255, 255, 255)

LABELS = [
    ("仓储与验收区", (170, 40), (520, 255)),
    ("后厨加工与食安区", (1500, 42), (1850, 300)),
    ("运营管理中心", (3060, 115), (3290, 365)),
    ("前厅就餐与结算区", (1230, 1215), (1650, 1100)),
    ("办公与健康服务区", (3020, 970), (3335, 1215)),
]


def add_label(
    overlay: Image.Image,
    text: str,
    position: tuple[int, int],
    target: tuple[int, int],
) -> None:
    x, y = position
    label_font = ImageFont.truetype(FONT_PATH, 64)
    measure = ImageDraw.Draw(overlay)
    bbox = measure.textbbox((0, 0), text, font=label_font)
    width = bbox[2] - bbox[0] + 92
    height = 116
    center_x = x + width // 2
    bottom_y = y + height

    shadow = Image.new("RGBA", overlay.size, (0, 0, 0, 0))
    shadow_draw = ImageDraw.Draw(shadow)
    shadow_draw.rounded_rectangle(
        (x + 9, y + 10, x + width + 9, y + height + 10),
        radius=26,
        fill=(16, 43, 32, 92),
    )
    shadow = shadow.filter(ImageFilter.GaussianBlur(12))
    overlay.alpha_composite(shadow)

    draw = ImageDraw.Draw(overlay)
    draw.line((center_x, bottom_y, target[0], target[1]), fill=WHITE, width=20)
    draw.line((center_x, bottom_y, target[0], target[1]), fill=GREEN, width=12)
    draw.ellipse(
        (target[0] - 22, target[1] - 22, target[0] + 22, target[1] + 22),
        fill=GREEN,
        outline=WHITE,
        width=8,
    )
    draw.rounded_rectangle(
        (x, y, x + width, y + height),
        radius=26,
        fill=GREEN,
        outline=WHITE,
        width=6,
    )
    text_y = y + (height - (bbox[3] - bbox[1])) // 2 - bbox[1] - 1
    draw.text((x + 46, text_y), text, font=label_font, fill=WHITE)


def main() -> None:
    base = Image.open(SOURCE).convert("RGBA")
    if base.size != (3840, 2160):
        raise ValueError(f"Expected a 3840x2160 source, got {base.size}")
    overlay = Image.new("RGBA", base.size, (0, 0, 0, 0))
    for text, position, target in LABELS:
        add_label(overlay, text, position, target)
    Image.alpha_composite(base, overlay).convert("RGB").save(OUTPUT, "PNG", optimize=True)
    print(OUTPUT)


if __name__ == "__main__":
    main()