#!/usr/bin/env python3
"""Generate compact jxl-art DSL links from rendered flag previews."""

import argparse
import base64
import json
import zlib
from pathlib import Path

from PIL import Image

ART_BASE = "https://www.januschka.com/jxl-art/?zcode="
WIDTH = 192
HEIGHT = 144
GRID_WIDTH = 48
GRID_HEIGHT = 36
PALETTE_COLORS = 16


def leaf_lines(pixel, alpha):
    red, green, blue, opacity = pixel
    if alpha:
        color = color_lines(red, green, blue, 1)
        return ["if c > 2", f"  - Set {opacity}"] + color
    return color_lines(red, green, blue, 0)


def color_lines(red, green, blue, indent):
    prefix = "  " * indent
    if red == green == blue:
        return [f"{prefix}- Set {red}"]
    if red == green:
        return [
            f"{prefix}if c > 1",
            f"{prefix}  - Set {blue}",
            f"{prefix}  - Set {red}",
        ]
    if green == blue:
        return [
            f"{prefix}if c > 0",
            f"{prefix}  - Set {green}",
            f"{prefix}  - Set {red}",
        ]
    return [
        f"{prefix}if c > 1",
        f"{prefix}  - Set {blue}",
        f"{prefix}  if c > 0",
        f"{prefix}    - Set {green}",
        f"{prefix}    - Set {red}",
    ]


def indent(lines):
    return ["  " + line for line in lines]


def split_tree(parts, axis, scale):
    if len(parts) == 1:
        return parts[0][2]
    middle = len(parts) // 2
    lower = parts[:middle]
    upper = parts[middle:]
    threshold = upper[0][0] * scale - 1
    return [f"if {axis} > {threshold}"] + indent(
        split_tree(upper, axis, scale)
    ) + indent(split_tree(lower, axis, scale))


def row_tree(row, alpha):
    runs = []
    start = 0
    current = row[0]
    for x, pixel in enumerate(row[1:], 1):
        if pixel != current:
            runs.append((start, x, leaf_lines(current, alpha)))
            start = x
            current = pixel
    runs.append((start, len(row), leaf_lines(current, alpha)))
    return split_tree(runs, "x", WIDTH // GRID_WIDTH)


def make_art(image):
    preview = image.resize((GRID_WIDTH, GRID_HEIGHT), Image.Resampling.NEAREST)
    source_alpha = preview.getchannel("A")
    preview = preview.convert("RGB").quantize(
        colors=PALETTE_COLORS,
        method=Image.Quantize.FASTOCTREE,
        dither=Image.Dither.NONE,
    ).convert("RGBA")
    preview.putalpha(source_alpha)
    pixels = list(preview.getdata())
    alpha = any(pixel[3] != 255 for pixel in pixels)
    if alpha:
        pixels = [
            (red, green, blue, 255 if opacity >= 128 else 0)
            for red, green, blue, opacity in pixels
        ]
    expected = Image.new("RGBA", (GRID_WIDTH, GRID_HEIGHT))
    expected.putdata(pixels)
    expected = expected.resize((WIDTH, HEIGHT), Image.Resampling.NEAREST)
    rows = [
        pixels[y * GRID_WIDTH : (y + 1) * GRID_WIDTH]
        for y in range(GRID_HEIGHT)
    ]
    bands = []
    start = 0
    current = rows[0]
    for y, row in enumerate(rows[1:], 1):
        if row != current:
            bands.append((start, y, row_tree(current, alpha)))
            start = y
            current = row
    bands.append((start, len(rows), row_tree(current, alpha)))
    header = ["Bitdepth 8", f"Width {WIDTH}", f"Height {HEIGHT}", "RCT 0"]
    if alpha:
        header.append("Alpha")
    code = "\n".join(header + [""] + split_tree(bands, "y", HEIGHT // GRID_HEIGHT))
    return code, expected


def make_code(image):
    return make_art(image)[0]


def encode_url(code):
    compressed = zlib.compress(code.encode())
    encoded = base64.urlsafe_b64encode(compressed).decode().rstrip("=")
    return ART_BASE + encoded


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--preview-dir", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, default=Path("manifest.json"))
    parser.add_argument("--code-dir", type=Path, default=Path("art-code"))
    args = parser.parse_args()

    items = json.loads(args.manifest.read_text())
    args.code_dir.mkdir(parents=True, exist_ok=True)
    lengths = []
    overrides_dir = Path(__file__).parent / "art-overrides"
    for item in items:
        image = Image.open(args.preview_dir / f"{item['code']}.png").convert("RGBA")
        override = overrides_dir / f"{item['code']}.tree"
        code = override.read_text().strip() if override.exists() else make_code(image)
        (args.code_dir / f"{item['code']}.tree").write_text(code + "\n")
        item["art"] = encode_url(code)
        item["artCodeBytes"] = len(code.encode())
        lengths.append((len(item["art"]), item["code"]))
    args.manifest.write_text(json.dumps(items, ensure_ascii=True, separators=(",", ":")) + "\n")
    print(f"generated {len(items)} links; longest URL: {max(lengths)}")


if __name__ == "__main__":
    main()