package chash

import (
	"math/rand"
	"os"
	"runtime"
	"sort"
	"strconv"
	"sync"
	"time"
)

const (
	CHASH_MAGIC    uint32 = 0x48414843
	CHASH_REPLICAS        = 128
)

type item struct {
	hash   uint32
	target uint16
}

type CHash struct {
	targets  map[string]uint8
	names    []string
	ring     []item
	ringSize uint32
	replicas uint8
	frozen   bool
	sync.RWMutex
}

var cores int

func init() {
	rand.Seed(time.Now().UnixNano() + int64(os.Getpid()))
}

func mmhash2(key []byte, keySize int) uint32 {
	var magic, hash, current, value uint32 = 0x5bd1e995, uint32(0x4d4d4832 ^ keySize), 0, 0

	if keySize < 0 {
		keySize = len(key)
	}
	for keySize >= 4 {
		value = uint32(key[current]) | uint32(key[current+1])<<8 |
			uint32(key[current+2])<<16 | uint32(key[current+3])<<24
		value *= magic
		value ^= value >> 24
		value *= magic
		hash *= magic
		hash ^= value
		current += 4
		keySize -= 4
	}
	if keySize >= 3 {
		hash ^= uint32(key[current+2]) << 16
	}
	if keySize >= 2 {
		hash ^= uint32(key[current+1]) << 8
	}
	if keySize >= 1 {
		hash ^= uint32(key[current])
	}
	if keySize != 0 {
		hash *= magic
	}
	hash ^= hash >> 13
	hash *= magic
	hash ^= hash >> 15
	return hash
}

type ByHash []item

func (a ByHash) Len() int           { return len(a) }
func (a ByHash) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByHash) Less(i, j int) bool { return a[i].hash < a[j].hash }

func (this *CHash) freeze() {
	if cores > 1 {
		this.Lock()
		defer this.Unlock()
	}
	if this.frozen {
		return
	}
	this.ringSize = 0
	for _, tweight := range this.targets {
		this.ringSize += uint32(tweight) * uint32(this.replicas)
	}
	if this.ringSize == 0 {
		this.frozen = true
		return
	}

	var (
		target uint16 = 0
		offset uint32 = 0
		key    []byte = make([]byte, 128)
	)
	this.names = make([]string, len(this.targets))
	this.ring = make([]item, this.ringSize)
	for tname, tweight := range this.targets {
		this.names[target] = tname
		for weight := uint8(0); weight < tweight; weight++ {
			for replica := uint8(0); replica < this.replicas; replica++ {
				key = append(key[:0], tname...)
				key = strconv.AppendInt(key, int64(weight), 10)
				key = strconv.AppendInt(key, int64(replica), 10)
				this.ring[offset] = item{mmhash2(key, -1), target}
				offset++
			}
		}
		target++
	}
	sort.Sort(ByHash(this.ring))
	this.frozen = true
}

func New(replicas ...uint8) *CHash {
	if cores == 0 {
		cores = runtime.NumCPU()
	}
	chash := &CHash{
		targets:  make(map[string]uint8),
		names:    nil,
		ring:     nil,
		ringSize: 0,
		replicas: CHASH_REPLICAS,
		frozen:   false,
	}
	if len(replicas) > 0 {
		chash.replicas = replicas[0]
	}
	if chash.replicas < 1 {
		chash.replicas = 1
	}
	if chash.replicas > CHASH_REPLICAS {
		chash.replicas = CHASH_REPLICAS
	}
	return chash
}

func (this *CHash) AddTarget(name string, weight uint8) bool {
	if weight > 0 && weight <= 100 && len(name) <= 128 && this.targets[name] != weight {
		if cores > 1 {
			this.Lock()
			defer this.Unlock()
		}
		this.targets[name] = weight
		this.frozen = false
		return true
	}
	return false
}
func (this *CHash) RemoveTarget(name string) bool {
	if cores > 1 {
		this.Lock()
		defer this.Unlock()
	}
	delete(this.targets, name)
	this.frozen = false
	return true
}
func (this *CHash) ClearTargets() bool {
	if cores > 1 {
		this.Lock()
		defer this.Unlock()
	}
	this.targets = make(map[string]uint8)
	this.frozen = false
	return true
}

func (this *CHash) Serialize() []byte {
	this.freeze()
	if cores > 1 {
		this.RLock()
		defer this.RUnlock()
	}
	size := uint32(4) + 4 + 1 + 2 + 4
	for _, name := range this.names {
		size += 1 + 1 + uint32(len(name))
	}
	size += (this.ringSize * 6)
	serialized := make([]byte, size)
	offset := uint32(0)
	serialized[offset] = byte(CHASH_MAGIC & 0xff)
	serialized[offset+1] = byte((CHASH_MAGIC >> 8) & 0xff)
	serialized[offset+2] = byte((CHASH_MAGIC >> 16) & 0xff)
	serialized[offset+3] = byte((CHASH_MAGIC >> 24) & 0xff)
	serialized[offset+4] = byte(size & 0xff)
	serialized[offset+5] = byte((size >> 8) & 0xff)
	serialized[offset+6] = byte((size >> 16) & 0xff)
	serialized[offset+7] = byte((size >> 24) & 0xff)
	serialized[offset+8] = this.replicas
	serialized[offset+9] = byte(uint16(len(this.names)) & 0xff)
	serialized[offset+10] = byte(((uint16(len(this.names))) >> 8) & 0xff)
	serialized[offset+11] = byte(this.ringSize & 0xff)
	serialized[offset+12] = byte((this.ringSize >> 8) & 0xff)
	serialized[offset+13] = byte((this.ringSize >> 16) & 0xff)
	serialized[offset+14] = byte((this.ringSize >> 24) & 0xff)
	offset += 15
	for _, name := range this.names {
		serialized[offset] = this.targets[name]
		serialized[offset+1] = byte(len(name) & 0xff)
		copy(serialized[offset+2:offset+2+uint32(serialized[offset+1])], []byte(name))
		offset += 2 + uint32(serialized[offset+1])
	}
	for _, item := range this.ring {
		serialized[offset] = byte(item.hash & 0xff)
		serialized[offset+1] = byte((item.hash >> 8) & 0xff)
		serialized[offset+2] = byte((item.hash >> 16) & 0xff)
		serialized[offset+3] = byte((item.hash >> 24) & 0xff)
		serialized[offset+4] = byte(item.target & 0xff)
		serialized[offset+5] = byte((item.target >> 8) & 0xff)
		offset += 6
	}
	return serialized
}
func (this *CHash) FileSerialize(path string) bool {
	handle, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return false
	}
	defer handle.Close()
	if _, err := handle.Write(this.Serialize()); err != nil {
		return false
	}
	return true
}

func (this *CHash) Unserialize(serialized []byte) bool {
	if cores > 1 {
		this.Lock()
		defer this.Unlock()
	}
	if len(serialized) < 15 {
		return false
	}
	magic := uint32(serialized[0]) + (uint32(serialized[1]) << 8) + (uint32(serialized[2]) << 16) + (uint32(serialized[3]) << 24)
	size := uint32(serialized[4]) + (uint32(serialized[5]) << 8) + (uint32(serialized[6]) << 16) + (uint32(serialized[7]) << 24)
	replicas := serialized[8]
	names := uint16(serialized[9]) + (uint16(serialized[10]) << 8)
	ringSize := uint32(serialized[11]) + (uint32(serialized[12]) << 8) + (uint32(serialized[13]) << 16) + (uint32(serialized[14]) << 24)
	if magic != CHASH_MAGIC || size != uint32(len(serialized)) {
		return false
	}
	this.targets = make(map[string]uint8)
	this.names = make([]string, names)
	this.ring = make([]item, ringSize)
	this.ringSize = ringSize
	this.replicas = replicas
	offset := uint32(15)
	for index := uint16(0); index < names && offset < size; index++ {
		len := uint32(serialized[offset+1])
		this.names[index] = string(serialized[offset+2 : offset+2+len])
		this.targets[this.names[index]] = serialized[offset]
		offset += 2 + len
	}
	if offset > size {
		return false
	}
	for item := uint32(0); item < ringSize && offset < size; item++ {
		this.ring[item].hash = uint32(serialized[offset]) + (uint32(serialized[offset+1]) << 8) + (uint32(serialized[offset+2]) << 16) + (uint32(serialized[offset+3]) << 24)
		this.ring[item].target = uint16(serialized[offset+4]) + (uint16(serialized[offset+5]) << 8)
		offset += 6
	}
	if offset != size {
		return false
	}
	this.frozen = true
	return true
}
func (this *CHash) FileUnserialize(path string) bool {
	handle, err := os.OpenFile(path, os.O_RDONLY, 0)
	if err != nil {
		return false
	}
	defer handle.Close()
	info, err := handle.Stat()
	if err != nil {
		return false
	}
	if info.Size() > 128*1024*1024 {
		return false
	}
	serialized := make([]byte, info.Size())
	read, err := handle.Read(serialized)
	if int64(read) != info.Size() || err != nil {
		return false
	}
	return this.Unserialize(serialized)
}

func (this *CHash) Lookup(candidate string, count int) []string {
	var start uint32 = 0

	this.freeze()
	if cores > 1 {
		this.RLock()
		defer this.RUnlock()
	}
	if count > len(this.targets) {
		count = len(this.targets)
	}
	if this.ringSize == 0 || count < 1 {
		return []string{}
	}
	hash := mmhash2([]byte(candidate), -1)
	if hash > this.ring[0].hash && hash <= this.ring[this.ringSize-1].hash {
		start = this.ringSize / 2
		span := start / 2
		for {
			if hash > this.ring[start].hash && hash <= this.ring[start+1].hash {
				break
			}
			if hash > this.ring[start].hash {
				start += span
			} else {
				start -= span
			}
			span /= 2
			if span < 1 {
				span = 1
			}
		}
	}
	result := make([]string, count)
	rank := 0
	for rank < count {
		index := 0
		for index = 0; index < rank; index++ {
			if result[index] == this.names[this.ring[start].target] {
				break
			}
		}
		if index >= rank {
			result[rank] = this.names[this.ring[start].target]
			rank++
		}
		start++
		if start >= this.ringSize {
			start = 0
		}
	}
	return result
}
func (this *CHash) LookupBalance(candidate string, count int) string {
	result := this.Lookup(candidate, count)
	if len(result) > 0 {
		return result[rand.Intn(len(result))]
	}
	return ""
}