// Kafka wire protocol primitive codecs.
//
// Fixed-width integers are big-endian. RPC-layer compact lengths and tag
// buffers use unsigned LEB128 varints; record-batch internals use zig-zag
// varints (see record.mbt).

///|
pub(all) suberror DecodeError {
  UnexpectedEof
  Malformed(String)
}

///|
pub struct Encoder {
  buf : Array[Byte]
}

///|
pub fn Encoder::new() -> Encoder {
  { buf: [], }
}

///|
/// Preallocate the backing array to hold roughly `n` bytes up front, so
/// the hot request/record-batch encoding paths avoid repeated array
/// growth reallocations. Growth stalls when the estimate is exceeded.
/// `n` < 0 inserts a default growable capacity (16 bytes).
pub fn Encoder::with_capacity(n : Int) -> Encoder {
  let capacity = if n < 0 { 16 } else { n }
  { buf: Array(capacity~), }
}

///|
pub fn Encoder::to_bytes(self : Encoder) -> Bytes {
  Bytes::from_array(self.buf)
}

///|
pub fn Encoder::write_byte(self : Encoder, b : Byte) -> Unit {
  self.buf.push(b)
}

///|
pub fn Encoder::write_bool(self : Encoder, v : Bool) -> Unit {
  self.buf.push(if v { b'\x01' } else { b'\x00' })
}

///|
pub fn Encoder::write_i8(self : Encoder, v : Int) -> Unit {
  self.buf.push(v.to_byte())
}

///|
pub fn Encoder::write_i16(self : Encoder, v : Int) -> Unit {
  self.buf.push((v >> 8).to_byte())
  self.buf.push(v.to_byte())
}

///|
pub fn Encoder::write_i32(self : Encoder, v : Int) -> Unit {
  for shift in [24, 16, 8, 0] {
    self.buf.push((v >> shift).to_byte())
  }
}

///|
pub fn Encoder::write_i64(self : Encoder, v : Int64) -> Unit {
  for shift in [56, 48, 40, 32, 24, 16, 8, 0] {
    self.buf.push((v >> shift).to_byte())
  }
}

///|
/// FLOAT64: big-endian IEEE 754 double (quota config values).
pub fn Encoder::write_f64(self : Encoder, v : Double) -> Unit {
  let bits = v.reinterpret_as_uint64()
  for shift in [56, 48, 40, 32, 24, 16, 8, 0] {
    self.buf.push((bits >> shift).to_byte())
  }
}

///|
pub fn Encoder::write_bytes(self : Encoder, b : Bytes) -> Unit {
  for byte in b {
    self.buf.push(byte)
  }
}

///|
/// Unsigned LEB128 varint.
pub fn Encoder::write_uvarint(self : Encoder, v : UInt) -> Unit {
  let mut v = v
  while v >= 0x80U {
    self.buf.push((v | 0x80U).to_byte())
    v = v >> 7
  }
  self.buf.push(v.to_byte())
}

///|
/// Zig-zag varint (used inside record batches).
pub fn Encoder::write_varint(self : Encoder, v : Int) -> Unit {
  self.write_uvarint(((v << 1) ^ (v >> 31)).reinterpret_as_uint())
}

///|
/// Zig-zag varlong (used inside record batches).
pub fn Encoder::write_varlong(self : Encoder, v : Int64) -> Unit {
  let mut n = ((v << 1) ^ (v >> 63)).reinterpret_as_uint64()
  while n >= 0x80UL {
    self.buf.push((n | 0x80UL).to_byte())
    n = n >> 7
  }
  self.buf.push(n.to_byte())
}

///|
/// Compact array/bytes length: count + 1 as unsigned varint.
pub fn Encoder::write_compact_len(self : Encoder, n : Int) -> Unit {
  self.write_uvarint((n + 1).reinterpret_as_uint())
}

///|
/// Legacy NULLABLE_STRING: INT16 length, -1 = null.
pub fn Encoder::write_nullable_string(self : Encoder, s : String?) -> Unit {
  match s {
    None => self.write_i16(-1)
    Some(s) => {
      let b = @utf8.encode(s)
      self.write_i16(b.length())
      self.write_bytes(b)
    }
  }
}

///|
pub fn Encoder::write_compact_string(self : Encoder, s : String) -> Unit {
  let b = @utf8.encode(s)
  self.write_compact_len(b.length())
  self.write_bytes(b)
}

///|
pub fn Encoder::write_compact_nullable_string(
  self : Encoder,
  s : String?,
) -> Unit {
  match s {
    None => self.write_uvarint(0U)
    Some(s) => self.write_compact_string(s)
  }
}

///|
/// Empty tag buffer.
pub fn Encoder::write_tag_buffer(self : Encoder) -> Unit {
  self.buf.push(b'\x00')
}

///|
/// Tag buffer with the given tagged fields. Count, tag numbers, and payload
/// sizes are plain unsigned varints (not compact +1). Fields are written in
/// ascending tag order as the protocol requires; tag numbers must be unique.
pub fn Encoder::write_tagged_fields(
  self : Encoder,
  tags : Array[(Int, Bytes)],
) -> Unit {
  let sorted = tags.copy()
  sorted.sort_by(fn(a, b) { a.0.compare(b.0) })
  self.write_uvarint(sorted.length().reinterpret_as_uint())
  for pair in sorted {
    let (tag, data) = pair
    self.write_uvarint(tag.reinterpret_as_uint())
    self.write_uvarint(data.length().reinterpret_as_uint())
    self.write_bytes(data)
  }
}

///|
pub struct Decoder {
  data : Bytes
  mut pos : Int
}

///|
pub fn Decoder::new(data : Bytes, start? : Int = 0) -> Decoder {
  { data, pos: start, }
}

///|
pub fn Decoder::remaining(self : Decoder) -> Int {
  self.data.length() - self.pos
}

///|
/// Bounds check computed by subtraction (never overflows); negative
/// lengths from malformed input raise instead of rewinding or panicking.
fn Decoder::need(self : Decoder, n : Int) -> Unit raise DecodeError {
  if n < 0 || n > self.data.length() - self.pos {
    raise DecodeError::UnexpectedEof
  }
}

///|
pub fn Decoder::read_byte(self : Decoder) -> Byte raise DecodeError {
  self.need(1)
  let b = self.data[self.pos]
  self.pos += 1
  b
}

///|
pub fn Decoder::read_bool(self : Decoder) -> Bool raise DecodeError {
  self.read_byte() != b'\x00'
}

///|
pub fn Decoder::read_i8(self : Decoder) -> Int raise DecodeError {
  let v = self.read_byte().to_int()
  if v >= 128 {
    v - 256
  } else {
    v
  }
}

///|
pub fn Decoder::read_i16(self : Decoder) -> Int raise DecodeError {
  self.need(2)
  let v = (self.data[self.pos].to_int() << 8) | self.data[self.pos + 1].to_int()
  self.pos += 2
  Int16::from_int(v).to_int()
}

///|
pub fn Decoder::read_i32(self : Decoder) -> Int raise DecodeError {
  self.need(4)
  let mut v = 0
  for _ in 0..<4 {
    v = (v << 8) | self.data[self.pos].to_int()
    self.pos += 1
  }
  v
}

///|
pub fn Decoder::read_i64(self : Decoder) -> Int64 raise DecodeError {
  self.need(8)
  let mut v = 0L
  for _ in 0..<8 {
    v = (v << 8) | Int64::from_int(self.data[self.pos].to_int())
    self.pos += 1
  }
  v
}

///|
/// FLOAT64: big-endian IEEE 754 double.
pub fn Decoder::read_f64(self : Decoder) -> Double raise DecodeError {
  self.need(8)
  let mut bits = 0UL
  for _ in 0..<8 {
    bits = (bits << 8) |
      Int64::from_int(self.data[self.pos].to_int()).reinterpret_as_uint64()
    self.pos += 1
  }
  bits.reinterpret_as_double()
}

///|
/// Skip n bytes. Negative n and overruns raise; bounds are computed by
/// subtraction so a malformed huge length cannot overflow the check.
pub fn Decoder::skip(self : Decoder, n : Int) -> Unit raise DecodeError {
  if n < 0 || n > self.data.length() - self.pos {
    raise DecodeError::UnexpectedEof
  }
  self.pos += n
}

///|
pub fn Decoder::read_bytes(self : Decoder, n : Int) -> Bytes raise DecodeError {
  self.need(n)
  let copy = self.data[self.pos:self.pos + n].to_owned()
  self.pos += n
  copy
}

///|
/// Zero-copy view over the next `n` bytes of the backing buffer — no
/// copy is made, so callers that only inspect or immediately re-read the
/// bytes (e.g. blobbing a whole compressed records region before
/// building the owned `Bytes` a decompressor needs) avoid a defensive
/// copy. The returned `BytesView` borrows `self.data` and stays valid
/// while that `Bytes` outlives the decoder.
pub fn Decoder::read_bytes_view(
  self : Decoder,
  n : Int,
) -> BytesView raise DecodeError {
  self.need(n)
  let view = self.data[self.pos:self.pos + n]
  self.pos += n
  view
}

///|
fn Decoder::read_raw_uvarint(self : Decoder) -> UInt raise DecodeError {
  let mut result = 0U
  for shift = 0; ; shift = shift + 7 {
    if shift >= 35 {
      raise DecodeError::Malformed("uvarint too long")
    }
    let b = self.read_byte().to_int().reinterpret_as_uint()
    result = result | ((b & 0x7FU) << shift)
    if (b & 0x80U) == 0U {
      break
    }
  }
  result
}

///|
/// Unsigned LEB128 varint.
pub fn Decoder::read_uvarint(self : Decoder) -> Int raise DecodeError {
  let v = self.read_raw_uvarint()
  if v > 0x7FFFFFFFU {
    raise DecodeError::Malformed("uvarint overflows int32")
  }
  v.reinterpret_as_int()
}

///|
/// Zig-zag varint.
pub fn Decoder::read_varint(self : Decoder) -> Int raise DecodeError {
  let u = self.read_raw_uvarint()
  ((u >> 1) ^ (0U - (u & 1U))).reinterpret_as_int()
}

///|
/// Zig-zag varlong.
pub fn Decoder::read_varlong(self : Decoder) -> Int64 raise DecodeError {
  let mut result = 0UL
  for shift = 0; ; shift = shift + 7 {
    if shift >= 70 {
      raise DecodeError::Malformed("uvarint too long")
    }
    let b = self.read_byte().to_int64().reinterpret_as_uint64()
    result = result | ((b & 0x7FUL) << shift)
    if (b & 0x80UL) == 0UL {
      break
    }
  }
  ((result >> 1) ^ (0UL - (result & 1UL))).reinterpret_as_int64()
}

///|
/// Compact length prefix (value + 1); returns -1 for the null marker.
pub fn Decoder::read_compact_len(self : Decoder) -> Int raise DecodeError {
  self.read_uvarint() - 1
}

///|
pub fn Decoder::read_compact_string(self : Decoder) -> String raise DecodeError {
  match self.read_compact_nullable_string() {
    Some(s) => s
    None => raise DecodeError::Malformed("unexpected null compact string")
  }
}

///|
pub fn Decoder::read_compact_nullable_string(
  self : Decoder,
) -> String? raise DecodeError {
  let len = self.read_compact_len()
  if len < 0 {
    None
  } else {
    Some(@utf8.decode_lossy(self.read_bytes(len)[:]))
  }
}

///|
/// Legacy NULLABLE_STRING: INT16 length, -1 = null.
pub fn Decoder::read_nullable_string(
  self : Decoder,
) -> String? raise DecodeError {
  let len = self.read_i16()
  if len < 0 {
    None
  } else {
    Some(@utf8.decode_lossy(self.read_bytes(len)[:]))
  }
}

///|
/// Read and discard a tag buffer.
pub fn Decoder::skip_tag_buffer(self : Decoder) -> Unit raise DecodeError {
  let count = self.read_uvarint()
  for _ in 0..