///|
/// Supply images in TileDescription::levels order. Scanline sources have one image.
pub(all) struct PartSource {
  name : String
  compression : Compression
  tiles : TileDescription?
  images : Array[Image]
}

///|
fn prepare_part(
  source : PartSource,
  multipart : Bool,
  limits : Limits,
  prefix : Sink,
) -> (Header, Array[LevelInfo], Int32) raise ExrError {
  if source.images.is_empty() {
    raise Invalid("missing base image", 0)
  }
  let base = source.images[0]
  let attrs = image_attributes(base, source.compression)
  if source.tiles is Some(t) {
    attrs.push(t.attribute())
  }
  let levels = match source.tiles {
    Some(t) => t.levels(base.data_window)
    None =>
      [
        {
          x: 0,
          y: 0,
          window: base.data_window,
          columns: 1,
          rows: (base.data_window.height() - 1) / source.compression.scanlines() +
          1,
        },
      ]
  }
  if source.images.length() != levels.length() {
    raise Invalid("wrong number of source levels", 0)
  }
  let mut count = 0
  for i = 0; i < levels.length(); i = i + 1 {
    let level = levels[i]
    let image = source.images[i]
    if image.data_window != level.window ||
      image.channels() != base.channels() ||
      image.display_window != base.display_window ||
      image.pixel_aspect != base.pixel_aspect ||
      image.screen() != base.screen() ||
      image.metadata != base.metadata {
      raise Invalid(
        "source level geometry, channels or shared header differs", 0,
      )
    }
    count = add(count, mul(level.columns, level.rows))
  }
  if multipart {
    attrs.push(Attribute::text("name", source.name))
    let kind = if source.tiles is Some(_) {
      "tiledimage"
    } else {
      "scanlineimage"
    }
    attrs.push(Attribute::text("type", kind))
    attrs.push(Attribute::integer("chunkCount", count))
  } else if source.name != "" {
    attrs.push(Attribute::text("name", source.name))
  }
  for attr in attrs {
    let size = add(
      @utf8.encode(attr.name).length() + @utf8.encode(attr.kind).length() + 6,
      attr.value.length(),
    )
    if size >= limits.max_header_bytes - prefix.size() ||
      size >= limits.max_file_bytes - prefix.size() {
      raise Limit("encoded header bytes")
    }
    write_attribute(prefix, attr)
  }
  prefix.uint(0UL, 1)
  let flags = 1026U |
    (if multipart {
      4096U
    } else if source.tiles is Some(_) {
      512U
    } else {
      0U
    })
  // Reuse the same typed header validation for both encoding and decoding.
  let r = reader(header_bytes(attrs))
  r.pos = 8
  (read_header(r, flags, limits), levels, count)
}

///|
fn encode_sources(
  parts : Array[PartSource],
  multipart : Bool,
  limits : Limits,
) -> Bytes raise ExrError {
  limits.validate()
  if parts.is_empty() || parts.length() > limits.max_chunks {
    raise Limit("encoded part count")
  }
  let prefix = sink()
  prefix.uint(20000630UL, 4)
  let flags = 1026UL |
    (if multipart {
      4096UL
    } else if parts[0].tiles is Some(_) {
      512UL
    } else {
      0UL
    })
  prefix.uint(flags, 4)
  let names : Map[String, Bool] = Map([])
  let mut logical_bytes = 0
  let mut chunks = 0
  let prepared = parts.map(source => {
    let (h, levels, count) = prepare_part(source, multipart, limits, prefix)
    if multipart {
      if names.contains(source.name) {
        raise Invalid("duplicate output part name", 0)
      }
      names[source.name] = true
    }
    chunks = add(chunks, count)
    if chunks > limits.max_chunks {
      raise Limit("encoded chunk count")
    }
    let stride = channel_stride(h.channels)
    for level in levels {
      logical_bytes = add(logical_bytes, mul(level.window.pixels(), stride))
      if logical_bytes > limits.max_decoded_bytes {
        raise Limit("encoded document decoded bytes")
      }
    }
    (h, levels)
  })
  if multipart {
    for (h, _) in prepared {
      check_shared_headers(prepared[0].0, h)
    }
    prefix.uint(0UL, 1)
  }
  if prefix.size() > limits.max_header_bytes {
    raise Limit("encoded header bytes")
  }
  let table_end = add(prefix.size(), mul(chunks, 8))
  if table_end > limits.max_file_bytes {
    raise Limit("encoded chunk offset table")
  }
  let body = sink()
  for part = 0; part < prepared.length(); part = part + 1 {
    let (h, levels) = prepared[part]
    let size = block_size(h)
    let chunk_header = (if multipart { 4 } else { 0 }) +
      (if h.tiles is Some(_) { 20 } else { 8 })
    for i = 0; i < levels.length(); i = i + 1 {
      let level = levels[i]
      let image = parts[part].images[i]
      for y = 0; y < level.rows; y = y + 1 {
        for x = 0; x < level.columns; x = x + 1 {
          let w = block_window(level, size, x, y)
          let raw = sink()
          for row = 0; row < w.height(); row = row + 1 {
            let start = (w.y0 - level.window.y0 + row) * level.window.width() +
              (w.x0 - level.window.x0)
            for p in image.planes {
              let sample_bytes = p.channel.kind.bytes()
              let begin = start * sample_bytes
              let end = begin + w.width() * sample_bytes
              for b in p.data[begin:end] {
                raw.bytes.push(b)
              }
            }
          }
          let packed = encode_block(raw.finish(), h.compression, h.channels, w)
          let offset = add(table_end, body.size())
          if add(offset, add(chunk_header, packed.length())) >
            limits.max_file_bytes {
            raise Limit("encoded file bytes")
          }
          prefix.uint(offset.to_uint64(), 8)
          if multipart {
            body.i32(part)
          }
          if h.tiles is Some(_) {
            body.i32(x)
            body.i32(y)
            body.i32(level.x)
            body.i32(level.y)
          } else {
            body.i32(w.y0)
          }
          body.i32(packed.length())
          body.raw(packed)
        }
      }
    }
  }
  for b in body.bytes {
    prefix.bytes.push(b)
  }
  prefix.finish()
}

///|
/// Encode single-part tiled data, including caller-supplied mipmap or ripmap levels.
pub fn encode_tiled(
  images : Array[Image],
  tiles : TileDescription,
  compression~ : Compression = Piz,
  limits~ : Limits = Limits::default(),
) -> Bytes raise ExrError {
  encode_sources(
    [{ name: "", compression, tiles: Some(tiles), images, }],
    false,
    limits,
  )
}

///|
/// Encode named flat scanline and tiled parts, each with its own compression.
pub fn encode_multipart(
  parts : Array[PartSource],
  limits~ : Limits = Limits::default(),
) -> Bytes raise ExrError {
  encode_sources(parts, true, limits)
}

///|
/// Extract a part or resolution to an ordinary single-part file, with optional
/// region and channel selection. Tiled output contains one level.
pub fn ExrFile::extract(
  self : ExrFile,
  level_x~ : Int32 = 0,
  level_y~ : Int32 = 0,
  region~ : Window? = None,
  channels~ : Array[String]? = None,
  tiles~ : TileDescription? = None,
  compression~ : Compression = Piz,
) -> Bytes raise ExrError {
  let image = self.read_level(level_x, level_y, region~, channels~)
  match tiles {
    Some(t) => encode_tiled([image], t, compression~, limits=self.limits)
    None => encode(image, compression~, limits=self.limits)
  }
}

///|
/// Recompress all parts and resolutions, preserving storage layout, names and sample bits.
pub fn ExrDocument::transcode(
  self : ExrDocument,
  compression : Compression,
) -> Bytes raise ExrError {
  let parts = self.files.map(file => {
    let images = file.levels().map(level => file.read_level(level.x, level.y))
    {
      name: file.header.name.unwrap_or(""),
      compression,
      tiles: file.header.tiles,
      images,
    }
  })
  encode_sources(parts, self.multipart, self.files[0].limits)
}