#!/usr/bin/env python3
"""
Compile text-grid part sources into palette-key PNGs.

Sources live in parts_src/<slot>/<name>.txt and compile to parts/<slot>/<name>.png.
Every non-blank line is one pixel row; lines starting with '#' are comments. All rows
must be the same width.

Legend
  .  transparent        o  outline
  1  main_light         2  main_base         3  main_dark
  a  accent_light       b  accent_base       c  accent_dark
  x  type_light         y  type_base         z  type_dark
  w  white              p  pupil
  P  pivot   T  tail   B  back   A  aura   ^  head_top

Layout sources (parts_src/<slot>/<name>.layout) assemble stamps onto a fixed canvas,
used for the 84x84 auras:
  size 84 84
  stamp <stamp-name> <x> <y> [flip]     # parts_src/<slot>/stamps/<stamp-name>.txt, top-left at x,y
  marker <char> <x> <y>                 # e.g. "marker P 42 42"
Later stamps paint over earlier ones; '.' in a stamp never paints.

Usage
  python build_parts.py                 # build everything
  python build_parts.py body/quad_slim  # build some parts
  python build_parts.py mirror FILE     # print FILE + its horizontal mirror (shared center column)
  python build_parts.py mirror FILE --gap N   # mirror with N transparent columns between halves
  python build_parts.py flip FILE       # print FILE flipped horizontally
"""
import os
import sys
from pathlib import Path

from PIL import Image

from generate_sprite import KEYS, MARKERS, PUPIL, EYE_WHITE

# PARTS_SRC / PARTS_DIR let experiments build into a separate library (e.g. proto/).
SRC = Path(os.environ.get("PARTS_SRC", "parts_src"))
OUT = Path(os.environ.get("PARTS_DIR", "parts"))

LEGEND = {
    "o": KEYS["outline"],
    "1": KEYS["main_light"], "2": KEYS["main_base"], "3": KEYS["main_dark"],
    "a": KEYS["accent_light"], "b": KEYS["accent_base"], "c": KEYS["accent_dark"],
    "x": KEYS["type_light"], "y": KEYS["type_base"], "z": KEYS["type_dark"],
    "w": EYE_WHITE, "p": PUPIL,
    "P": MARKERS["pivot"], "T": MARKERS["tail"], "B": MARKERS["back"],
    "A": MARKERS["aura"], "^": MARKERS["head_top"],
}


class SourceError(Exception):
    pass


def read_grid(path):
    rows = []
    for n, line in enumerate(Path(path).read_text().splitlines(), 1):
        line = line.rstrip()
        if not line or line.startswith("#"):
            continue
        bad = sorted({c for c in line if c != "." and c not in LEGEND})
        if bad:
            raise SourceError(f"{path}:{n}: unknown characters {bad}")
        rows.append(line)
    if not rows:
        raise SourceError(f"{path}: empty grid")
    widths = {len(r) for r in rows}
    if len(widths) != 1:
        for n, r in enumerate(rows):
            if len(r) != len(rows[0]):
                raise SourceError(f"{path}: row {n} is {len(r)} wide, row 0 is {len(rows[0])}")
    return rows


def paint(canvas, rows, x0, y0, flip=False):
    for y, row in enumerate(rows):
        if flip:
            row = row[::-1]
        for x, c in enumerate(row):
            if c == ".":
                continue
            px, py = x0 + x, y0 + y
            if not (0 <= px < len(canvas[0]) and 0 <= py < len(canvas)):
                raise SourceError(f"stamp pixel ({px},{py}) outside canvas")
            canvas[py][px] = c


def read_layout(path):
    path = Path(path)
    canvas = None
    for n, line in enumerate(path.read_text().splitlines(), 1):
        parts = line.split("#")[0].split()
        if not parts:
            continue
        cmd, args = parts[0], parts[1:]
        try:
            if cmd == "size":
                w, h = int(args[0]), int(args[1])
                canvas = [["."] * w for _ in range(h)]
            elif cmd == "stamp":
                rows = read_grid(path.parent / "stamps" / f"{args[0]}.txt")
                paint(canvas, rows, int(args[1]), int(args[2]), flip=len(args) > 3 and args[3] == "flip")
            elif cmd == "marker":
                paint(canvas, [args[0]], int(args[1]), int(args[2]))
            else:
                raise SourceError(f"unknown command {cmd!r}")
        except (IndexError, ValueError, TypeError) as e:
            raise SourceError(f"{path}:{n}: {e}")
    if canvas is None:
        raise SourceError(f"{path}: missing 'size'")
    return ["".join(r) for r in canvas]


def to_image(rows):
    h, w = len(rows), len(rows[0])
    img = Image.new("RGBA", (w, h), (0, 0, 0, 0))
    px = img.load()
    for y, row in enumerate(rows):
        for x, c in enumerate(row):
            if c != ".":
                px[x, y] = LEGEND[c] + (255,)
    return img


def sources():
    for p in sorted(SRC.glob("*/*")):
        if p.suffix in (".txt", ".layout"):
            yield p


def build(path):
    rows = read_layout(path) if path.suffix == ".layout" else read_grid(path)
    out = OUT / path.parent.name / f"{path.stem}.png"
    out.parent.mkdir(parents=True, exist_ok=True)
    to_image(rows).save(out)
    return out, len(rows[0]), len(rows)


def main(argv):
    if argv and argv[0] in ("mirror", "flip"):
        rows = read_grid(argv[1])
        if argv[0] == "flip":
            print("\n".join(r[::-1] for r in rows))
        else:
            gap = int(argv[3]) if len(argv) > 3 and argv[2] == "--gap" else None
            for r in rows:
                print(r + "." * gap + r[::-1] if gap is not None else r + r[-2::-1])
        return 0
    wanted = set(argv)
    errors = 0
    built = 0
    for path in sources():
        key = f"{path.parent.name}/{path.stem}"
        if wanted and key not in wanted:
            continue
        try:
            out, w, h = build(path)
            built += 1
            if wanted:
                print(f"{out} {w}x{h}")
        except SourceError as e:
            print(f"ERROR {e}", file=sys.stderr)
            errors += 1
    print(f"built {built} part(s), {errors} error(s)")
    return 1 if errors else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))