#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///

from __future__ import annotations

import math
import struct
import sys
import zlib
from pathlib import Path


PIXELS_PER_BAYER_CELL = 16
TOTAL_DOT_AREA = 0.5
SDF_RADIUS_SCALE = 2.4
BINARY_THRESHOLD = 2.0 / 3.0
BASE_BAYER_POINTS = (
    (0.0, 0.0),
    (0.5, 0.5),
    (0.5, 0.0),
    (0.0, 0.5),
)


def main(argv: list[str]) -> int:
    if len(argv) != 2:
        print(f"usage: {Path(argv[0]).name} N", file=sys.stderr)
        return 2

    try:
        order = int(argv[1])
    except ValueError:
        print("N must be an integer.", file=sys.stderr)
        return 2

    if order < 1 or order & (order - 1):
        print("N must be a positive power of two.", file=sys.stderr)
        return 2

    output_dir = Path.cwd() / f"SSFD_{order}x{order}"
    output_dir.mkdir(parents=True, exist_ok=True)
    clear_previous_outputs(output_dir)

    size = order * PIXELS_PER_BAYER_CELL
    points = build_bayer_points(order)
    min_dist_sq = [math.inf] * (size * size)
    sample_x = [((x + 0.5) / size) for x in range(size)]
    sample_y = [((y + 0.5) / size) for y in range(size)]

    for layer_index, point in enumerate(points, start=1):
        update_min_distances(min_dist_sq, sample_x, sample_y, size, point)
        write_layer(output_dir, min_dist_sq, size, layer_index)

    print(f"Wrote {len(points)} textures to {output_dir}")
    return 0


def clear_previous_outputs(output_dir: Path) -> None:
    for pattern in (
        "dots_L*.png",
        "dots_L*-sdf.png",
        "dots_L*.png.meta",
        "dots_L*-sdf.png.meta",
    ):
        for path in output_dir.glob(pattern):
            if path.is_file():
                path.unlink()


def build_bayer_points(order: int) -> list[tuple[float, float]]:
    if order == 1:
        return [(0.0, 0.0)]

    recursion = int(math.log2(order))
    points = list(BASE_BAYER_POINTS)
    for r in range(recursion - 1):
        count = len(points)
        offset = 0.5 ** (r + 1)
        for i in range(1, 4):
            dx, dy = BASE_BAYER_POINTS[i]
            for j in range(count):
                px, py = points[j]
                points.append((px + dx * offset, py + dy * offset))
    return points


def update_min_distances(
    min_dist_sq: list[float],
    sample_x: list[float],
    sample_y: list[float],
    size: int,
    point: tuple[float, float],
) -> None:
    px, py = point
    for y, sy in enumerate(sample_y):
        dy = wrap_delta(sy - py)
        dy_sq = dy * dy
        row_offset = y * size
        for x, sx in enumerate(sample_x):
            dx = wrap_delta(sx - px)
            dist_sq = dx * dx + dy_sq
            idx = row_offset + x
            if dist_sq < min_dist_sq[idx]:
                min_dist_sq[idx] = dist_sq


def write_layer(output_dir: Path, min_dist_sq: list[float], size: int, layer_index: int) -> None:
    dot_radius = math.sqrt((TOTAL_DOT_AREA / layer_index) / math.pi)
    scale = dot_radius * SDF_RADIUS_SCALE
    sdf_pixels = bytearray()
    binary_pixels = bytearray()

    for dist_sq in min_dist_sq:
        dist = math.sqrt(dist_sq)
        value = clamp01(1.0 - dist / scale)
        sdf_byte = round(value * 255.0)
        sdf_pixels.append(sdf_byte)
        binary_pixels.append(255 if value >= BINARY_THRESHOLD else 0)

    write_grayscale_png(output_dir / f"dots_L{layer_index}-sdf.png", size, size, sdf_pixels)
    write_grayscale_png(output_dir / f"dots_L{layer_index}.png", size, size, binary_pixels)


def wrap_delta(value: float) -> float:
    return (value + 0.5) % 1.0 - 0.5


def clamp01(value: float) -> float:
    if value <= 0.0:
        return 0.0
    if value >= 1.0:
        return 1.0
    return value


def write_grayscale_png(path: Path, width: int, height: int, pixels: bytearray) -> None:
    raw = bytearray()
    stride = width
    for y in range(height):
        start = y * stride
        raw.append(0)
        raw.extend(pixels[start : start + stride])

    compressed = zlib.compress(bytes(raw), level=9)
    with path.open("wb") as handle:
        handle.write(b"\x89PNG\r\n\x1a\n")
        handle.write(make_chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)))
        handle.write(make_chunk(b"IDAT", compressed))
        handle.write(make_chunk(b"IEND", b""))


def make_chunk(chunk_type: bytes, payload: bytes) -> bytes:
    return (
        struct.pack(">I", len(payload))
        + chunk_type
        + payload
        + struct.pack(">I", zlib.crc32(chunk_type + payload) & 0xFFFFFFFF)
    )


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))