/*
This package contains help functions for stable marshaller. Their usage is
totally optional. One can implement fast stable marshaller without these
runtime function calls.
*/

package proto

import (
	"math"
	"math/bits"

	"google.golang.org/protobuf/encoding/protowire"
)

type (
	stableMarshaller interface {
		StableMarshal([]byte) []byte
		StableSize() int
	}
)

func BytesMarshal(field int, buf, v []byte) []byte {
	return bytesMarshal(field, buf, v)
}

func BytesSize(field int, v []byte) int {
	return bytesSize(field, v)
}

func bytesMarshal(field int, buf []byte, v []byte) []byte {
	if len(v) == 0 {
		return buf
	}
	return bytesMarshalNoCheck(field, buf, v)
}

func stringMarshal(field int, buf []byte, v string) []byte {
	if len(v) == 0 {
		return buf
	}
	return stringMarshalNoCheck(field, buf, v)
}

func stringMarshalNoCheck(field int, buf []byte, v string) []byte {
	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.BytesType)
	buf = protowire.AppendString(buf, v)
	return buf
}

func bytesMarshalNoCheck(field int, buf []byte, v []byte) []byte {
	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.BytesType)
	buf = protowire.AppendBytes(buf, v)
	return buf
}

func bytesSize[T ~[]byte | ~string](field int, v T) int {
	if len(v) == 0 {
		return 0
	}
	return bytesSizeNoCheck(field, v)
}

func bytesSizeNoCheck[T ~[]byte | ~string](field int, v T) int {
	return protowire.SizeGroup(protowire.Number(field), protowire.SizeBytes(len(v)))
}

func StringMarshal(field int, buf []byte, v string) []byte {
	return stringMarshal(field, buf, v)
}

func StringSize(field int, v string) int {
	return bytesSize(field, v)
}

func BoolMarshal(field int, buf []byte, v bool) []byte {
	if !v {
		return buf
	}

	const boolTrueValue = 0x1

	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.VarintType)
	buf = protowire.AppendVarint(buf, boolTrueValue)

	return buf
}

func BoolSize(field int, v bool) int {
	if !v {
		return 0
	}
	const boolLength = 1
	return protowire.SizeGroup(protowire.Number(field), boolLength)
}

func UInt64Marshal(field int, buf []byte, v uint64) []byte {
	if v == 0 {
		return buf
	}

	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.VarintType)
	buf = protowire.AppendVarint(buf, v)

	return buf
}

func UInt64Size(field int, v uint64) int {
	if v == 0 {
		return 0
	}
	return protowire.SizeGroup(protowire.Number(field), protowire.SizeVarint(v))
}

func Int64Marshal(field int, buf []byte, v int64) []byte {
	return UInt64Marshal(field, buf, uint64(v))
}

func Int64Size(field int, v int64) int {
	return UInt64Size(field, uint64(v))
}

func UInt32Marshal(field int, buf []byte, v uint32) []byte {
	return UInt64Marshal(field, buf, uint64(v))
}

func UInt32Size(field int, v uint32) int {
	return UInt64Size(field, uint64(v))
}

func Int32Marshal(field int, buf []byte, v int32) []byte {
	return UInt64Marshal(field, buf, uint64(v))
}

func Int32Size(field int, v int32) int {
	return UInt64Size(field, uint64(v))
}

func EnumMarshal(field int, buf []byte, v int32) []byte {
	return UInt64Marshal(field, buf, uint64(v))
}

func EnumSize(field int, v int32) int {
	return UInt64Size(field, uint64(v))
}

func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) []byte {
	for i := range v {
		buf = bytesMarshalNoCheck(field, buf, v[i])
	}

	return buf
}

func RepeatedBytesSize(field int, v [][]byte) (size int) {
	for i := range v {
		size += bytesSizeNoCheck(field, v[i])
	}

	return size
}

func RepeatedStringMarshal(field int, buf []byte, v []string) []byte {
	var offset int

	for i := range v {
		buf = stringMarshalNoCheck(field, buf[offset:], v[i])
	}

	return buf
}

func RepeatedStringSize(field int, v []string) (size int) {
	for i := range v {
		size += bytesSizeNoCheck(field, v[i])
	}

	return size
}

func repeatedUIntSize[T ~uint64 | ~int64 | ~uint32 | ~int32](field int, v []T) (size, arraySize int) {
	if len(v) == 0 {
		return 0, 0
	}

	for i := range v {
		arraySize += protowire.SizeVarint(uint64(v[i]))
	}

	size = protowire.SizeGroup(protowire.Number(field), protowire.SizeBytes(arraySize))

	return
}

func repeatedUIntMarshal[T ~uint64 | ~int64 | ~uint32 | ~int32](field int, buf []byte, v []T) []byte {
	if len(v) == 0 {
		return buf
	}

	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.BytesType)

	_, arrSize := repeatedUIntSize(field, v)
	buf = protowire.AppendVarint(buf, uint64(arrSize))
	for i := range v {
		buf = protowire.AppendVarint(buf, uint64(v[i]))
	}

	return buf
}

func RepeatedUInt64Marshal(field int, buf []byte, v []uint64) []byte {
	return repeatedUIntMarshal(field, buf, v)
}

func RepeatedUInt64Size(field int, v []uint64) (size, arraySize int) {
	return repeatedUIntSize(field, v)
}

func RepeatedInt64Marshal(field int, buf []byte, v []int64) []byte {
	return repeatedUIntMarshal(field, buf, v)
}

func RepeatedInt64Size(field int, v []int64) (size, arraySize int) {
	return repeatedUIntSize(field, v)
}

func RepeatedUInt32Marshal(field int, buf []byte, v []uint32) []byte {
	return repeatedUIntMarshal(field, buf, v)
}

func RepeatedUInt32Size(field int, v []uint32) (size, arraySize int) {
	return repeatedUIntSize(field, v)
}

func RepeatedInt32Marshal(field int, buf []byte, v []int32) []byte {
	return repeatedUIntMarshal(field, buf, v)
}

func RepeatedInt32Size(field int, v []int32) (size, arraySize int) {
	return repeatedUIntSize(field, v)
}

// VarUIntSize returns length of varint byte sequence for uint64 value 'x'.
func VarUIntSize(x uint64) int {
	return (bits.Len64(x|1) + 6) / 7
}

func NestedStructureMarshal[T stableMarshaller](field int64, buf []byte, v T) []byte {
	n := v.StableSize()
	if n == 0 {
		return buf
	}

	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.BytesType)
	buf = protowire.AppendVarint(buf, uint64(n))
	buf = v.StableMarshal(buf)

	return buf
}

func NestedStructureSize[T stableMarshaller](field int64, v T) (size int) {
	n := v.StableSize()
	if n == 0 {
		return 0
	}
	size = protowire.SizeGroup(protowire.Number(field), protowire.SizeBytes(n))
	return
}

func Fixed64Marshal(field int, buf []byte, v uint64) []byte {
	if v == 0 {
		return buf
	}

	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.Fixed64Type)
	buf = protowire.AppendFixed64(buf, v)

	return buf
}

func Fixed64Size(fNum int, v uint64) int {
	if v == 0 {
		return 0
	}
	return protowire.SizeGroup(protowire.Number(fNum), protowire.SizeFixed64())
}

func Float64Marshal(field int, buf []byte, v float64) []byte {
	if v == 0 {
		return buf
	}

	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.Fixed64Type)
	buf = protowire.AppendFixed64(buf, math.Float64bits(v))

	return buf
}

func Float64Size(fNum int, v float64) int {
	if v == 0 {
		return 0
	}
	return protowire.SizeGroup(protowire.Number(fNum), protowire.SizeFixed64())
}

// Fixed32Marshal encodes uint32 value to Protocol Buffers fixed32 field with specified number,
// and writes it to specified buffer. Returns number of bytes written.
//
// Panics if the buffer is undersized.
func Fixed32Marshal(field int, buf []byte, v uint32) []byte {
	if v == 0 {
		return buf
	}

	buf = protowire.AppendTag(buf, protowire.Number(field), protowire.Fixed32Type)
	buf = protowire.AppendFixed32(buf, v)

	return buf
}

// Fixed32Size returns number of bytes required to encode uint32 value to Protocol Buffers fixed32 field
// with specified number.
func Fixed32Size(fNum int, v uint32) int {
	if v == 0 {
		return 0
	}
	return protowire.SizeGroup(protowire.Number(fNum), protowire.SizeFixed32())
}