#!/usr/bin/env python3
"""Build the deterministic Discord avatar test-canvas study.

One 900 x 1200 synthetic source is exported at 128, 256 and 512 square
test canvases in Fill and Fit modes. The generator writes PNG and JPG exports,
32 px square previews and 96 px circular previews, then reads every file back.

Mask counts map export or preview pixels back to exact source coordinates.
They record whether a marker or small-text pixel can survive the geometry. They
are not an OCR result, a readability score or a claim about Discord rendering.

There is no clock, randomness or network access. Python standard library plus
the local `cjpeg` and `djpeg` commands are required.
"""

import hashlib
import json
import math
import os
import struct
import subprocess
import tempfile
import zlib

HERE = os.path.dirname(os.path.abspath(__file__))
SOURCE = {"width": 900, "height": 1200, "file": "sources/source-900x1200.png"}
CANVASES = [128, 256, 512]
PREVIEW_SQUARE = 32
PREVIEW_CIRCLE = 96

BG = (27, 31, 43, 255)
GRID = (47, 55, 75, 255)
INK = (244, 239, 228, 255)
TEAL = (45, 212, 191, 255)
CORAL = (251, 113, 133, 255)
LAVENDER = (167, 139, 250, 255)
MARKER = (255, 45, 120, 255)
TEXT = (255, 214, 64, 255)
MARKER_SIZE = 36

# 5 x 7 bitmap glyphs keep the source reproducible without a font dependency.
FONT = {
    " ": ["00000"] * 7,
    "1": ["00100", "01100", "00100", "00100", "00100", "00100", "01110"],
    "2": ["01110", "10001", "00001", "00010", "00100", "01000", "11111"],
    "7": ["11111", "00001", "00010", "00100", "01000", "01000", "01000"],
    "A": ["01110", "10001", "10001", "11111", "10001", "10001", "10001"],
    "C": ["01111", "10000", "10000", "10000", "10000", "10000", "01111"],
    "D": ["11110", "10001", "10001", "10001", "10001", "10001", "11110"],
    "F": ["11111", "10000", "10000", "11110", "10000", "10000", "10000"],
    "I": ["01110", "00100", "00100", "00100", "00100", "00100", "01110"],
    "L": ["10000", "10000", "10000", "10000", "10000", "10000", "11111"],
    "M": ["10001", "11011", "10101", "10101", "10001", "10001", "10001"],
    "O": ["01110", "10001", "10001", "10001", "10001", "10001", "01110"],
    "P": ["11110", "10001", "10001", "11110", "10000", "10000", "10000"],
    "R": ["11110", "10001", "10001", "11110", "10100", "10010", "10001"],
    "S": ["01111", "10000", "10000", "01110", "00001", "00001", "11110"],
    "X": ["10001", "10001", "01010", "00100", "01010", "10001", "10001"],
}


def rel(path):
    return os.path.relpath(path, HERE).replace(os.sep, "/")


def canvas(width, height, colour=(0, 0, 0, 0)):
    row = bytes(colour) * width
    return [bytearray(row) for _ in range(height)]


def pixel(px, width, height, x, y, colour):
    if 0 <= x < width and 0 <= y < height:
        px[y][x * 4:x * 4 + 4] = bytes(colour)


def rect(px, width, height, x0, y0, w, h, colour):
    for y in range(max(0, y0), min(height, y0 + h)):
        for x in range(max(0, x0), min(width, x0 + w)):
            pixel(px, width, height, x, y, colour)


def disc(px, width, height, cx, cy, radius, colour):
    r2 = radius * radius
    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)):
            if (x - cx) ** 2 + (y - cy) ** 2 <= r2:
                pixel(px, width, height, x, y, colour)


def text(px, width, height, x, y, value, scale, colour):
    cursor = x
    for char in value:
        if char not in FONT:
            raise ValueError("missing glyph: " + repr(char))
        for gy, line in enumerate(FONT[char]):
            for gx, bit in enumerate(line):
                if bit == "1":
                    rect(px, width, height, cursor + gx * scale, y + gy * scale, scale, scale, colour)
        cursor += 6 * scale


def centered_text(px, width, height, y, value, scale, colour):
    x = (width - (len(value) * 6 * scale - scale)) // 2
    text(px, width, height, x, y, value, scale, colour)


def make_source():
    width, height = SOURCE["width"], SOURCE["height"]
    px = canvas(width, height, BG)

    for x in range(0, width, 60):
        rect(px, width, height, x, 0, 1, height, GRID)
    for y in range(0, height, 60):
        rect(px, width, height, 0, y, width, 1, GRID)
    rect(px, width, height, 0, 0, width, 5, INK)
    rect(px, width, height, 0, height - 5, width, 5, INK)
    rect(px, width, height, 0, 0, 5, height, INK)
    rect(px, width, height, width - 5, 0, 5, height, INK)

    # A simple portrait-like subject remains inside Fill's central square.
    disc(px, width, height, 450, 360, 132, INK)
    disc(px, width, height, 450, 360, 108, LAVENDER)
    disc(px, width, height, 402, 338, 16, BG)
    disc(px, width, height, 498, 338, 16, BG)
    rect(px, width, height, 405, 398, 90, 12, BG)
    rect(px, width, height, 258, 526, 384, 330, TEAL)
    rect(px, width, height, 304, 526, 292, 330, INK)
    rect(px, width, height, 348, 526, 204, 330, LAVENDER)
    rect(px, width, height, 430, 540, 40, 316, CORAL)

    # Four exact-colour markers sit outside Fill's retained 900 px-high square.
    corners = [
        (28, 28), (width - 28 - MARKER_SIZE, 28),
        (28, height - 28 - MARKER_SIZE), (width - 28 - MARKER_SIZE, height - 28 - MARKER_SIZE),
    ]
    for x, y in corners:
        rect(px, width, height, x, y, MARKER_SIZE, MARKER_SIZE, MARKER)

    centered_text(px, width, height, 930, "DISCORD", 4, INK)
    centered_text(px, width, height, 1080, "SMALL 7 PX", 1, TEXT)
    return px


def mask_pixels(px, width, height, colour):
    wanted = bytes(colour)
    return {
        (x, y)
        for y in range(height)
        for x in range(width)
        if bytes(px[y][x * 4:x * 4 + 4]) == wanted
    }


def resize(src, x0, y0, source_width, source_height, out_width, out_height):
    src_height, src_width = len(src), len(src[0]) // 4
    out = canvas(out_width, out_height)
    for y in range(out_height):
        sy = min(src_height - 1, max(0, int(math.floor(y0 + (y + 0.5) * source_height / out_height))))
        for x in range(out_width):
            sx = min(src_width - 1, max(0, int(math.floor(x0 + (x + 0.5) * source_width / out_width))))
            out[y][x * 4:x * 4 + 4] = src[sy][sx * 4:sx * 4 + 4]
    return out


def fill_export(src, size, shift=0.0):
    scale = max(size / SOURCE["width"], size / SOURCE["height"])
    kept = size / scale
    source_x = (SOURCE["width"] - kept) / 2
    source_y = (SOURCE["height"] - kept) / 2 + shift * SOURCE["height"]
    source_y = max(0.0, min(SOURCE["height"] - kept, source_y))
    out = resize(src, source_x, source_y, kept, kept, size, size)
    return out, {
        "mode": "fill",
        "scale": scale,
        "source_region_px": [kept, kept],
        "source_x": source_x,
        "source_y": source_y,
        "crop_area_px": SOURCE["width"] * SOURCE["height"] - kept * kept,
        "crop_share": 1 - (kept * kept) / (SOURCE["width"] * SOURCE["height"]),
        "pad_area_px": 0,
        "pad_share": 0,
    }


def fit_export(src, size):
    scale = min(size / SOURCE["width"], size / SOURCE["height"])
    drawn_w, drawn_h = SOURCE["width"] * scale, SOURCE["height"] * scale
    raster_w, raster_h = int(round(drawn_w)), int(round(drawn_h))
    x0, y0 = int(round((size - drawn_w) / 2)), int(round((size - drawn_h) / 2))
    drawn = resize(src, 0, 0, SOURCE["width"], SOURCE["height"], raster_w, raster_h)
    out = canvas(size, size)
    for y, row in enumerate(drawn):
        if 0 <= y0 + y < size:
            out[y0 + y][x0 * 4:x0 * 4 + len(row)] = row
    pad = size * size - drawn_w * drawn_h
    return out, {
        "mode": "fit",
        "scale": scale,
        "source_region_px": [SOURCE["width"], SOURCE["height"]],
        "source_x": 0,
        "source_y": 0,
        "draw_offset_px": [x0, y0],
        "drawn_raster_px": [raster_w, raster_h],
        "drawn_continuous_px": [drawn_w, drawn_h],
        "crop_area_px": 0,
        "crop_share": 0,
        "pad_area_px": pad,
        "pad_share": pad / (size * size),
    }


def export_to_source(geometry, size, x, y):
    if geometry["mode"] == "fill":
        kept_w, kept_h = geometry["source_region_px"]
        sx = int(math.floor(geometry["source_x"] + (x + 0.5) * kept_w / size))
        sy = int(math.floor(geometry["source_y"] + (y + 0.5) * kept_h / size))
    else:
        x0, y0 = geometry["draw_offset_px"]
        raster_w, raster_h = geometry["drawn_raster_px"]
        if not (x0 <= x < x0 + raster_w and y0 <= y < y0 + raster_h):
            return None
        sx = int(math.floor((x - x0 + 0.5) * SOURCE["width"] / raster_w))
        sy = int(math.floor((y - y0 + 0.5) * SOURCE["height"] / raster_h))
    if not (0 <= sx < SOURCE["width"] and 0 <= sy < SOURCE["height"]):
        return None
    return sx, sy


def geometry_masks(geometry, size, markers, small_text):
    marker_count = text_count = 0
    for y in range(size):
        for x in range(size):
            source = export_to_source(geometry, size, x, y)
            if source in markers:
                marker_count += 1
            if source in small_text:
                text_count += 1
    return {"marker_pixels": marker_count, "small_text_pixels": text_count}


def resize_preview(src, size, target):
    out = canvas(target, target)
    for y in range(target):
        sy = min(size - 1, max(0, int(math.floor((y + 0.5) * size / target))))
        for x in range(target):
            sx = min(size - 1, max(0, int(math.floor((x + 0.5) * size / target))))
            out[y][x * 4:x * 4 + 4] = src[sy][sx * 4:sx * 4 + 4]
    return out


def preview_masks(preview_px, preview_size, export_size, geometry, markers, small_text):
    marker_count = text_count = 0
    for y in range(preview_size):
        ey = int(math.floor((y + 0.5) * export_size / preview_size))
        for x in range(preview_size):
            if preview_px[y][x * 4 + 3] == 0:
                continue
            ex = int(math.floor((x + 0.5) * export_size / preview_size))
            source = export_to_source(geometry, export_size, ex, ey)
            if source in markers:
                marker_count += 1
            if source in small_text:
                text_count += 1
    return {"marker_pixels": marker_count, "small_text_pixels": text_count}


def circle_preview(src, size):
    out = canvas(PREVIEW_CIRCLE, PREVIEW_CIRCLE)
    radius = PREVIEW_CIRCLE / 2
    for y in range(PREVIEW_CIRCLE):
        sy = min(size - 1, max(0, int(math.floor((y + 0.5) * size / PREVIEW_CIRCLE))))
        for x in range(PREVIEW_CIRCLE):
            if (x + 0.5 - radius) ** 2 + (y + 0.5 - radius) ** 2 > radius * radius:
                continue
            sx = min(size - 1, max(0, int(math.floor((x + 0.5) * size / PREVIEW_CIRCLE))))
            out[y][x * 4:x * 4 + 4] = src[sy][sx * 4:sx * 4 + 4]
    return out


def write_png(path, pixels, width, height):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    raw = b"".join(b"\x00" + bytes(row) for row in pixels)

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

    with open(path, "wb") as stream:
        stream.write(b"\x89PNG\r\n\x1a\n")
        stream.write(chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)))
        stream.write(chunk(b"IDAT", zlib.compress(raw, 9)))
        stream.write(chunk(b"IEND", b""))


def read_png(path):
    data = open(path, "rb").read()
    if data[:8] != b"\x89PNG\r\n\x1a\n":
        raise ValueError("not a PNG: " + rel(path))
    offset, idat = 8, []
    width = height = colour_type = None
    while offset < len(data):
        length = struct.unpack(">I", data[offset:offset + 4])[0]
        kind = data[offset + 4:offset + 8]
        body = data[offset + 8:offset + 8 + length]
        expected = struct.unpack(">I", data[offset + 8 + length:offset + 12 + length])[0]
        if zlib.crc32(kind + body) & 0xFFFFFFFF != expected:
            raise ValueError("bad PNG chunk CRC: " + rel(path))
        if kind == b"IHDR":
            width, height = struct.unpack(">II", body[:8])
            colour_type = body[9]
        if kind == b"IDAT":
            idat.append(body)
        offset += length + 12
    if colour_type != 6:
        raise ValueError("expected RGBA PNG: " + rel(path))
    raster = zlib.decompress(b"".join(idat))
    stride = width * 4
    if len(raster) != height * (stride + 1):
        raise ValueError("decoded PNG length mismatch: " + rel(path))
    out = canvas(width, height)
    previous = bytearray(stride)
    position = 0
    filters = set()
    for y in range(height):
        filter_type = raster[position]
        filters.add(filter_type)
        position += 1
        row = bytearray(raster[position:position + stride])
        position += stride
        for i in range(stride):
            left = row[i - 4] if i >= 4 else 0
            above = previous[i]
            upper_left = previous[i - 4] if i >= 4 else 0
            if filter_type == 1:
                row[i] = (row[i] + left) & 255
            elif filter_type == 2:
                row[i] = (row[i] + above) & 255
            elif filter_type == 3:
                row[i] = (row[i] + ((left + above) // 2)) & 255
            elif filter_type == 4:
                p = left + above - upper_left
                pa, pb, pc = abs(p - left), abs(p - above), abs(p - upper_left)
                predictor = left if pa <= pb and pa <= pc else above if pb <= pc else upper_left
                row[i] = (row[i] + predictor) & 255
            elif filter_type != 0:
                raise ValueError("unsupported PNG filter: " + rel(path))
        out[y][:] = row
        previous = row
    return out, {
        "width": width,
        "height": height,
        "colour_type": "RGBA",
        "chunk_crc_ok": True,
        "filter_types": sorted(filters),
        "rgba_sha256": hashlib.sha256(b"".join(bytes(row) for row in out)).hexdigest(),
    }


def write_ppm(path, pixels, width, height):
    with open(path, "wb") as stream:
        stream.write(f"P6\n{width} {height}\n255\n".encode("ascii"))
        for row in pixels:
            rgb = bytearray(width * 3)
            for x in range(width):
                source = row[x * 4:x * 4 + 4]
                alpha = source[3] / 255
                rgb[x * 3:x * 3 + 3] = bytes(
                    int(round(source[channel] * alpha + 255 * (1 - alpha)))
                    for channel in range(3)
                )
            stream.write(rgb)


def read_ppm(path):
    data = open(path, "rb").read()
    magic, dimensions, maximum, raster = data.split(b"\n", 3)
    if magic != b"P6" or maximum != b"255":
        raise ValueError("unexpected djpeg PPM: " + path)
    width, height = map(int, dimensions.split())
    if len(raster) != width * height * 3:
        raise ValueError("decoded JPEG length mismatch: " + path)
    return {
        "width": width,
        "height": height,
        "rgb_sha256": hashlib.sha256(raster).hexdigest(),
        "pixel_count": width * height,
    }


def output_jpg(pixels, width, height, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with tempfile.TemporaryDirectory() as tmp:
        ppm = os.path.join(tmp, "source.ppm")
        write_ppm(ppm, pixels, width, height)
        subprocess.run(["cjpeg", "-quality", "92", "-outfile", path, ppm], check=True)
    data = open(path, "rb").read()
    with tempfile.TemporaryDirectory() as tmp:
        decoded = os.path.join(tmp, "decoded.ppm")
        subprocess.run(["djpeg", "-ppm", "-outfile", decoded, path], check=True)
        readback = read_ppm(decoded)
    return {
        "file": rel(path),
        "bytes": len(data),
        "sha256": hashlib.sha256(data).hexdigest(),
        "readback": readback,
    }


def rgba_sha(pixels):
    return hashlib.sha256(b"".join(bytes(row) for row in pixels)).hexdigest()


def png_record(pixels, width, height, path):
    write_png(path, pixels, width, height)
    readback, metadata = read_png(path)
    if bytes(b"".join(bytes(row) for row in readback)) != bytes(b"".join(bytes(row) for row in pixels)):
        raise ValueError("PNG pixel mismatch: " + rel(path))
    metadata["pixel_match"] = True
    return {
        "file": rel(path),
        "width": width,
        "height": height,
        "bytes": os.path.getsize(path),
        "sha256": hashlib.sha256(open(path, "rb").read()).hexdigest(),
        "rgba_sha256": rgba_sha(pixels),
        "readback": metadata,
    }


def main():
    source = make_source()
    markers = mask_pixels(source, SOURCE["width"], SOURCE["height"], MARKER)
    small_text = mask_pixels(source, SOURCE["width"], SOURCE["height"], TEXT)
    if len(markers) != 4 * MARKER_SIZE * MARKER_SIZE:
        raise ValueError("marker mask count changed")
    if not small_text:
        raise ValueError("small-text mask is empty")

    source_path = os.path.join(HERE, SOURCE["file"])
    source_record = png_record(source, SOURCE["width"], SOURCE["height"], source_path)
    exports = []
    readbacks = [dict(source_record["readback"], file=rel(source_path))]

    for size in CANVASES:
        for mode, builder in (("fill", fill_export), ("fit", lambda image, target: fit_export(image, target))):
            output, geometry = builder(source, size)
            export = dict(geometry)
            export.update({
                "canvas_px": [size, size],
                "mask_counts": geometry_masks(geometry, size, markers, small_text),
            })
            stem = f"exports/{size}-{mode}"
            png = png_record(output, size, size, os.path.join(HERE, stem + ".png"))
            readbacks.append(dict(png["readback"], file=png["file"]))
            jpg = output_jpg(output, size, size, os.path.join(HERE, stem + ".jpg"))

            square = resize_preview(output, size, PREVIEW_SQUARE)
            square_path = os.path.join(HERE, f"previews/{size}-{mode}-square-{PREVIEW_SQUARE}.png")
            square_record = png_record(square, PREVIEW_SQUARE, PREVIEW_SQUARE, square_path)
            readbacks.append(dict(square_record["readback"], file=square_record["file"]))

            circle = circle_preview(output, size)
            circle_path = os.path.join(HERE, f"previews/{size}-{mode}-circle-{PREVIEW_CIRCLE}.png")
            circle_record = png_record(circle, PREVIEW_CIRCLE, PREVIEW_CIRCLE, circle_path)
            readbacks.append(dict(circle_record["readback"], file=circle_record["file"]))
            outside_alpha = sum(
                1
                for y in range(PREVIEW_CIRCLE)
                for x in range(PREVIEW_CIRCLE)
                if (x + 0.5 - PREVIEW_CIRCLE / 2) ** 2 + (y + 0.5 - PREVIEW_CIRCLE / 2) ** 2
                > (PREVIEW_CIRCLE / 2) ** 2
                and circle[y][x * 4 + 3] == 0
            )

            export.update({
                "png": {key: png[key] for key in ("file", "bytes", "sha256", "readback")},
                "jpg": jpg,
                "preview_square": {
                    **square_record,
                    "mask_counts": preview_masks(
                        square, PREVIEW_SQUARE, size, geometry, markers, small_text
                    ),
                },
                "preview_circle": {
                    **circle_record,
                    "outside_alpha_zero": outside_alpha,
                    "mask_counts": preview_masks(
                        circle, PREVIEW_CIRCLE, size, geometry, markers, small_text
                    ),
                },
            })
            exports.append(export)

    # Negative controls: a wrong-canvas export and a shifted Fill crop must not
    # be accepted as the same result.
    correct, _ = fill_export(source, 256)
    wrong = resize(correct, 0, 0, 256, 256, 255, 256)
    wrong_path = os.path.join(HERE, "negative/wrong-canvas-255x256.png")
    wrong_record = png_record(wrong, 255, 256, wrong_path)
    shifted, shifted_geometry = fill_export(source, 256, shift=0.02)
    shifted_path = os.path.join(HERE, "negative/shifted-fill-256.png")
    shifted_record = png_record(shifted, 256, 256, shifted_path)
    differing_pixels = sum(
        1
        for y in range(256)
        for x in range(256)
        if correct[y][x * 4:x * 4 + 4] != shifted[y][x * 4:x * 4 + 4]
    )

    circle_mask_pixels = sum(
        1
        for y in range(PREVIEW_CIRCLE)
        for x in range(PREVIEW_CIRCLE)
        if (x + 0.5 - PREVIEW_CIRCLE / 2) ** 2 + (y + 0.5 - PREVIEW_CIRCLE / 2) ** 2
        <= (PREVIEW_CIRCLE / 2) ** 2
    )
    results = {
        "generator": "make_samples.py",
        "method": (
            "One deterministic 900 x 1200 RGBA source; nearest-neighbour Fill/Fit geometry; "
            "PNG plus cjpeg -quality 92 / djpeg readback; mask counts map preview pixels back "
            "to exact source marker and small-text pixels."
        ),
        "source": source_record,
        "source_masks": {
            "marker_pixels": len(markers),
            "small_text_pixels": len(small_text),
            "small_text_band_y": [1080, 1086],
        },
        "canvases": CANVASES,
        "preview_sizes": {"square": PREVIEW_SQUARE, "circle": PREVIEW_CIRCLE},
        "circle": {
            "canvas_px": PREVIEW_CIRCLE,
            "inside_pixels_continuous": PREVIEW_CIRCLE * PREVIEW_CIRCLE * math.pi / 4,
            "inside_pixels_counted": circle_mask_pixels,
            "corner_pixels_excluded_continuous": PREVIEW_CIRCLE * PREVIEW_CIRCLE * (1 - math.pi / 4),
        },
        "exports": exports,
        "negative_controls": {
            "wrong_canvas": {
                "file": wrong_record["file"],
                "declared_width": 255,
                "declared_height": 256,
                "criterion": "width === 256 and height === 256",
                "passes": False,
                "expected_acceptance": False,
            },
            "shifted_fill": {
                "file": shifted_record["file"],
                "source_y_px": shifted_geometry["source_y"],
                "differing_pixels_from_correct_256_fill": differing_pixels,
                "expected_file": "exports/256-fill.png",
                "actual_rgba_sha256": rgba_sha(shifted),
                "expected_rgba_sha256": rgba_sha(correct),
                "criterion": "decoded RGBA equals exports/256-fill.png",
                "passes": False,
                "expected_acceptance": False,
            },
        },
        "png_readbacks": readbacks,
    }
    out_path = os.path.join(HERE, "results.json")
    with open(out_path, "w", encoding="utf-8") as stream:
        json.dump(results, stream, indent=2, sort_keys=True)
        stream.write("\n")
    print(
        f"wrote {len(exports)} exports, {len(readbacks)} PNG readbacks, "
        f"source marker pixels={len(markers)}, small-text pixels={len(small_text)}"
    )


if __name__ == "__main__":
    main()
