package kafkas

import (
  "github.com/segmentio/kafka-go"
)

// This file is all implemented by https://github.com/segmentio/kafka-go/blob/main/message.go
const timestampSize = 8

func varIntLen(i int64) int {
  u := uint64((i << 1) ^ (i >> 63)) // zig-zag encoding
  n := 0

  for u >= 0x80 {
  	u >>= 7
  	n++
  }

  return n + 1
}

func varBytesLen(b []byte) int {
  return varIntLen(int64(len(b))) + len(b)
}

func varStringLen(s string) int {
  return varIntLen(int64(len(s))) + len(s)
}

func varArrayLen(n int, f func(int) int) int {
  size := varIntLen(int64(n))
  for i := 0; i < n; i++ {
  	size += f(i)
  }
  return size
}

func sizeofBytes(b []byte) int32 {
  return 4 + int32(len(b))
}

func size(msg *kafka.Message) int32 {
  return 4 + 1 + 1 + sizeofBytes(msg.Key) + sizeofBytes(msg.Value) + timestampSize
}

func headerSize(msg *kafka.Message) int {
  return varArrayLen(len(msg.Headers), func(i int) int {
  	h := &msg.Headers[i]
  	return varStringLen(h.Key) + varBytesLen(h.Value)
  })
}

// Make it public for external packages
func TotalSize(msg *kafka.Message) int32 {
  return int32(headerSize(msg)) + size(msg)
}