from __future__ import annotations

import json
import math
import struct
import zlib
from pathlib import Path

WIDTH = 256
HEIGHT = 256
STATES = ("idle", "listening", "recording", "thinking", "speaking", "error")


def blend(dst: tuple[int, int, int, int], src: tuple[int, int, int, int]) -> tuple[int, int, int, int]:
    sr, sg, sb, sa = src
    dr, dg, db, da = dst
    alpha = sa / 255
    out_a = int(sa + da * (1 - alpha))
    if out_a == 0:
        return 0, 0, 0, 0
    return (
        int((sr * alpha + dr * (da / 255) * (1 - alpha)) / (out_a / 255)),
        int((sg * alpha + dg * (da / 255) * (1 - alpha)) / (out_a / 255)),
        int((sb * alpha + db * (da / 255) * (1 - alpha)) / (out_a / 255)),
        out_a,
    )


def circle(pixels: list[list[tuple[int, int, int, int]]], cx: int, cy: int, radius: int, color: tuple[int, int, int, int]) -> None:
    for y in range(max(0, cy - radius), min(HEIGHT, cy + radius + 1)):
        for x in range(max(0, cx - radius), min(WIDTH, cx + radius + 1)):
            dist = math.hypot(x - cx, y - cy)
            if dist <= radius:
                shade = max(0.72, 1 - dist / (radius * 2.8))
                shaded = (int(color[0] * shade), int(color[1] * shade), int(color[2] * shade), color[3])
                pixels[y][x] = blend(pixels[y][x], shaded)


def line_circle(pixels: list[list[tuple[int, int, int, int]]], cx: int, cy: int, radius: int, color: tuple[int, int, int, int]) -> None:
    for y in range(max(0, cy - radius - 2), min(HEIGHT, cy + radius + 3)):
        for x in range(max(0, cx - radius - 2), min(WIDTH, cx + radius + 3)):
            dist = math.hypot(x - cx, y - cy)
            if radius - 2 <= dist <= radius + 2:
                pixels[y][x] = blend(pixels[y][x], color)


def rect(pixels: list[list[tuple[int, int, int, int]]], x0: int, y0: int, x1: int, y1: int, color: tuple[int, int, int, int]) -> None:
    for y in range(max(0, y0), min(HEIGHT, y1)):
        for x in range(max(0, x0), min(WIDTH, x1)):
            pixels[y][x] = blend(pixels[y][x], color)


def draw_pet(state: str) -> list[list[tuple[int, int, int, int]]]:
    pixels = [[(0, 0, 0, 0) for _ in range(WIDTH)] for _ in range(HEIGHT)]
    accent = {
        "idle": (78, 142, 154, 255),
        "listening": (60, 122, 214, 255),
        "recording": (224, 90, 64, 255),
        "thinking": (142, 96, 210, 255),
        "speaking": (44, 165, 108, 255),
        "error": (220, 70, 70, 255),
    }[state]
    body = (236, 230, 210, 255)
    shadow = (70, 82, 86, 90)
    circle(pixels, 128, 142, 74, shadow)
    circle(pixels, 128, 132, 70, body)
    circle(pixels, 86, 112, 20, accent)
    circle(pixels, 170, 112, 20, accent)
    circle(pixels, 103, 132, 9, (35, 44, 48, 255))
    circle(pixels, 153, 132, 9, (35, 44, 48, 255))
    circle(pixels, 100, 129, 3, (255, 255, 255, 220))
    circle(pixels, 150, 129, 3, (255, 255, 255, 220))
    rect(pixels, 92, 190, 111, 207, accent)
    rect(pixels, 145, 190, 164, 207, accent)
    if state == "speaking":
        circle(pixels, 128, 158, 15, (45, 48, 52, 255))
        circle(pixels, 128, 163, 7, (255, 130, 150, 255))
        line_circle(pixels, 198, 132, 18, accent)
        line_circle(pixels, 205, 132, 30, accent)
    elif state == "thinking":
        circle(pixels, 128, 158, 12, (45, 48, 52, 255))
        rect(pixels, 116, 153, 140, 160, body)
        circle(pixels, 180, 70, 7, accent)
        circle(pixels, 197, 52, 11, accent)
        circle(pixels, 216, 35, 15, accent)
    elif state == "listening":
        rect(pixels, 116, 157, 140, 163, (45, 48, 52, 255))
        line_circle(pixels, 58, 132, 18, accent)
        line_circle(pixels, 51, 132, 30, accent)
    elif state == "recording":
        circle(pixels, 128, 160, 11, (45, 48, 52, 255))
        circle(pixels, 197, 58, 13, accent)
    elif state == "error":
        rect(pixels, 124, 149, 132, 168, (45, 48, 52, 255))
        circle(pixels, 128, 178, 4, (45, 48, 52, 255))
        rect(pixels, 190, 50, 198, 84, accent)
        circle(pixels, 194, 96, 5, accent)
    else:
        rect(pixels, 114, 157, 142, 162, (45, 48, 52, 255))
    return pixels


def chunk(kind: bytes, data: bytes) -> bytes:
    return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF)


def save_png(path: Path, pixels: list[list[tuple[int, int, int, int]]]) -> None:
    raw = bytearray()
    for row in pixels:
        raw.append(0)
        for rgba in row:
            raw.extend(bytes(rgba))
    data = b"\x89PNG\r\n\x1a\n"
    data += chunk(b"IHDR", struct.pack(">IIBBBBB", WIDTH, HEIGHT, 8, 6, 0, 0, 0))
    data += chunk(b"IDAT", zlib.compress(bytes(raw), 9))
    data += chunk(b"IEND", b"")
    path.write_bytes(data)


def main() -> int:
    root = Path("assets/pet")
    root.mkdir(parents=True, exist_ok=True)
    manifest = {"generated_by": "scripts/generate_pet_assets.py", "states": {}}
    for state in STATES:
        name = f"pet-{state}.png"
        save_png(root / name, draw_pet(state))
        manifest["states"][state] = name
    (root / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())