package decoder

import (
	"encoding/binary"
	"fmt"
	"math"

	"github.com/oschwald/maxminddb-golang/v2/internal/mmdberrors"
)

// Kind constants for the different MMDB data kinds.
type Kind int

// MMDB data kind constants.
const (
	// KindExtended indicates an extended kind.
	KindExtended Kind = iota
	// KindPointer is a pointer to another location in the data section.
	KindPointer
	// KindString is a UTF-8 string.
	KindString
	// KindFloat64 is a 64-bit floating point number.
	KindFloat64
	// KindBytes is a byte slice.
	KindBytes
	// KindUint16 is a 16-bit unsigned integer.
	KindUint16
	// KindUint32 is a 32-bit unsigned integer.
	KindUint32
	// KindMap is a map from strings to other data types.
	KindMap
	// KindInt32 is a 32-bit signed integer.
	KindInt32
	// KindUint64 is a 64-bit unsigned integer.
	KindUint64
	// KindUint128 is a 128-bit unsigned integer.
	KindUint128
	// KindSlice is an array of values.
	KindSlice
	// KindContainer is a data cache container.
	KindContainer
	// KindEndMarker marks the end of the data section.
	KindEndMarker
	// KindBool is a boolean value.
	KindBool
	// KindFloat32 is a 32-bit floating point number.
	KindFloat32
)

// String returns a human-readable name for the Kind.
func (k Kind) String() string {
	switch k {
	case KindExtended:
		return "Extended"
	case KindPointer:
		return "Pointer"
	case KindString:
		return "String"
	case KindFloat64:
		return "Float64"
	case KindBytes:
		return "Bytes"
	case KindUint16:
		return "Uint16"
	case KindUint32:
		return "Uint32"
	case KindMap:
		return "Map"
	case KindInt32:
		return "Int32"
	case KindUint64:
		return "Uint64"
	case KindUint128:
		return "Uint128"
	case KindSlice:
		return "Slice"
	case KindContainer:
		return "Container"
	case KindEndMarker:
		return "EndMarker"
	case KindBool:
		return "Bool"
	case KindFloat32:
		return "Float32"
	default:
		return fmt.Sprintf("Unknown(%d)", int(k))
	}
}

// IsContainer returns true if the Kind represents a container type (Map or Slice).
func (k Kind) IsContainer() bool {
	return k == KindMap || k == KindSlice
}

// IsScalar returns true if the Kind represents a scalar value type.
func (k Kind) IsScalar() bool {
	switch k {
	case KindString, KindFloat64, KindBytes, KindUint16, KindUint32,
		KindInt32, KindUint64, KindUint128, KindBool, KindFloat32:
		return true
	default:
		return false
	}
}

// DataDecoder is a decoder for the MMDB data section.
// This is exported so mmdbdata package can use it, but still internal.
type DataDecoder struct {
	stringCache *stringCache
	buffer      []byte
}

const (
	// This is the value used in libmaxminddb.
	maximumDataStructureDepth = 512
	pointerBase2              = 2048
	pointerBase3              = 526336
)

// NewDataDecoder creates a [DataDecoder].
func NewDataDecoder(buffer []byte) DataDecoder {
	return DataDecoder{
		buffer:      buffer,
		stringCache: newStringCache(),
	}
}

// getBuffer returns the underlying buffer for direct access.
func (d *DataDecoder) getBuffer() []byte {
	return d.buffer
}

// decodeCtrlData decodes the control byte and data info at the given offset.
// Encoding follows the MaxMind DB spec: the control byte's high 3 bits
// encode the type (or KindExtended if zero, in which case the next byte
// holds kind-7 and the decoder adds 7 back), and the low 5 bits encode
// size. Sizes 0..28 are encoded directly; 29 reads 1 extra byte (+29),
// 30 reads 2 (+285), and 31 reads 3 (+65821).
func (d *DataDecoder) decodeCtrlData(offset uint) (Kind, uint, uint, error) {
	bufferLen := uint(len(d.buffer))
	newOffset := offset + 1
	if offset >= bufferLen {
		return 0, 0, 0, mmdberrors.NewOffsetError()
	}
	ctrlByte := d.buffer[offset]

	kindNum := Kind(ctrlByte >> 5)
	if kindNum == KindExtended {
		if newOffset >= bufferLen {
			return 0, 0, 0, mmdberrors.NewOffsetError()
		}
		kindNum = Kind(d.buffer[newOffset] + 7)
		newOffset++
	}

	size := uint(ctrlByte & 0x1f)
	if size < 29 {
		return kindNum, size, newOffset, nil
	}

	endOffset := newOffset + size - 28
	if endOffset > bufferLen {
		return 0, 0, 0, mmdberrors.NewOffsetError()
	}

	switch size {
	case 29:
		return kindNum, 29 + uint(d.buffer[newOffset]), newOffset + 1, nil
	case 30:
		value := uint(d.buffer[newOffset])<<8 | uint(d.buffer[newOffset+1])
		return kindNum, 285 + value, endOffset, nil
	default: // size == 31
		value := uint(d.buffer[newOffset])<<16 |
			uint(d.buffer[newOffset+1])<<8 |
			uint(d.buffer[newOffset+2])
		return kindNum, 65821 + value, endOffset, nil
	}
}

// decodeBytes decodes a byte slice from the given offset with the given size.
func (d *DataDecoder) decodeBytes(size, offset uint) ([]byte, uint, error) {
	if offset+size > uint(len(d.buffer)) {
		return nil, 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	bytes := make([]byte, size)
	copy(bytes, d.buffer[offset:newOffset])
	return bytes, newOffset, nil
}

// DecodeFloat64 decodes a 64-bit float from the given offset.
func (d *DataDecoder) decodeFloat64(size, offset uint) (float64, uint, error) {
	if size != 8 {
		return 0, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (float 64 size of %v)",
			size,
		)
	}
	if offset+size > uint(len(d.buffer)) {
		return 0, 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	bits := binary.BigEndian.Uint64(d.buffer[offset:newOffset])
	return math.Float64frombits(bits), newOffset, nil
}

// DecodeFloat32 decodes a 32-bit float from the given offset.
func (d *DataDecoder) decodeFloat32(size, offset uint) (float32, uint, error) {
	if size != 4 {
		return 0, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (float32 size of %v)",
			size,
		)
	}
	if offset+size > uint(len(d.buffer)) {
		return 0, 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	bits := binary.BigEndian.Uint32(d.buffer[offset:newOffset])
	return math.Float32frombits(bits), newOffset, nil
}

// DecodeInt32 decodes a 32-bit signed integer from the given offset.
func (d *DataDecoder) decodeInt32(size, offset uint) (int32, uint, error) {
	if size > 4 {
		return 0, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (int32 size of %v)",
			size,
		)
	}
	if offset+size > uint(len(d.buffer)) {
		return 0, 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	var val int32
	for _, b := range d.buffer[offset:newOffset] {
		val = (val << 8) | int32(b)
	}
	return val, newOffset, nil
}

// DecodePointer decodes a pointer from the given offset.
func (d *DataDecoder) decodePointer(
	size uint,
	offset uint,
) (uint, uint, error) {
	pointerSize := ((size >> 3) & 0x3) + 1
	newOffset := offset + pointerSize
	if newOffset > uint(len(d.buffer)) {
		return 0, 0, mmdberrors.NewOffsetError()
	}
	pointerBytes := d.buffer[offset:newOffset]
	var prefix uint
	if pointerSize == 4 {
		prefix = 0
	} else {
		prefix = size & 0x7
	}
	unpacked := uintFromBytes(prefix, pointerBytes)

	var pointerValueOffset uint
	switch pointerSize {
	case 1, 4:
		pointerValueOffset = 0
	case 2:
		pointerValueOffset = pointerBase2
	case 3:
		pointerValueOffset = pointerBase3
	default:
		return 0, 0, mmdberrors.NewInvalidDatabaseError("invalid pointer size: %d", pointerSize)
	}

	pointer := unpacked + pointerValueOffset

	return pointer, newOffset, nil
}

// DecodeBool decodes a boolean from the given offset.
func (*DataDecoder) decodeBool(size, offset uint) (bool, uint, error) {
	if size > 1 {
		return false, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (bool size of %v)",
			size,
		)
	}
	value, newOffset := decodeBool(size, offset)
	return value, newOffset, nil
}

// DecodeString decodes a string from the given offset.
func (d *DataDecoder) decodeString(size, offset uint) (string, uint, error) {
	if offset+size > uint(len(d.buffer)) {
		return "", 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	value := d.stringCache.internAt(offset, size, d.buffer)
	return value, newOffset, nil
}

// DecodeUint16 decodes a 16-bit unsigned integer from the given offset.
func (d *DataDecoder) decodeUint16(size, offset uint) (uint16, uint, error) {
	if size > 2 {
		return 0, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (uint16 size of %v)",
			size,
		)
	}
	if offset+size > uint(len(d.buffer)) {
		return 0, 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	bytes := d.buffer[offset:newOffset]

	var val uint16
	for _, b := range bytes {
		val = (val << 8) | uint16(b)
	}
	return val, newOffset, nil
}

// DecodeUint32 decodes a 32-bit unsigned integer from the given offset.
func (d *DataDecoder) decodeUint32(size, offset uint) (uint32, uint, error) {
	if size > 4 {
		return 0, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (uint32 size of %v)",
			size,
		)
	}
	if offset+size > uint(len(d.buffer)) {
		return 0, 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	bytes := d.buffer[offset:newOffset]

	var val uint32
	for _, b := range bytes {
		val = (val << 8) | uint32(b)
	}
	return val, newOffset, nil
}

// DecodeUint64 decodes a 64-bit unsigned integer from the given offset.
func (d *DataDecoder) decodeUint64(size, offset uint) (uint64, uint, error) {
	if size > 8 {
		return 0, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (uint64 size of %v)",
			size,
		)
	}
	if offset+size > uint(len(d.buffer)) {
		return 0, 0, mmdberrors.NewOffsetError()
	}

	newOffset := offset + size
	bytes := d.buffer[offset:newOffset]

	var val uint64
	for _, b := range bytes {
		val = (val << 8) | uint64(b)
	}
	return val, newOffset, nil
}

// DecodeUint128 decodes a 128-bit unsigned integer from the given offset.
// Returns the value as high and low 64-bit unsigned integers.
func (d *DataDecoder) decodeUint128(size, offset uint) (hi, lo uint64, newOffset uint, err error) {
	if size > 16 {
		return 0, 0, 0, mmdberrors.NewInvalidDatabaseError(
			"the MaxMind DB file's data section contains bad data (uint128 size of %v)",
			size,
		)
	}
	if offset+size > uint(len(d.buffer)) {
		return 0, 0, 0, mmdberrors.NewOffsetError()
	}

	newOffset = offset + size

	// Process bytes from most significant to least significant
	for _, b := range d.buffer[offset:newOffset] {
		var carry byte
		lo, carry = append64(lo, b)
		hi, _ = append64(hi, carry)
	}

	return hi, lo, newOffset, nil
}

func append64(val uint64, b byte) (uint64, byte) {
	return (val << 8) | uint64(b), byte(val >> 56)
}

// DecodeKey decodes a map key into []byte slice. We use a []byte so that we
// can take advantage of https://github.com/golang/go/issues/3512 to avoid
// copying the bytes when decoding a struct. Previously, we achieved this by
// using unsafe.
func (d *DataDecoder) decodeKey(offset uint) ([]byte, uint, error) {
	bufferLen := uint(len(d.buffer))
	if offset >= bufferLen {
		return nil, 0, mmdberrors.NewOffsetError()
	}

	ctrlByte := d.buffer[offset]
	if Kind(ctrlByte>>5) == KindString {
		size := uint(ctrlByte & 0x1f)
		if size < 29 {
			dataOffset := offset + 1
			newOffset := dataOffset + size
			if newOffset > bufferLen {
				return nil, 0, mmdberrors.NewOffsetError()
			}
			return d.buffer[dataOffset:newOffset], newOffset, nil
		}
	}

	kindNum, size, dataOffset, err := d.decodeCtrlData(offset)
	if err != nil {
		return nil, 0, err
	}

	// Follow pointer if present (but only once, per spec)
	nextOffset := dataOffset + size // default return offset
	if kindNum == KindPointer {
		pointer, newNextOffset, err := d.decodePointer(size, dataOffset)
		if err != nil {
			return nil, 0, err
		}
		nextOffset = newNextOffset

		// Decode the pointed-to data
		kindNum, size, dataOffset, err = d.decodeCtrlData(pointer)
		if err != nil {
			return nil, 0, err
		}

		// Check for pointer-to-pointer, which is invalid per spec
		if kindNum == KindPointer {
			return nil, 0, mmdberrors.NewInvalidDatabaseError(
				"invalid pointer to pointer at offset %d",
				pointer,
			)
		}
	}

	if kindNum != KindString {
		return nil, 0, mmdberrors.NewInvalidDatabaseError(
			"unexpected type when decoding string: %v",
			kindNum,
		)
	}
	newOffset := dataOffset + size
	if newOffset > bufferLen {
		return nil, 0, mmdberrors.NewOffsetError()
	}
	return d.buffer[dataOffset:newOffset], nextOffset, nil
}

// NextValueOffset skips ahead to the next value without decoding
// the one at the offset passed in. The size bits have different meanings for
// different data types.
func (d *DataDecoder) nextValueOffset(offset, numberToSkip uint) (uint, error) {
	for numberToSkip > 0 {
		kindNum, size, newOffset, err := d.decodeCtrlData(offset)
		if err != nil {
			return 0, err
		}

		switch kindNum {
		case KindPointer:
			// A pointer value is represented by its pointer token only.
			// To skip it, just move past the pointer bytes; do NOT follow
			// the pointer target here.
			_, ptrEndOffset, err2 := d.decodePointer(size, newOffset)
			if err2 != nil {
				return 0, err2
			}
			newOffset = ptrEndOffset
		case KindMap:
			numberToSkip += 2 * size
		case KindSlice:
			numberToSkip += size
		case KindBool:
			// size encodes the boolean; nothing else to skip
		default:
			newOffset += size
		}

		offset = newOffset
		numberToSkip--
	}
	return offset, nil
}

func decodeBool(size, offset uint) (bool, uint) {
	return size != 0, offset
}

func uintFromBytes(prefix uint, uintBytes []byte) uint {
	val := prefix
	for _, b := range uintBytes {
		val = (val << 8) | uint(b)
	}
	return val
}
