package msgp

import (
	"encoding/binary"
	"encoding/json"
	"errors"
	"math"
	"reflect"
	"time"
)

// ensure 'sz' extra bytes in 'b' btw len(b) and cap(b)
func ensure(b []byte, sz int) ([]byte, int) {
	l := len(b)
	c := cap(b)
	if c-l < sz {
		o := make([]byte, (2*c)+sz) // exponential growth
		n := copy(o, b)
		return o[:n+sz], n
	}
	return b[:l+sz], l
}

// AppendMapHeader appends a map header with the
// given size to the slice
func AppendMapHeader(b []byte, sz uint32) []byte {
	switch {
	case sz <= 15:
		return append(b, wfixmap(uint8(sz)))

	case sz <= math.MaxUint16:
		o, n := ensure(b, 3)
		prefixu16(o[n:], mmap16, uint16(sz))
		return o

	default:
		o, n := ensure(b, 5)
		prefixu32(o[n:], mmap32, sz)
		return o
	}
}

// AppendArrayHeader appends an array header with
// the given size to the slice
func AppendArrayHeader(b []byte, sz uint32) []byte {
	switch {
	case sz <= 15:
		return append(b, wfixarray(uint8(sz)))

	case sz <= math.MaxUint16:
		o, n := ensure(b, 3)
		prefixu16(o[n:], marray16, uint16(sz))
		return o

	default:
		o, n := ensure(b, 5)
		prefixu32(o[n:], marray32, sz)
		return o
	}
}

// AppendNil appends a 'nil' byte to the slice
func AppendNil(b []byte) []byte { return append(b, mnil) }

// AppendFloat appends a float to the slice as either float64
// or float32 when it represents the exact same value
func AppendFloat(b []byte, f float64) []byte {
	f32 := float32(f)
	if float64(f32) == f {
		return AppendFloat32(b, f32)
	}
	return AppendFloat64(b, f)
}

// AppendFloat64 appends a float64 to the slice
func AppendFloat64(b []byte, f float64) []byte {
	o, n := ensure(b, Float64Size)
	prefixu64(o[n:], mfloat64, math.Float64bits(f))
	return o
}

// AppendFloat32 appends a float32 to the slice
func AppendFloat32(b []byte, f float32) []byte {
	o, n := ensure(b, Float32Size)
	prefixu32(o[n:], mfloat32, math.Float32bits(f))
	return o
}

// AppendDuration appends a time.Duration to the slice
func AppendDuration(b []byte, d time.Duration) []byte {
	return AppendInt64(b, int64(d))
}

// AppendInt64 appends an int64 to the slice
func AppendInt64(b []byte, i int64) []byte {
	if i >= 0 {
		switch {
		case i <= math.MaxInt8:
			return append(b, wfixint(uint8(i)))
		case i <= math.MaxInt16:
			o, n := ensure(b, 3)
			putMint16(o[n:], int16(i))
			return o
		case i <= math.MaxInt32:
			o, n := ensure(b, 5)
			putMint32(o[n:], int32(i))
			return o
		default:
			o, n := ensure(b, 9)
			putMint64(o[n:], i)
			return o
		}
	}
	switch {
	case i >= -32:
		return append(b, wnfixint(int8(i)))
	case i >= math.MinInt8:
		o, n := ensure(b, 2)
		putMint8(o[n:], int8(i))
		return o
	case i >= math.MinInt16:
		o, n := ensure(b, 3)
		putMint16(o[n:], int16(i))
		return o
	case i >= math.MinInt32:
		o, n := ensure(b, 5)
		putMint32(o[n:], int32(i))
		return o
	default:
		o, n := ensure(b, 9)
		putMint64(o[n:], i)
		return o
	}
}

// AppendInt appends an int to the slice
func AppendInt(b []byte, i int) []byte { return AppendInt64(b, int64(i)) }

// AppendInt8 appends an int8 to the slice
func AppendInt8(b []byte, i int8) []byte { return AppendInt64(b, int64(i)) }

// AppendInt16 appends an int16 to the slice
func AppendInt16(b []byte, i int16) []byte { return AppendInt64(b, int64(i)) }

// AppendInt32 appends an int32 to the slice
func AppendInt32(b []byte, i int32) []byte { return AppendInt64(b, int64(i)) }

// AppendUint64 appends a uint64 to the slice
func AppendUint64(b []byte, u uint64) []byte {
	switch {
	case u <= (1<<7)-1:
		return append(b, wfixint(uint8(u)))

	case u <= math.MaxUint8:
		o, n := ensure(b, 2)
		putMuint8(o[n:], uint8(u))
		return o

	case u <= math.MaxUint16:
		o, n := ensure(b, 3)
		putMuint16(o[n:], uint16(u))
		return o

	case u <= math.MaxUint32:
		o, n := ensure(b, 5)
		putMuint32(o[n:], uint32(u))
		return o

	default:
		o, n := ensure(b, 9)
		putMuint64(o[n:], u)
		return o

	}
}

// AppendUint appends a uint to the slice
func AppendUint(b []byte, u uint) []byte { return AppendUint64(b, uint64(u)) }

// AppendUint8 appends a uint8 to the slice
func AppendUint8(b []byte, u uint8) []byte { return AppendUint64(b, uint64(u)) }

// AppendByte is analogous to AppendUint8
func AppendByte(b []byte, u byte) []byte { return AppendUint8(b, u) }

// AppendUint16 appends a uint16 to the slice
func AppendUint16(b []byte, u uint16) []byte { return AppendUint64(b, uint64(u)) }

// AppendUint32 appends a uint32 to the slice
func AppendUint32(b []byte, u uint32) []byte { return AppendUint64(b, uint64(u)) }

// AppendBytes appends bytes to the slice as MessagePack 'bin' data
func AppendBytes(b []byte, bts []byte) []byte {
	sz := len(bts)
	var o []byte
	var n int
	switch {
	case sz <= math.MaxUint8:
		o, n = ensure(b, 2+sz)
		prefixu8(o[n:], mbin8, uint8(sz))
		n += 2
	case sz <= math.MaxUint16:
		o, n = ensure(b, 3+sz)
		prefixu16(o[n:], mbin16, uint16(sz))
		n += 3
	default:
		o, n = ensure(b, 5+sz)
		prefixu32(o[n:], mbin32, uint32(sz))
		n += 5
	}
	return o[:n+copy(o[n:], bts)]
}

// AppendBytesHeader appends an 'bin' header with
// the given size to the slice.
func AppendBytesHeader(b []byte, sz uint32) []byte {
	var o []byte
	var n int
	switch {
	case sz <= math.MaxUint8:
		o, n = ensure(b, 2)
		prefixu8(o[n:], mbin8, uint8(sz))
		return o
	case sz <= math.MaxUint16:
		o, n = ensure(b, 3)
		prefixu16(o[n:], mbin16, uint16(sz))
		return o
	}
	o, n = ensure(b, 5)
	prefixu32(o[n:], mbin32, sz)
	return o
}

// AppendBool appends a bool to the slice
func AppendBool(b []byte, t bool) []byte {
	if t {
		return append(b, mtrue)
	}
	return append(b, mfalse)
}

// AppendString appends a string as a MessagePack 'str' to the slice
func AppendString(b []byte, s string) []byte {
	sz := len(s)
	var n int
	var o []byte
	switch {
	case sz <= 31:
		o, n = ensure(b, 1+sz)
		o[n] = wfixstr(uint8(sz))
		n++
	case sz <= math.MaxUint8:
		o, n = ensure(b, 2+sz)
		prefixu8(o[n:], mstr8, uint8(sz))
		n += 2
	case sz <= math.MaxUint16:
		o, n = ensure(b, 3+sz)
		prefixu16(o[n:], mstr16, uint16(sz))
		n += 3
	default:
		o, n = ensure(b, 5+sz)
		prefixu32(o[n:], mstr32, uint32(sz))
		n += 5
	}
	return o[:n+copy(o[n:], s)]
}

// AppendStringFromBytes appends a []byte
// as a MessagePack 'str' to the slice 'b.'
func AppendStringFromBytes(b []byte, str []byte) []byte {
	sz := len(str)
	var n int
	var o []byte
	switch {
	case sz <= 31:
		o, n = ensure(b, 1+sz)
		o[n] = wfixstr(uint8(sz))
		n++
	case sz <= math.MaxUint8:
		o, n = ensure(b, 2+sz)
		prefixu8(o[n:], mstr8, uint8(sz))
		n += 2
	case sz <= math.MaxUint16:
		o, n = ensure(b, 3+sz)
		prefixu16(o[n:], mstr16, uint16(sz))
		n += 3
	default:
		o, n = ensure(b, 5+sz)
		prefixu32(o[n:], mstr32, uint32(sz))
		n += 5
	}
	return o[:n+copy(o[n:], str)]
}

// AppendComplex64 appends a complex64 to the slice as a MessagePack extension
func AppendComplex64(b []byte, c complex64) []byte {
	o, n := ensure(b, Complex64Size)
	o[n] = mfixext8
	o[n+1] = Complex64Extension
	big.PutUint32(o[n+2:], math.Float32bits(real(c)))
	big.PutUint32(o[n+6:], math.Float32bits(imag(c)))
	return o
}

// AppendComplex128 appends a complex128 to the slice as a MessagePack extension
func AppendComplex128(b []byte, c complex128) []byte {
	o, n := ensure(b, Complex128Size)
	o[n] = mfixext16
	o[n+1] = Complex128Extension
	big.PutUint64(o[n+2:], math.Float64bits(real(c)))
	big.PutUint64(o[n+10:], math.Float64bits(imag(c)))
	return o
}

// AppendTime appends a time.Time to the slice as a MessagePack extension
func AppendTime(b []byte, t time.Time) []byte {
	o, n := ensure(b, TimeSize)
	t = t.UTC()
	o[n] = mext8
	o[n+1] = 12
	o[n+2] = TimeExtension
	putUnix(o[n+3:], t.Unix(), int32(t.Nanosecond()))
	return o
}

// AppendTimeExt will write t using the official msgpack extension spec.
// https://github.com/msgpack/msgpack/blob/master/spec.md#timestamp-extension-type
func AppendTimeExt(b []byte, t time.Time) []byte {
	// Time rounded towards zero.
	secPrec := t.Truncate(time.Second)
	remain := t.Sub(secPrec).Nanoseconds()
	asSecs := secPrec.Unix()
	switch {
	case remain == 0 && asSecs > 0 && asSecs <= math.MaxUint32:
		// 4 bytes
		o, n := ensure(b, 2+4)
		o[n+0] = mfixext4
		o[n+1] = byte(msgTimeExtension)
		binary.BigEndian.PutUint32(o[n+2:], uint32(asSecs))
		return o
	case asSecs < 0 || asSecs >= (1<<34):
		// 12 bytes
		o, n := ensure(b, 3+12)
		o[n+0] = mext8
		o[n+1] = 12
		o[n+2] = byte(msgTimeExtension)
		binary.BigEndian.PutUint32(o[n+3:], uint32(remain))
		binary.BigEndian.PutUint64(o[n+3+4:], uint64(asSecs))
		return o
	default:
		// 8 bytes
		o, n := ensure(b, 2+8)
		o[n+0] = mfixext8
		o[n+1] = byte(msgTimeExtension)
		binary.BigEndian.PutUint64(o[n+2:], uint64(asSecs)|(uint64(remain)<<34))
		return o
	}
}

// AppendMapStrStr appends a map[string]string to the slice
// as a MessagePack map with 'str'-type keys and values
func AppendMapStrStr(b []byte, m map[string]string) []byte {
	sz := uint32(len(m))
	b = AppendMapHeader(b, sz)
	for key, val := range m {
		b = AppendString(b, key)
		b = AppendString(b, val)
	}
	return b
}

// AppendMapStrIntf appends a map[string]interface{} to the slice
// as a MessagePack map with 'str'-type keys.
func AppendMapStrIntf(b []byte, m map[string]any) ([]byte, error) {
	sz := uint32(len(m))
	b = AppendMapHeader(b, sz)
	var err error
	for key, val := range m {
		b = AppendString(b, key)
		b, err = AppendIntf(b, val)
		if err != nil {
			return b, err
		}
	}
	return b, nil
}

// AppendIntf appends the concrete type of 'i' to the
// provided []byte. 'i' must be one of the following:
//   - 'nil'
//   - A bool, float, string, []byte, int, uint, or complex
//   - A map[string]T where T is another supported type
//   - A []T, where T is another supported type
//   - A *T, where T is another supported type
//   - A type that satisfies the msgp.Marshaler interface
//   - A type that satisfies the msgp.Extension interface
func AppendIntf(b []byte, i any) ([]byte, error) {
	if i == nil {
		return AppendNil(b), nil
	}

	// all the concrete types
	// for which we have methods
	switch i := i.(type) {
	case Marshaler:
		return i.MarshalMsg(b)
	case Extension:
		return AppendExtension(b, i)
	case bool:
		return AppendBool(b, i), nil
	case float32:
		return AppendFloat32(b, i), nil
	case float64:
		return AppendFloat64(b, i), nil
	case complex64:
		return AppendComplex64(b, i), nil
	case complex128:
		return AppendComplex128(b, i), nil
	case string:
		return AppendString(b, i), nil
	case []byte:
		return AppendBytes(b, i), nil
	case int8:
		return AppendInt8(b, i), nil
	case int16:
		return AppendInt16(b, i), nil
	case int32:
		return AppendInt32(b, i), nil
	case int64:
		return AppendInt64(b, i), nil
	case int:
		return AppendInt64(b, int64(i)), nil
	case uint:
		return AppendUint64(b, uint64(i)), nil
	case uint8:
		return AppendUint8(b, i), nil
	case uint16:
		return AppendUint16(b, i), nil
	case uint32:
		return AppendUint32(b, i), nil
	case uint64:
		return AppendUint64(b, i), nil
	case time.Time:
		return AppendTime(b, i), nil
	case time.Duration:
		return AppendDuration(b, i), nil
	case map[string]any:
		return AppendMapStrIntf(b, i)
	case map[string]string:
		return AppendMapStrStr(b, i), nil
	case json.Number:
		return AppendJSONNumber(b, i)
	case []any:
		b = AppendArrayHeader(b, uint32(len(i)))
		var err error
		for _, k := range i {
			b, err = AppendIntf(b, k)
			if err != nil {
				return b, err
			}
		}
		return b, nil
	}

	var err error
	v := reflect.ValueOf(i)
	switch v.Kind() {
	case reflect.Map:
		if v.Type().Key().Kind() != reflect.String {
			return b, errors.New("msgp: map keys must be strings")
		}
		ks := v.MapKeys()
		b = AppendMapHeader(b, uint32(len(ks)))
		for _, key := range ks {
			val := v.MapIndex(key)
			b = AppendString(b, key.String())
			b, err = AppendIntf(b, val.Interface())
			if err != nil {
				return nil, err
			}
		}
		return b, nil
	case reflect.Array, reflect.Slice:
		l := v.Len()
		b = AppendArrayHeader(b, uint32(l))
		for i := range l {
			b, err = AppendIntf(b, v.Index(i).Interface())
			if err != nil {
				return b, err
			}
		}
		return b, nil
	case reflect.Ptr:
		if v.IsNil() {
			return AppendNil(b), err
		}
		b, err = AppendIntf(b, v.Elem().Interface())
		return b, err
	default:
		return b, &ErrUnsupportedType{T: v.Type()}
	}
}

// AppendJSONNumber appends a json.Number to the slice.
// An error will be returned if the json.Number returns error as both integer and float.
func AppendJSONNumber(b []byte, n json.Number) ([]byte, error) {
	if n == "" {
		// The zero value outputs the 0 integer.
		return append(b, 0), nil
	}
	ii, err := n.Int64()
	if err == nil {
		return AppendInt64(b, ii), nil
	}
	ff, err := n.Float64()
	if err == nil {
		return AppendFloat(b, ff), nil
	}
	return b, err
}

// AppendBytesTwoPrefixed will add the length to a bin section written with
// 2 bytes of space saved for a bin8 header.
// If the sz cannot fit inside a bin8, the data will be moved to make space for the header.
func AppendBytesTwoPrefixed(b []byte, sz int) []byte {
	off := len(b) - sz - 2
	switch {
	case sz <= math.MaxUint8:
		// Just write header...
		prefixu8(b[off:], mbin8, uint8(sz))
	case sz <= math.MaxUint16:
		// Scoot one
		b = append(b, 0)
		copy(b[off+1:], b[off:])
		prefixu16(b[off:], mbin16, uint16(sz))
	default:
		// Scoot three
		b = append(b, 0, 0, 0)
		copy(b[off+3:], b[off:])
		prefixu32(b[off:], mbin32, uint32(sz))
	}
	return b
}

// AppendBytesStringTwoPrefixed will add the length to a string section written with
// 2 bytes of space saved for a str8 header.
// If the sz cannot fit inside a str8, the data will be moved to make space for the header.
func AppendBytesStringTwoPrefixed(b []byte, sz int) []byte {
	off := len(b) - sz - 2
	switch {
	case sz <= math.MaxUint8:
		// Just write header...
		prefixu8(b[off:], mstr8, uint8(sz))
	case sz <= math.MaxUint16:
		// Scoot one
		b = append(b, 0)
		copy(b[off+1:], b[off:])
		prefixu16(b[off:], mstr16, uint16(sz))
	default:
		// Scoot three
		b = append(b, 0, 0, 0)
		copy(b[off+3:], b[off:])
		prefixu32(b[off:], mstr32, uint32(sz))
	}
	return b
}