# png_write.ludic - raw 8- or 16-bit pixels (1 to 4 channels) as a PNG anything opens: filter 0 on every row,
# deflate's stored blocks (no zlib here), the CRCs and the Adler-32. Needs nothing else of the lab, so a
# program with a lab of its own imports this file alone: import "ludic.lab/png_write.ludic"

export state LabPngState {
  table: []long = lp_crc_table()
}

# the CRC-32 table, made once with a state
function lp_crc_table() -> []long {
  let t = new []long
  let poly: long = 3988292384            # 0xEDB88320: eight hex digits would be a 32-bit pattern
  for n in 0 .. 256 {
    var c = long(n)
    for k in 0 .. 8 {
      if (c & long(1)) != long(0) { c = (c >> long(1)) ^ poly } else { c = c >> long(1) }
    }
    push(t, c)
  }
  return t
}
# CRC-32 of out[from .. from + n)
function lp_crc_of(table: []long, out: []byte, from: int, n: int) -> long {
  let mask: long = 4294967295
  var c = mask
  for i in from .. from + n {
    let k = int((c ^ long(out[i])) & long(255))
    c = table[k] ^ (c >> long(8))
  }
  return c ^ mask
}
function lp_put32(out: []byte, at: int, v: long) -> void {
  out[at] = int((v >> long(24)) & long(255))
  out[at + 1] = int((v >> long(16)) & long(255))
  out[at + 2] = int((v >> long(8)) & long(255))
  out[at + 3] = int(v & long(255))
}
function lp_tag(out: []byte, at: int, tag: string) -> void {
  for i in 0 .. 4 { out[at + i] = tag[i] }
}
# PNG's colour type for 1..4 channels: grey, grey + alpha, RGB, RGBA
function lp_colour_type(channels: int) -> int {
  if channels == 1 { return 0 }
  if channels == 2 { return 4 }
  if channels == 3 { return 2 }
  return 6
}

# the header, then the pixels, then the end, into one buffer; the length written, or -1
@alloc_ok("a bake's preview or a lab shot: one image written, once")
function lp_encode(table: []long, w: int, h: int, channels: int, px: []byte, at: int, out: []byte) -> int { return lp_encode_depth(table, w, h, channels, 8, px, at, out) }
# ... with `depth` bits a sample (8, or 16: two bytes each, most significant first, as PNG stores them)
@alloc_ok("a bake's preview or a lab shot: one image written, once")
function lp_encode_depth(table: []long, w: int, h: int, channels: int, depth: int, px: []byte, at: int, out: []byte) -> int {
  let sig = [137, 80, 78, 71, 13, 10, 26, 10]
  for i in 0 .. 8 { out[i] = sig[i] }
  lp_put32(out, 8, long(13))
  lp_tag(out, 12, "IHDR")
  lp_put32(out, 16, long(w))
  lp_put32(out, 20, long(h))
  out[24] = depth
  out[25] = lp_colour_type(channels)
  out[26] = 0
  out[27] = 0
  out[28] = 0
  lp_put32(out, 29, lp_crc_of(table, out, 12, 17))
  let o = lp_idat(table, w * channels * depth / 8, h, px, at, out, 33)
  lp_put32(out, o, long(0))
  lp_tag(out, o + 4, "IEND")
  lp_put32(out, o + 8, lp_crc_of(table, out, o + 4, 4))
  return o + 12
}
# the IDAT chunk at `at0`: every row (`stride` bytes) behind a filter byte of 0, in stored blocks; where it ends
function lp_idat(table: []long, stride: int, h: int, px: []byte, at: int, out: []byte, at0: int) -> int {
  let row = stride + 1
  let raw = row * h
  let blocks = (raw + 65534) / 65535
  lp_put32(out, at0, long(2 + blocks * 5 + raw + 4))
  lp_tag(out, at0 + 4, "IDAT")
  var o = at0 + 8
  out[o] = 120
  out[o + 1] = 1
  o += 2
  var a = 1
  var b = 0
  var left = raw
  var k = 0                                # the byte of the raw stream being written
  while left > 0 {
    let n = Math.min(left, 65535)
    left -= n
    var fin = 0
    if left == 0 { fin = 1 }
    out[o] = fin
    out[o + 1] = n & 255
    out[o + 2] = (n >> 8) & 255
    out[o + 3] = (~n) & 255
    out[o + 4] = ((~n) >> 8) & 255
    o += 5
    for j in 0 .. n {
      let x = k % row
      var v = 0
      if x > 0 { v = px[at + (k / row) * stride + x - 1] & 255 }
      out[o] = v
      a = (a + v) % 65521
      b = (b + a) % 65521
      o += 1
      k += 1
    }
  }
  lp_put32(out, o, long(b) * long(65536) + long(a))
  o += 4
  lp_put32(out, o, lp_crc_of(table, out, at0 + 4, o - at0 - 4))
  return o + 4
}

# the bytes a PNG of this size takes, stored
function lp_png_size(w: int, h: int, channels: int) -> int { return lp_png_bytes(w * channels, h) }
# ... of rows `stride` bytes long
function lp_png_bytes(stride: int, h: int) -> int {
  let raw = (stride + 1) * h
  return 8 + 25 + 12 + 2 + ((raw + 65534) / 65535) * 5 + raw + 4 + 12
}

# w x h pixels of `channels` 8-bit samples each (1 grey, 2 grey + alpha, 3 RGB, 4 RGBA), rows top down,
# starting at px[at], written as a PNG at `path`
@alloc_ok("a bake's preview or a lab shot: one image written, once")
export function lab_png_write_from(lab_png_st: LabPngState, path: string, w: int, h: int, channels: int, px: []byte, at: int) -> bool {
  if channels < 1 or channels > 4 or w <= 0 or h <= 0 or at + w * h * channels > len(px) { return false }
  let out = buffer(lp_png_size(w, h, channels))
  let n = lp_encode(lab_png_st.table, w, h, channels, px, at, out)
  let ok = Fs.write_bytes(path, out, n)
  free(out)
  return ok
}
export function lab_png_write(lab_png_st: LabPngState, path: string, w: int, h: int, channels: int, px: []byte) -> bool {
  return lab_png_write_from(lab_png_st, path, w, h, channels, px, 0)
}
# w x h pixels of `channels` 16-bit samples each, two bytes a sample most significant first (a height map:
# 1 channel, read back by render3d's decoder as R16), rows top down, starting at px[at]
@alloc_ok("a bake or a test's fixture: one image written, once")
export function lab_png_write16_from(lab_png_st: LabPngState, path: string, w: int, h: int, channels: int, px: []byte, at: int) -> bool {
  if channels < 1 or channels > 4 or w <= 0 or h <= 0 or at + w * h * channels * 2 > len(px) { return false }
  let out = buffer(lp_png_bytes(w * channels * 2, h))
  let n = lp_encode_depth(lab_png_st.table, w, h, channels, 16, px, at, out)
  let ok = Fs.write_bytes(path, out, n)
  free(out)
  return ok
}