package memcache

import (
	"crypto/md5"
	"fmt"
	"math"
	"sort"
)

type HashKeyOrder []uint32

func (h HashKeyOrder) Len() int           { return len(h) }
func (h HashKeyOrder) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h HashKeyOrder) Less(i, j int) bool { return h[i] < h[j] }

// HashRing is based on https://github.com/serialx/hashring
// It's been simplified by removing weighting (just repeat a node instead)
// and lacks any Add/Remove/Update methods (make a new hashring and atomically swap it)
// It adds the ability to control the number of virtual nodes per server (use to be hardcoded to 120)
// It also implements multi-probing, but only with a K of 4, since MD5 gives us 4 keys per hash anyway.
// For more about Multi-probe consistent hashing, see http://arxiv.org/pdf/1505.00062.pdf
type HashRing struct {
	ring       map[uint32]string
	sortedKeys []uint32
}

// NewHashRing takes a list of servers and a duplication value (nodesPerServer) and constructs a
// consistent hashing ring out of them. See https://en.wikipedia.org/wiki/Consistent_hashing#Technique
func NewHashRing(servers []string, nodesPerServer int) *HashRing {
	h := &HashRing{
		ring:       map[uint32]string{},
		sortedKeys: []uint32{},
	}

	weights := map[string]int{}
	for _, server := range servers {
		weights[server]++
	}

	// This will be constant for a compiled binary, but can vary if GenKey's hash algorithm is changed
	genKeyLen := len(h.GenKey(""))

	// For each server, add it to the ring nodesPerServer times, using each key provided by GenKey for performance
	for server, weight := range weights {
		for j := 0; j <= (weight*nodesPerServer)/genKeyLen; j++ {
			nodeKey := fmt.Sprintf("%s-%d", server, j)
			for i, key := range h.GenKey(nodeKey) {
				if j*genKeyLen+i < (weight * nodesPerServer) {
					h.ring[key] = server
					h.sortedKeys = append(h.sortedKeys, key)
				}
			}
		}
	}

	// Having a sorted list of keys allows us to use binary searching in GetNodePos
	sort.Sort(HashKeyOrder(h.sortedKeys))

	return h
}

func (h *HashRing) GetNode(stringKey string) (node string, ok bool) {
	pos, ok := h.GetNodePos(stringKey)
	if !ok {
		return "", false
	}
	return h.ring[h.sortedKeys[pos]], true
}

func (h *HashRing) GetNodePos(stringKey string) (pos int, ok bool) {
	if len(h.ring) == 0 {
		return 0, false
	}

	nodes := h.sortedKeys
	index := 0
	distance := uint32(math.MaxUint32)

	// For each hash, find the nearest node (either forwards or backwards)
	// then see if it's closer than the closest found node
	// By using multiple keys we get better results and can reduce the memory consumed by the ring
	for _, key := range h.GenKey(stringKey) {
		p := sort.Search(len(nodes), func(i int) bool { return nodes[i] > key })

		// Find distance to next node
		if p != len(nodes) {
			d := nodes[p] - key
			if d < distance {
				index, distance = p, d
			}
		} else {
			d := key - nodes[0]
			if d < distance {
				index, distance = p, d
			}
		}

		// Find distance to previous node
		if p != 0 {
			d := key - nodes[p-1]
			if d < distance {
				index, distance = p-1, d
			}
		} else {
			d := nodes[len(nodes)-1] - key
			if d < distance {
				index, distance = p, d
			}
		}
	}

	return index, true
}

// GenKey returns a number of hashes from a given key
// It must return at least one hash
// Returning more than one allows for performance improvements
// by populating the ring faster, and reducing the needed nodesPerServer
// via multi-probing.
func (_ *HashRing) GenKey(key string) []uint32 {
	h := md5.Sum([]byte(key))
	e := func(h []byte) uint32 {
		return ((uint32(h[3]) << 24) |
			(uint32(h[2]) << 16) |
			(uint32(h[1]) << 8) |
			(uint32(h[0])))
	}
	return []uint32{uint32(e(h[0:4])), uint32(e(h[4:8])), uint32(e(h[8:12])), uint32(e(h[12:16]))}
}
