package producer

import (
	"bytes"
	"hash"
	"hash/fnv"
	"strconv"

	"github.com/Shopify/sarama"
	"github.com/juju/errors"
	"github.com/ngaut/log"
)

const (
	trickyPartitionPrefix = "tricky_partition:"
)

type hashPartitioner struct {
	random sarama.Partitioner
	hasher hash.Hash32
}

// NewTrickyHashPartitioner creates a tricky Partitioner for fanout usage.
func NewTrickyHashPartitioner(topic string) sarama.Partitioner {
	p := new(hashPartitioner)
	p.random = sarama.NewRandomPartitioner(topic)
	p.hasher = fnv.New32a()
	return p
}

func (p *hashPartitioner) Partition(message *sarama.ProducerMessage, numPartitions int32) (int32, error) {
	if message.Key == nil || message.Key.Length() == 0 {
		//FIXME: random partition is not allowed which potentially leads to data corruption.
		return p.random.Partition(message, numPartitions)
	}
	key, err := message.Key.Encode()
	if err != nil {
		return -1, errors.Trace(err)
	}

	if bytes.HasPrefix(key, []byte(trickyPartitionPrefix)) {
		spt := bytes.Split(key, []byte{':'})
		if len(spt) == 2 {
			var partition int64
			partition, err = strconv.ParseInt(string(spt[1]), 10, 64)
			if err == nil {
				return int32(partition), nil
			}
			log.Warnf("[producer] parse int failed. source: %s", string(spt[1]))
		}
		log.Warnf("[producer] invalid tricky key: %s", string(key))
	}

	p.hasher.Reset()
	_, err = p.hasher.Write(key)
	if err != nil {
		return -1, errors.Trace(err)
	}
	partition := int32(p.hasher.Sum32()) % numPartitions
	return partition, nil
}

func (p *hashPartitioner) RequiresConsistency() bool {
	return true
}