///|
pub(all) struct ChunkInfo {
  y : Int32
  rows : Int32
  offset : Int32
  encoded_bytes : Int32
  decoded_bytes : Int32
} derive(Debug, Eq)

///|
pub extend ChunkInfo with @debug.Debug::{to_repr}

///|
pub extend ChunkInfo with Eq::{equal, not_equal}

///|
priv struct Block {
  info : ChunkInfo
  window : Window
  payload : Int32
}

///|
priv struct LevelDirectory {
  geometry : LevelInfo
  start : Int32
}

///|
pub struct ExrFile {
  priv bytes : Bytes
  priv header : Header
  priv blocks : Array[Block]
  priv directory : Array[LevelDirectory]
  priv limits : Limits
}

///|
pub struct ExrDocument {
  priv files : Array[ExrFile]
  priv multipart : Bool
}

///|
fn channel_stride(channels : Array[Channel]) -> Int32 raise ExrError {
  let mut stride = 0
  for ch in channels {
    stride = add(stride, ch.kind.bytes())
  }
  stride
}

///|
fn check_shared_headers(first : Header, other : Header) -> Unit raise ExrError {
  for name in ["displayWindow", "pixelAspectRatio"] {
    if first.attribute(name) != other.attribute(name) {
      raise Invalid("multipart shared attribute differs: " + name, other.end)
    }
  }
  for name in ["chromaticities", "timeCode", "timecode"] {
    if other.attribute(name) is Some(value) &&
      first.attribute(name) != Some(value) {
      raise Invalid("multipart shared attribute differs: " + name, other.end)
    }
  }
}

///|
/// Parse all headers and offset tables once; compressed pixels remain in the input.
pub fn ExrDocument::open(
  bytes : Bytes,
  limits~ : Limits = Limits::default(),
) -> ExrDocument raise ExrError {
  let (r, flags) = file_prefix(bytes, limits)
  let multipart = (flags & 4096U) != 0U
  let headers = [read_header(r, flags, limits)]
  if multipart {
    let names : Map[String, Bool] = Map([])
    names[headers[0].name.unwrap()] = true
    while r.pos < r.end && bytes[r.pos] != 0 {
      if headers.length() >= limits.max_chunks {
        raise Limit("part count")
      }
      let h = read_header(r, flags, limits)
      let name = h.name.unwrap()
      if names.contains(name) {
        raise Invalid("duplicate part name", h.end)
      }
      names[name] = true
      check_shared_headers(headers[0], h)
      headers.push(h)
    }
    ignore(r.u8())
    if r.pos > limits.max_header_bytes {
      raise Limit("multipart header bytes")
    }
  }
  let mut chunk_count = 0
  let mut logical_bytes = 0
  let directories = headers.map(h => {
    let mut count = 0
    let stride = channel_stride(h.channels)
    let directory = block_levels(h).map(level => {
      let start = count
      count = add(count, mul(level.columns, level.rows))
      logical_bytes = add(logical_bytes, mul(level.window.pixels(), stride))
      if logical_bytes > limits.max_decoded_bytes {
        raise Limit("document decoded bytes across parts and levels")
      }
      { geometry: level, start, }
    })
    if h.declared_chunks is Some(declared) && count != declared {
      raise Invalid("chunkCount disagrees with part geometry", h.end)
    }
    chunk_count = add(chunk_count, count)
    if chunk_count > limits.max_chunks {
      raise Limit("document chunk count")
    }
    directory
  })
  let table_end = add(r.pos, mul(chunk_count, 8))
  if table_end > bytes.length() {
    raise Invalid("truncated chunk offset table", r.pos)
  }
  let all_blocks : Array[Block] = []
  let files = headers.mapi((part, h) => {
    let blocks = []
    let stride = channel_stride(h.channels)
    let head_bytes = (if multipart { 4 } else { 0 }) +
      (if h.tiles is Some(_) { 20 } else { 8 })
    let size = block_size(h)
    for entry in directories[part] {
      let level = entry.geometry
      for y = 0; y < level.rows; y = y + 1 {
        for x = 0; x < level.columns; x = x + 1 {
          let offset = index(r.uint(8))
          if offset < table_end || offset > bytes.length() - head_bytes {
            raise Invalid("chunk offset outside file body", offset)
          }
          let c = { data: bytes, pos: offset, end: offset + head_bytes, }
          if multipart && c.i32() != part {
            raise Invalid("chunk belongs to another part", offset)
          }
          let window = block_window(level, size, x, y)
          if h.tiles is Some(_) {
            if c.i32() != x ||
              c.i32() != y ||
              c.i32() != level.x ||
              c.i32() != level.y {
              raise Invalid("offset table tile coordinate mismatch", offset)
            }
          } else if c.i32() != window.y0 {
            raise Invalid("offset table row mismatch", offset)
          }
          let encoded_bytes = c.i32()
          let decoded_bytes = mul(window.pixels(), stride)
          if encoded_bytes <= 0 ||
            encoded_bytes > decoded_bytes ||
            encoded_bytes > bytes.length() - c.end ||
            (h.compression == None && encoded_bytes != decoded_bytes) {
            raise Invalid("chunk payload size", c.end - 4)
          }
          let info = {
            y: window.y0,
            rows: window.height(),
            offset,
            encoded_bytes,
            decoded_bytes,
          }
          let block = { info, window, payload: c.end, }
          blocks.push(block)
          all_blocks.push(block)
        }
      }
    }
    { bytes, header: h, blocks, directory: directories[part], limits, }
  })
  all_blocks.sort_by((a, b) => a.info.offset.compare(b.info.offset))
  let mut previous_end = table_end
  for block in all_blocks {
    if block.info.offset < previous_end {
      raise Invalid("duplicate/overlapping chunk range", block.info.offset)
    }
    previous_end = block.payload + block.info.encoded_bytes
  }
  { files, multipart, }
}

///|
pub fn ExrDocument::is_multipart(self : ExrDocument) -> Bool {
  self.multipart
}

///|
pub fn ExrDocument::part_count(self : ExrDocument) -> Int32 {
  self.files.length()
}

///|
pub fn ExrDocument::part(
  self : ExrDocument,
  index : Int32,
) -> ExrFile raise ExrError {
  if index < 0 || index >= self.files.length() {
    raise Invalid("part index out of bounds", 0)
  }
  self.files[index]
}

///|
pub fn ExrDocument::named_part(
  self : ExrDocument,
  name : String,
) -> ExrFile raise ExrError {
  for file in self.files {
    if file.header.name == Some(name) {
      return file
    }
  }
  raise Invalid("unknown part: " + name, 0)
}

///|
/// Open a single-part scanline or tiled file.
pub fn ExrFile::open(
  bytes : Bytes,
  limits~ : Limits = Limits::default(),
) -> ExrFile raise ExrError {
  let doc = ExrDocument::open(bytes, limits~)
  if doc.multipart {
    raise Unsupported("use ExrDocument to select a multipart image")
  }
  doc.files[0]
}

///|
pub fn ExrFile::header(self : ExrFile) -> Header {
  self.header
}

///|
pub fn ExrFile::levels(self : ExrFile) -> Array[LevelInfo] {
  self.directory.map(entry => entry.geometry)
}

///|
fn ExrFile::level_entry(
  self : ExrFile,
  x : Int32,
  y : Int32,
) -> LevelDirectory raise ExrError {
  for entry in self.directory {
    if entry.geometry.x == x && entry.geometry.y == y {
      return entry
    }
  }
  raise Invalid("level does not exist", 0)
}

///|
pub fn ExrFile::chunks(self : ExrFile) -> Array[ChunkInfo] {
  self.blocks.map(block => block.info)
}

///|
fn ExrFile::plan_blocks(
  self : ExrFile,
  entry : LevelDirectory,
  region : Window,
) -> Array[Block] {
  let level = entry.geometry
  let base = level.window
  let (width, height) = block_size(self.header)
  let result = []
  if base.intersection(region) is Some(w) {
    let x0 = (w.x0 - base.x0) / width
    let x1 = (w.x1 - base.x0) / width
    let y0 = (w.y0 - base.y0) / height
    let y1 = (w.y1 - base.y0) / height
    for y = y0; y <= y1; y = y + 1 {
      for x = x0; x <= x1; x = x + 1 {
        result.push(self.blocks[entry.start + y * level.columns + x])
      }
    }
  }
  result
}

///|
/// Exact encoded chunks needed for a region of the base level.
pub fn ExrFile::plan(self : ExrFile, region : Window) -> Array[ChunkInfo] {
  self.plan_blocks(self.directory[0], region).map(block => block.info)
}

///|
pub fn ExrFile::plan_level(
  self : ExrFile,
  x : Int32,
  y : Int32,
  region : Window,
) -> Array[ChunkInfo] raise ExrError {
  self.plan_blocks(self.level_entry(x, y), region).map(block => block.info)
}

///|
fn ExrFile::unpack_block(self : ExrFile, block : Block) -> Bytes raise ExrError {
  let bytes = self.bytes[block.payload:block.payload + block.info.encoded_bytes].to_owned()
  decode_block(
    bytes,
    self.header.compression,
    self.header.channels,
    block.window,
    block.info.decoded_bytes,
  )
}