// SPDX-License-Identifier: AGPL-3.0-only
// Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/cortexpb/timeseries.go
// Provenance-includes-license: Apache-2.0
// Provenance-includes-copyright: The Cortex Authors.

package mimirpb

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"unsafe"

	"github.com/prometheus/prometheus/model/labels"
)

const (
	expectedTimeseries         = 100
	expectedLabels             = 20
	expectedSamplesPerSeries   = 10
	expectedExemplarsPerSeries = 1
)

var (

	/*
		We cannot pool these as pointer-to-slice because the place we use them is in WriteRequest which is generated from Protobuf
		and we don't have an option to make it a pointer. There is overhead here 24 bytes of garbage every time a PreallocTimeseries
		is re-used. But since the slices are far far larger, we come out ahead.
	*/
	slicePool = sync.Pool{
		New: func() interface{} {
			return make([]PreallocTimeseries, 0, expectedTimeseries)
		},
	}

	timeSeriesPool = sync.Pool{
		New: func() interface{} {
			return &TimeSeries{
				Labels:    make([]LabelAdapter, 0, expectedLabels),
				Samples:   make([]Sample, 0, expectedSamplesPerSeries),
				Exemplars: make([]Exemplar, 0, expectedExemplarsPerSeries),
			}
		},
	}

	// yoloSlicePool is a pool of byte slices which are used to back the yoloStrings of this package.
	yoloSlicePool = sync.Pool{
		New: func() interface{} {
			// The initial cap of 200 is an arbitrary number which has been chosen because the default
			// of 0 is guaranteed to be insufficient, so any number greater than 0 would be better.
			// 200 should be enough to back all the strings of one TimeSeries in many cases.
			val := make([]byte, 0, 200)
			return &val
		},
	}
)

// PreallocWriteRequest is a WriteRequest which preallocs slices on Unmarshal.
type PreallocWriteRequest struct {
	WriteRequest
}

// Unmarshal implements proto.Message.
func (p *PreallocWriteRequest) Unmarshal(dAtA []byte) error {
	p.Timeseries = PreallocTimeseriesSliceFromPool()
	return p.WriteRequest.Unmarshal(dAtA)
}

// PreallocTimeseries is a TimeSeries which preallocs slices on Unmarshal.
type PreallocTimeseries struct {
	*TimeSeries

	// yoloSlice may contain a reference to the byte slice which was used to back the yolo strings
	// of the properties and sub-properties of this TimeSeries.
	// If it is set to a non-nil value then it must be returned to the yoloSlicePool on cleanup,
	// if it is set to nil then it can be ignored because the backing byte slice came from somewhere else.
	yoloSlice *[]byte
}

// Unmarshal implements proto.Message.
func (p *PreallocTimeseries) Unmarshal(dAtA []byte) error {
	p.TimeSeries = TimeseriesFromPool()
	return p.TimeSeries.Unmarshal(dAtA)
}

// LabelAdapter is a labels.Label that can be marshalled to/from protos.
type LabelAdapter labels.Label

// Marshal implements proto.Marshaller.
func (bs *LabelAdapter) Marshal() ([]byte, error) {
	size := bs.Size()
	buf := make([]byte, size)
	n, err := bs.MarshalToSizedBuffer(buf[:size])
	if err != nil {
		return nil, err
	}
	return buf[:n], err
}

func (bs *LabelAdapter) MarshalTo(dAtA []byte) (int, error) {
	size := bs.Size()
	return bs.MarshalToSizedBuffer(dAtA[:size])
}

// MarshalTo implements proto.Marshaller.
func (bs *LabelAdapter) MarshalToSizedBuffer(buf []byte) (n int, err error) {
	ls := (*labels.Label)(bs)
	i := len(buf)
	if len(ls.Value) > 0 {
		i -= len(ls.Value)
		copy(buf[i:], ls.Value)
		i = encodeVarintMimir(buf, i, uint64(len(ls.Value)))
		i--
		buf[i] = 0x12
	}
	if len(ls.Name) > 0 {
		i -= len(ls.Name)
		copy(buf[i:], ls.Name)
		i = encodeVarintMimir(buf, i, uint64(len(ls.Name)))
		i--
		buf[i] = 0xa
	}
	return len(buf) - i, nil
}

// Unmarshal a LabelAdapter, implements proto.Unmarshaller.
// NB this is a copy of the autogenerated code to unmarshal a LabelPair,
// with the byte copying replaced with a yoloString.
func (bs *LabelAdapter) Unmarshal(dAtA []byte) error {
	l := len(dAtA)
	iNdEx := 0
	for iNdEx < l {
		preIndex := iNdEx
		var wire uint64
		for shift := uint(0); ; shift += 7 {
			if shift >= 64 {
				return ErrIntOverflowMimir
			}
			if iNdEx >= l {
				return io.ErrUnexpectedEOF
			}
			b := dAtA[iNdEx]
			iNdEx++
			wire |= uint64(b&0x7F) << shift
			if b < 0x80 {
				break
			}
		}
		fieldNum := int32(wire >> 3)
		wireType := int(wire & 0x7)
		if wireType == 4 {
			return fmt.Errorf("proto: LabelPair: wiretype end group for non-group")
		}
		if fieldNum <= 0 {
			return fmt.Errorf("proto: LabelPair: illegal tag %d (wire type %d)", fieldNum, wire)
		}
		switch fieldNum {
		case 1:
			if wireType != 2 {
				return fmt.Errorf("proto: wrong wireType = %d for field Name", wireType)
			}
			var byteLen int
			for shift := uint(0); ; shift += 7 {
				if shift >= 64 {
					return ErrIntOverflowMimir
				}
				if iNdEx >= l {
					return io.ErrUnexpectedEOF
				}
				b := dAtA[iNdEx]
				iNdEx++
				byteLen |= int(b&0x7F) << shift
				if b < 0x80 {
					break
				}
			}
			if byteLen < 0 {
				return ErrInvalidLengthMimir
			}
			postIndex := iNdEx + byteLen
			if postIndex < 0 {
				return ErrInvalidLengthMimir
			}
			if postIndex > l {
				return io.ErrUnexpectedEOF
			}
			bs.Name = yoloString(dAtA[iNdEx:postIndex])
			iNdEx = postIndex
		case 2:
			if wireType != 2 {
				return fmt.Errorf("proto: wrong wireType = %d for field Value", wireType)
			}
			var byteLen int
			for shift := uint(0); ; shift += 7 {
				if shift >= 64 {
					return ErrIntOverflowMimir
				}
				if iNdEx >= l {
					return io.ErrUnexpectedEOF
				}
				b := dAtA[iNdEx]
				iNdEx++
				byteLen |= int(b&0x7F) << shift
				if b < 0x80 {
					break
				}
			}
			if byteLen < 0 {
				return ErrInvalidLengthMimir
			}
			postIndex := iNdEx + byteLen
			if postIndex < 0 {
				return ErrInvalidLengthMimir
			}
			if postIndex > l {
				return io.ErrUnexpectedEOF
			}
			bs.Value = yoloString(dAtA[iNdEx:postIndex])
			iNdEx = postIndex
		default:
			iNdEx = preIndex
			skippy, err := skipMimir(dAtA[iNdEx:])
			if err != nil {
				return err
			}
			if skippy < 0 {
				return ErrInvalidLengthMimir
			}
			if (iNdEx + skippy) < 0 {
				return ErrInvalidLengthMimir
			}
			if (iNdEx + skippy) > l {
				return io.ErrUnexpectedEOF
			}
			iNdEx += skippy
		}
	}

	if iNdEx > l {
		return io.ErrUnexpectedEOF
	}
	return nil
}

func yoloString(buf []byte) string {
	return *((*string)(unsafe.Pointer(&buf)))
}

// Size implements proto.Sizer.
func (bs *LabelAdapter) Size() (n int) {
	ls := (*labels.Label)(bs)
	if bs == nil {
		return 0
	}
	var l int
	_ = l
	l = len(ls.Name)
	if l > 0 {
		n += 1 + l + sovMimir(uint64(l))
	}
	l = len(ls.Value)
	if l > 0 {
		n += 1 + l + sovMimir(uint64(l))
	}
	return n
}

// Equal implements proto.Equaler.
func (bs *LabelAdapter) Equal(other LabelAdapter) bool {
	return bs.Name == other.Name && bs.Value == other.Value
}

// Compare implements proto.Comparer.
func (bs *LabelAdapter) Compare(other LabelAdapter) int {
	if c := strings.Compare(bs.Name, other.Name); c != 0 {
		return c
	}
	return strings.Compare(bs.Value, other.Value)
}

// PreallocTimeseriesSliceFromPool retrieves a slice of PreallocTimeseries from a sync.Pool.
// ReuseSlice should be called once done.
func PreallocTimeseriesSliceFromPool() []PreallocTimeseries {
	return slicePool.Get().([]PreallocTimeseries)
}

// ReuseSlice puts the slice back into a sync.Pool for reuse.
func ReuseSlice(ts []PreallocTimeseries) {
	for i := range ts {
		ReuseTimeseries(ts[i].TimeSeries)

		if ts[i].yoloSlice != nil {
			reuseYoloSlice(ts[i].yoloSlice)
			ts[i].yoloSlice = nil
		}
	}

	slicePool.Put(ts[:0]) //nolint:staticcheck //see comment on slicePool for more details
}

// TimeseriesFromPool retrieves a pointer to a TimeSeries from a sync.Pool.
// ReuseTimeseries should be called once done, unless ReuseSlice was called on the slice that contains this TimeSeries.
func TimeseriesFromPool() *TimeSeries {
	return timeSeriesPool.Get().(*TimeSeries)
}

// ReuseTimeseries puts the timeseries back into a sync.Pool for reuse.
func ReuseTimeseries(ts *TimeSeries) {
	// Name and Value may point into a large gRPC buffer, so clear the reference to allow GC
	for i := 0; i < len(ts.Labels); i++ {
		ts.Labels[i].Name = ""
		ts.Labels[i].Value = ""
	}
	ts.Labels = ts.Labels[:0]
	ts.Samples = ts.Samples[:0]
	// Name and Value may point into a large gRPC buffer, so clear the reference in each exemplar to allow GC
	for i := range ts.Exemplars {
		for j := range ts.Exemplars[i].Labels {
			ts.Exemplars[i].Labels[j].Name = ""
			ts.Exemplars[i].Labels[j].Value = ""
		}
	}
	ts.Exemplars = ts.Exemplars[:0]
	timeSeriesPool.Put(ts)
}

func yoloSliceFromPool() *[]byte {
	return yoloSlicePool.Get().(*[]byte)
}

func reuseYoloSlice(val *[]byte) {
	*val = (*val)[:0]
	yoloSlicePool.Put(val)
}

// DeepCopyTimeseries copies the timeseries of one PreallocTimeseries into another one.
// It copies all the properties, sub-properties and strings by value to ensure that the two timeseries are not sharing
// anything after the deep copying.
// The returned PreallocTimeseries has a yoloSlice property which should be returned to the yoloSlicePool on cleanup.
func DeepCopyTimeseries(dst, src PreallocTimeseries, keepExemplars bool) PreallocTimeseries {
	if dst.TimeSeries == nil {
		dst.TimeSeries = TimeseriesFromPool()
	}

	srcTs := src.TimeSeries
	dstTs := dst.TimeSeries

	// Prepare a buffer which is large enough to hold all the label names and values of src.
	requiredYoloSliceCap := countTotalLabelLen(src.TimeSeries, keepExemplars)
	dst.yoloSlice = yoloSliceFromPool()
	buf := ensureCap(dst.yoloSlice, requiredYoloSliceCap)

	// Copy the time series labels by using the prepared buffer.
	dst.TimeSeries.Labels, buf = copyToYoloLabels(buf, dstTs.Labels, srcTs.Labels)

	// Copy the samples.
	if cap(dst.TimeSeries.Samples) < len(src.TimeSeries.Samples) {
		dstTs.Samples = make([]Sample, len(src.Samples))
	} else {
		dstTs.Samples = dstTs.Samples[:len(src.Samples)]
	}
	copy(dstTs.Samples, srcTs.Samples)

	// Prepare the slice of exemplars.
	if keepExemplars {
		if cap(dstTs.Exemplars) < len(srcTs.Exemplars) {
			dstTs.Exemplars = make([]Exemplar, len(srcTs.Exemplars))
		} else {
			dstTs.Exemplars = dstTs.Exemplars[:len(srcTs.Exemplars)]
		}

		for exemplarIdx := range src.Exemplars {
			// Copy the exemplar labels by using the prepared buffer.
			dstTs.Exemplars[exemplarIdx].Labels, buf = copyToYoloLabels(buf, dstTs.Exemplars[exemplarIdx].Labels, src.Exemplars[exemplarIdx].Labels)

			// Copy the other exemplar properties.
			dstTs.Exemplars[exemplarIdx].Value = src.Exemplars[exemplarIdx].Value
			dstTs.Exemplars[exemplarIdx].TimestampMs = src.Exemplars[exemplarIdx].TimestampMs
		}
	} else {
		dstTs.Exemplars = dstTs.Exemplars[:0]
	}

	return dst
}

// ensureCap takes a pointer to a byte slice and ensures that the capacity of the referred slice is at least equal to
// the given capacity, if not then the byte slice referred to by the pointer gets replaced with a new, larger one.
// The return value is the byte slice which is now referred by the given pointer which has at least the given capacity.
func ensureCap(bufRef *[]byte, requiredCap int) []byte {
	buf := *bufRef
	if cap(buf) >= requiredCap {
		return buf[:0]
	}

	buf = make([]byte, 0, requiredCap)
	*bufRef = buf
	return buf
}

// countTotalLabelLen takes a time series and calculates the sum of the lengths of all label names and values.
func countTotalLabelLen(ts *TimeSeries, includeExemplars bool) int {
	var labelLen int
	for _, label := range ts.Labels {
		labelLen += len(label.Name) + len(label.Value)
	}

	if includeExemplars {
		for _, exemplar := range ts.Exemplars {
			for _, label := range exemplar.Labels {
				labelLen += len(label.Name) + len(label.Value)
			}
		}
	}

	return labelLen
}

// copyToYoloLabels copies the values of src to dst, it uses the given buffer to store all the string values in it.
// The returned buffer is the remainder of the given buffer, which remains unused after the copying is complete.
func copyToYoloLabels(buf []byte, dst, src []LabelAdapter) ([]LabelAdapter, []byte) {
	if cap(dst) < len(src) {
		dst = make([]LabelAdapter, len(src))
	} else {
		dst = dst[:len(src)]
	}

	for labelIdx := range src {
		dst[labelIdx].Name, buf = copyToYoloString(buf, src[labelIdx].Name)
		dst[labelIdx].Value, buf = copyToYoloString(buf, src[labelIdx].Value)
	}

	return dst, buf
}

// copyToYoloString takes a string and creates a new string which uses the given buffer as underlying byte array.
// It requires that the buffer has a capacitity which is greater than or equal to the length of the source string.
func copyToYoloString(buf []byte, src string) (string, []byte) {
	buf = buf[:len(src)]
	copy(buf, *((*[]byte)(unsafe.Pointer(&src))))
	return yoloString(buf), buf[len(buf):]
}
