// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package memberlist

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/go-msgpack/v2/codec"
	"github.com/sean-/seed"
)

// pushPullScale is the minimum number of nodes
// before we start scaling the push/pull timing. The scale
// effect is the log2(Nodes) - log2(pushPullScale). This means
// that the 33rd node will cause us to double the interval,
// while the 65th will triple it.
const pushPullScaleThreshold = 32

func init() {
	_, _ = seed.Init()
}

// Decode reverses the encode operation on a byte slice input
func decode(buf []byte, out interface{}) error {
	r := bytes.NewReader(buf)
	hd := codec.MsgpackHandle{}
	dec := codec.NewDecoder(r, &hd)
	return dec.Decode(out)
}

// pushPullBufPool recycles *bytes.Buffer values used by decryptRemoteState
// to stage cipher text as it streams in from the peer. Push-pull state
// can approach maxPushStateBytes (20 MiB), so pooling here lets the
// io.CopyN growth steps amortize across receives instead of starting
// from zero capacity every time.
//
// Pool callers MUST releasePushPullBuffer once they are done with the
// returned buffer's bytes.
var pushPullBufPool = sync.Pool{
	New: func() any {
		return new(bytes.Buffer)
	},
}

// maxPooledPushPullBufCap bounds the capacity of buffers retained by
// pushPullBufPool. Sized to cover maxPushStateBytes (20 MiB) plus headroom
// for the AEAD overhead added by encryptLocalState. Buffers that grow
// above this — degenerate or attack payloads — are dropped on release.
//
// This cap is for pool sizing only; it is NOT a security boundary. Size
// enforcement against malicious peers lives at decryptRemoteState's
// explicit `moreBytes > maxPushStateBytes` check — do not rely
// on this constant to reject oversized incoming state.
const maxPooledPushPullBufCap = 32 * 1024 * 1024

// pushPullBufInitialCap is the initial backing-slice capacity for fresh
// buffers from pushPullBufPool. Push-pull state on a small cluster fits
// well under 64 KiB, so 64 KiB lets typical encodes proceed without the
// first few growth steps. Buffers grow naturally beyond this when a peer
// has more state to ship.
const pushPullBufInitialCap = 64 * 1024

func getPushPullBuffer() *bytes.Buffer {
	b := pushPullBufPool.Get().(*bytes.Buffer)
	b.Reset()
	if b.Cap() < pushPullBufInitialCap {
		b.Grow(pushPullBufInitialCap)
	}
	return b
}

func releasePushPullBuffer(b *bytes.Buffer) {
	if b.Cap() > maxPooledPushPullBufCap {
		return
	}
	b.Reset()
	pushPullBufPool.Put(b)
}

// encode writes an encoded object and returns a freshly-allocated byte
// slice owned by the caller.
//
// The input `in` is not retained by the returned slice.
func encode(msgType messageType, in any, msgpackUseNewTimeFormat bool) ([]byte, error) {
	return encodeWithSizeHint(msgType, in, msgpackUseNewTimeFormat, 0)
}

// encodeWithSizeHint is like encode but pre-grows the internal working
// buffer to sizeHint bytes when sizeHint > 0. Use it from call sites
// that know the upcoming msgpack output size up front.
// sizeHint == 0 makes it equivalent to encode().
func encodeWithSizeHint(msgType messageType, in any, msgpackUseNewTimeFormat bool, sizeHint int) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	if sizeHint > 0 {
		buf.Grow(sizeHint)
	}
	buf.WriteByte(uint8(msgType))
	hd := codec.MsgpackHandle{}
	hd.TimeNotBuiltin = !msgpackUseNewTimeFormat

	if err := codec.NewEncoder(buf, &hd).Encode(in); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// Returns a random offset between 0 and n
func randomOffset(n int) int {
	if n == 0 {
		return 0
	}
	return int(rand.Uint32() % uint32(n))
}

// suspicionTimeout computes the timeout that should be used when
// a node is suspected
func suspicionTimeout(suspicionMult, n int, interval time.Duration) time.Duration {
	nodeScale := math.Max(1.0, math.Log10(math.Max(1.0, float64(n))))
	// multiply by 1000 to keep some precision because time.Duration is an int64 type
	timeout := time.Duration(suspicionMult) * time.Duration(nodeScale*1000) * interval / 1000
	return timeout
}

// retransmitLimit computes the limit of retransmissions
func retransmitLimit(retransmitMult, n int) int {
	nodeScale := math.Ceil(math.Log10(float64(n + 1)))
	limit := retransmitMult * int(nodeScale)
	return limit
}

// shuffleNodes randomly shuffles the input nodes using the Fisher-Yates shuffle
func shuffleNodes(nodes []*NodeState) {
	n := len(nodes)
	rand.Shuffle(n, func(i, j int) {
		nodes[i], nodes[j] = nodes[j], nodes[i]
	})
}

// pushPushScale is used to scale the time interval at which push/pull
// syncs take place. It is used to prevent network saturation as the
// cluster size grows
func pushPullScale(interval time.Duration, n int) time.Duration {
	// Don't scale until we cross the threshold
	if n <= pushPullScaleThreshold {
		return interval
	}

	multiplier := math.Ceil(math.Log2(float64(n))-math.Log2(pushPullScaleThreshold)) + 1.0
	return time.Duration(multiplier) * interval
}

// moveDeadNodes moves dead and left nodes that that have not changed during the gossipToTheDeadTime interval
// to the end of the slice and returns the index of the first moved node.
func moveDeadNodes(nodes []*NodeState, gossipToTheDeadTime time.Duration) int {
	numDead := 0
	n := len(nodes)
	for i := 0; i < n-numDead; i++ {
		if !nodes[i].DeadOrLeft() {
			continue
		}

		// Respect the gossip to the dead interval
		if time.Since(nodes[i].StateChange) <= gossipToTheDeadTime {
			continue
		}

		// Move this node to the end
		nodes[i], nodes[n-numDead-1] = nodes[n-numDead-1], nodes[i]
		numDead++
		i--
	}
	return n - numDead
}

// kRandomNodes is used to select up to k random Nodes, excluding any nodes where
// the exclude function returns true. It is possible that less than k nodes are
// returned.
func kRandomNodes(k int, nodes []*NodeState, delegate NodeSelectionDelegate, exclude func(*NodeState) bool) []Node {
	kNodes := make([]Node, 0, k)

	// Filter the nodes using the delegate. This allows downstream projects
	// to implement custom routing logics (e.g. zone-aware gossiping).
	if delegate != nil {
		selected, preferred := delegate.SelectNodes(nodes)
		nodes = selected

		// Add the preferred node first to guarantee it's in the result set.
		if preferred != nil && k > 0 {
			// Ensure it's not excluded by the filter.
			if exclude == nil || !exclude(preferred) {
				kNodes = append(kNodes, preferred.Node)
			}
		}
	}

	n := len(nodes)

OUTER:
	// Probe up to 3*n times, with large n this is not necessary
	// since k << n, but with small n we want search to be
	// exhaustive
	for i := 0; i < 3*n && len(kNodes) < k; i++ {
		// Get random nodeState
		idx := randomOffset(n)
		state := nodes[idx]

		// Give the filter a shot at it.
		if exclude != nil && exclude(state) {
			continue OUTER
		}

		// Check if we have this node already
		for j := 0; j < len(kNodes); j++ {
			if state.Name == kNodes[j].Name {
				continue OUTER
			}
		}

		// Append the node
		kNodes = append(kNodes, state.Node)
	}
	return kNodes
}

// makeCompoundMessages takes a list of messages and packs them into
// one or multiple messages based on the limitations of compound
// messages (255 messages each, 64KB max message size).
//
// The input msgs's underlying array is rearranged in place: small
// messages are kept and oversized ones are passed through to the
// output. Returned slices have mixed ownership: compounds produced
// by makeCompoundMessage are freshly-allocated; oversized messages
// alias the input. Callers must therefore treat returned slices,
// and the underlying bytes of every msgs entry, as read-only until
// the returned slices are consumed.
func makeCompoundMessages(msgs [][]byte) [][]byte {
	const (
		maxMsgs      = math.MaxUint8
		maxMsgLength = math.MaxUint16
	)

	// Optimistically assume there will be no big message.
	bufs := make([][]byte, 0, (len(msgs)+(maxMsgs-1))/maxMsgs)

	// Do not add to a compound message any message bigger than the max message length
	// we can store.
	r, w := 0, 0
	for r < len(msgs) {
		if len(msgs[r]) <= maxMsgLength {
			// Keep it.
			msgs[w] = msgs[r]
			r++
			w++
			continue
		}

		// Oversized message — send it alone. Passes the input slice
		// through directly; callers do not mutate the result.
		bufs = append(bufs, msgs[r])
		r++
	}
	msgs = msgs[:w]

	// Group remaining messages in compound message(s).
	for ; len(msgs) > maxMsgs; msgs = msgs[maxMsgs:] {
		bufs = append(bufs, makeCompoundMessage(msgs[:maxMsgs]))
	}
	if len(msgs) > 0 {
		bufs = append(bufs, makeCompoundMessage(msgs))
	}

	return bufs
}

// makeCompoundMessage takes a list of messages and generates
// a single compound message containing all of them.
// Returns a freshly-allocated byte slice owned by the caller.
func makeCompoundMessage(msgs [][]byte) []byte {
	// Pre-size the buffer to the exact compound length (1 type byte +
	// 1 count byte + 2 bytes per length prefix + the message bodies).
	total := 2 + 2*len(msgs)
	for _, m := range msgs {
		total += len(m)
	}
	buf := bytes.NewBuffer(nil)
	buf.Grow(total)

	// Write out the type
	buf.WriteByte(uint8(compoundMsg))

	// Write out the number of message
	buf.WriteByte(uint8(len(msgs)))

	// Add the message lengths
	for _, m := range msgs {
		_ = binary.Write(buf, binary.BigEndian, uint16(len(m)))
	}

	// Append the messages
	for _, m := range msgs {
		buf.Write(m)
	}

	return buf.Bytes()
}

// decodeCompoundMessage splits a compound message and returns
// the slices of individual messages. Also returns the number
// of truncated messages and any potential error
func decodeCompoundMessage(buf []byte) (trunc int, parts [][]byte, err error) {
	if len(buf) < 1 {
		err = fmt.Errorf("missing compound length byte")
		return
	}
	numParts := int(buf[0])
	buf = buf[1:]

	// Check we have enough bytes
	if len(buf) < numParts*2 {
		err = fmt.Errorf("truncated len slice")
		return
	}

	// Decode the lengths
	lengths := make([]uint16, numParts)
	for i := 0; i < numParts; i++ {
		lengths[i] = binary.BigEndian.Uint16(buf[i*2 : i*2+2])
	}
	buf = buf[numParts*2:]

	// Split each message
	for idx, msgLen := range lengths {
		if len(buf) < int(msgLen) {
			trunc = numParts - idx
			return
		}

		// Extract the slice, seek past on the buffer
		slice := buf[:msgLen]
		buf = buf[msgLen:]
		parts = append(parts, slice)
	}
	return
}

// joinHostPort returns the host:port form of an address, for use with a
// transport.
func joinHostPort(host string, port uint16) string {
	return net.JoinHostPort(host, strconv.Itoa(int(port)))
}

// hasPort is given a string of the form "host", "host:port", "ipv6::address",
// or "[ipv6::address]:port", and returns true if the string includes a port.
func hasPort(s string) bool {
	// IPv6 address in brackets.
	if strings.LastIndex(s, "[") == 0 {
		return strings.LastIndex(s, ":") > strings.LastIndex(s, "]")
	}

	// Otherwise the presence of a single colon determines if there's a port
	// since IPv6 addresses outside of brackets (count > 1) can't have a
	// port.
	return strings.Count(s, ":") == 1
}

// ensurePort makes sure the given string has a port number on it, otherwise it
// appends the given port as a default.
func ensurePort(s string, port int) string {
	if hasPort(s) {
		return s
	}

	// If this is an IPv6 address, the join call will add another set of
	// brackets, so we have to trim before we add the default port.
	s = strings.Trim(s, "[]")
	s = net.JoinHostPort(s, strconv.Itoa(port))
	return s
}
