from __future__ import annotations

import struct
import zlib
from pathlib import Path

from .formats import ShapeFrame


DEFAULT_BACKGROUND = (10, 12, 18, 255)


def rgba_buffer(width: int, height: int, color: tuple[int, int, int, int]) -> bytearray:
    r, g, b, a = color
    row = bytes((r, g, b, a)) * width
    return bytearray(row * height)


def blit_frame(
    buffer: bytearray,
    canvas_width: int,
    canvas_height: int,
    left: int,
    top: int,
    frame: ShapeFrame,
    pixels: list[int],
    palette: list[tuple[int, int, int]],
    flipped: bool,
) -> None:
    for src_y in range(frame.height):
        dst_y = top + src_y
        if dst_y < 0 or dst_y >= canvas_height:
            continue
        row_base = src_y * frame.width
        for src_x in range(frame.width):
            color_index = pixels[row_base + (frame.width - 1 - src_x if flipped else src_x)]
            if color_index < 0:
                continue
            dst_x = left + src_x
            if dst_x < 0 or dst_x >= canvas_width:
                continue
            pixel_base = (dst_y * canvas_width + dst_x) * 4
            r, g, b = palette[color_index]
            buffer[pixel_base : pixel_base + 4] = bytes((r, g, b, 255))


def write_png_rgba(path: Path, width: int, height: int, pixels: bytearray) -> None:
    def 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)
        )

    rows = bytearray()
    stride = width * 4
    for row in range(height):
        rows.append(0)
        start = row * stride
        rows.extend(pixels[start : start + stride])

    payload = bytearray(b"\x89PNG\r\n\x1a\n")
    payload.extend(chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)))
    payload.extend(chunk(b"IDAT", zlib.compress(bytes(rows), level=9)))
    payload.extend(chunk(b"IEND", b""))
    path.write_bytes(payload)