#!/usr/bin/env python3
"""
Utility script to convert a PNG file (non-interlaced, 8-bit RGBA/RGB)
into a 24-bit BMP with optional letterboxing resize.

The script is intentionally lightweight and relies only on Python's
standard library so it can run in constrained build environments.
"""

from __future__ import annotations

import argparse
import struct
import sys
import zlib
from pathlib import Path


PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


def parse_png(path: Path):
  data = path.read_bytes()
  if not data.startswith(PNG_SIGNATURE):
    raise ValueError("Input is not a PNG file")

  idx = len(PNG_SIGNATURE)
  width = height = bit_depth = color_type = None
  compressed = bytearray()
  interlaced = False

  while idx < len(data):
    if idx + 8 > len(data):
      raise ValueError("Corrupted PNG (unexpected EOF)")
    length = struct.unpack(">I", data[idx : idx + 4])[0]
    idx += 4
    chunk_type = data[idx : idx + 4]
    idx += 4
    chunk_data = data[idx : idx + length]
    idx += length
    crc = data[idx : idx + 4]  # noqa: F841 - crc skipped (validated by reader)
    idx += 4

    if chunk_type == b"IHDR":
      width, height, bit_depth, color_type, compression, filter_method, interlace = struct.unpack(
        ">IIBBBBB", chunk_data
      )
      if compression != 0 or filter_method != 0:
        raise ValueError("Unsupported PNG compression/filter method")
      interlaced = interlace != 0
    elif chunk_type == b"IDAT":
      compressed.extend(chunk_data)
    elif chunk_type == b"IEND":
      break

  if interlaced:
    raise ValueError("Interlaced PNGs are not supported by this script")
  if bit_depth != 8:
    raise ValueError(f"Unsupported bit depth: {bit_depth}")
  if color_type not in (2, 6):
    raise ValueError(f"Unsupported color type: {color_type}")

  raw = zlib.decompress(bytes(compressed))
  bytes_per_pixel = 3 if color_type == 2 else 4
  stride = width * bytes_per_pixel
  expected = (stride + 1) * height
  if len(raw) != expected:
    raise ValueError("Corrupted PNG data")

  # Apply PNG scanline filters
  image = bytearray(width * height * 4)  # Force RGBA output
  prev_row = [0] * (stride)

  def paeth(a, b, c):
    p = a + b - c
    pa = abs(p - a)
    pb = abs(p - b)
    pc = abs(p - c)
    if pa <= pb and pa <= pc:
      return a
    if pb <= pc:
      return b
    return c

  out_idx = 0
  for y in range(height):
    offset = y * (stride + 1)
    filter_type = raw[offset]
    row = bytearray(raw[offset + 1 : offset + 1 + stride])
    if filter_type == 1:  # Sub
      for i in range(stride):
        left = row[i - bytes_per_pixel] if i >= bytes_per_pixel else 0
        row[i] = (row[i] + left) & 0xFF
    elif filter_type == 2:  # Up
      for i in range(stride):
        row[i] = (row[i] + prev_row[i]) & 0xFF
    elif filter_type == 3:  # Average
      for i in range(stride):
        left = row[i - bytes_per_pixel] if i >= bytes_per_pixel else 0
        up = prev_row[i]
        row[i] = (row[i] + ((left + up) >> 1)) & 0xFF
    elif filter_type == 4:  # Paeth
      for i in range(stride):
        left = row[i - bytes_per_pixel] if i >= bytes_per_pixel else 0
        up = prev_row[i]
        up_left = prev_row[i - bytes_per_pixel] if i >= bytes_per_pixel else 0
        row[i] = (row[i] + paeth(left, up, up_left)) & 0xFF
    elif filter_type != 0:
      raise ValueError(f"Unsupported PNG filter type: {filter_type}")

    # Convert to RGBA
    for x in range(width):
      if color_type == 2:
        r, g, b = row[x * 3 : x * 3 + 3]
        a = 255
      else:
        r, g, b, a = row[x * 4 : x * 4 + 4]
      image[out_idx : out_idx + 4] = bytes((r, g, b, a))
      out_idx += 4

    prev_row = list(row)

  return width, height, image


def resize_with_letterbox(image, width, height, target_w, target_h, background, scale_factor=1.0):
  if width == target_w and height == target_h and abs(scale_factor - 1.0) < 1e-6:
    return image, width, height

  bg_r, bg_g, bg_b = background
  base_scale = min(target_w / width, target_h / height)
  base_scale *= scale_factor
  base_scale = max(base_scale, 1 / max(width, height))  # avoid zero / collapse
  scaled_w = max(1, int(round(width * base_scale)))
  scaled_h = max(1, int(round(height * base_scale)))

  output = bytearray(target_w * target_h * 4)
  # Fill background
  for i in range(0, len(output), 4):
    output[i : i + 4] = bytes((bg_r, bg_g, bg_b, 255))

  offset_x = (target_w - scaled_w) // 2
  offset_y = (target_h - scaled_h) // 2

  for y in range(scaled_h):
    src_y = min(height - 1, int(round(y / base_scale)))
    for x in range(scaled_w):
      src_x = min(width - 1, int(round(x / base_scale)))
      src_idx = (src_y * width + src_x) * 4
      dst_idx = ((y + offset_y) * target_w + (x + offset_x)) * 4
      output[dst_idx : dst_idx + 4] = image[src_idx : src_idx + 4]

  return output, target_w, target_h


def blend_to_rgb(image):
  rgb = bytearray(len(image) // 4 * 3)
  for i in range(0, len(image), 4):
    r, g, b, a = image[i : i + 4]
    if a == 255:
      rgb[(i // 4) * 3 : (i // 4) * 3 + 3] = bytes((b, g, r))  # BMP stores BGR
    else:
      alpha = a / 255.0
      bg = (255, 255, 255)
      rr = int(round(r * alpha + bg[0] * (1 - alpha)))
      gg = int(round(g * alpha + bg[1] * (1 - alpha)))
      bb = int(round(b * alpha + bg[2] * (1 - alpha)))
      rgb[(i // 4) * 3 : (i // 4) * 3 + 3] = bytes((bb, gg, rr))
  return rgb


def write_bmp(path: Path, width: int, height: int, rgb: bytearray):
  row_stride = (width * 3 + 3) & ~3  # align to 4 bytes
  padding = row_stride - width * 3
  pixel_data = bytearray()

  for y in range(height - 1, -1, -1):
    start = y * width * 3
    end = start + width * 3
    pixel_data.extend(rgb[start:end])
    if padding:
      pixel_data.extend(b"\0" * padding)

  file_size = 14 + 40 + len(pixel_data)
  header = struct.pack("<2sIHHI", b"BM", file_size, 0, 0, 14 + 40)
  dib_header = struct.pack(
    "<IIIHHIIIIII",
    40,  # header size
    width,
    height,
    1,  # planes
    24,  # bits per pixel
    0,  # compression
    len(pixel_data),
    2835,  # horizontal resolution (px/m ~72dpi)
    2835,  # vertical resolution
    0,
    0,
  )

  path.write_bytes(header + dib_header + pixel_data)


def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("input", type=Path)
  parser.add_argument("output", type=Path)
  parser.add_argument("--width", type=int, help="Target width (px)")
  parser.add_argument("--height", type=int, help="Target height (px)")
  parser.add_argument(
    "--scale",
    type=float,
    default=1.0,
    help="Optional multiplier applied to the fitted image size (e.g. 0.7 adds padding).",
  )
  parser.add_argument(
    "--background",
    type=str,
    default="FFFFFF",
    help="Background hex color used for transparent pixels (default: FFFFFF)",
  )
  args = parser.parse_args()

  try:
    width, height, image = parse_png(args.input)
    if args.width and args.height:
      bg = tuple(int(args.background[i : i + 2], 16) for i in (0, 2, 4))
      image, width, height = resize_with_letterbox(
        image, width, height, args.width, args.height, bg, max(args.scale, 0.05)
      )
    rgb = blend_to_rgb(image)
    write_bmp(args.output, width, height, rgb)
  except Exception as exc:  # noqa: BLE001
    print(f"Error: {exc}", file=sys.stderr)
    sys.exit(1)


if __name__ == "__main__":
  main()