package producer

import (
	"fmt"

	"github.com/Shopify/sarama"
	"github.com/gogo/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/wormhole/pkg/etcdutil"
	"github.com/pingcap/wormhole/pkg/whpb"
)

// Producer is a producer to produce job to Message Queue.
type Producer interface {
	// Produce drip to specific topic. It's partition is determined by drip.Key. If drip.Key is not set, then random partition is chose.
	Produce(t *whpb.Topic, drip *whpb.Drip) (pos *whpb.Position, partition int32, err error)

	//FanoutProduce is used to produce a drip to all partitions in specific topic, that's fan-out. However, it's not atomic.
	FanoutProduce(topics []*whpb.Topic, drip *whpb.Drip) (err error)

	// Must call this function to close the Producer, or memory leak.
	Close() error
}

// NewProducer returns a Producer
func NewProducer(addrs []string, config *sarama.Config, client *etcdutil.Client, opt ...PositionManagerOption) (Producer, error) {
	return newKafkaProducer(addrs, config, client, opt...)
}

// KafkaProducer is the kafkaClient wraps the sarama library.
type KafkaProducer struct {
	posManager PositionManager
	producer   sarama.SyncProducer
	Client     sarama.Client
}

func newKafkaProducer(addrs []string, config *sarama.Config, client *etcdutil.Client, opt ...PositionManagerOption) (*KafkaProducer, error) {
	p := &KafkaProducer{}

	// use tricky hash partitioner for fan-out.
	config.Producer.Partitioner = NewTrickyHashPartitioner
	kclient, err := sarama.NewClient(addrs, config)
	if err != nil {
		return nil, errors.Trace(err)
	}
	p.Client = kclient

	producer, err := sarama.NewSyncProducerFromClient(kclient)
	if err != nil {
		return nil, errors.Trace(err)
	}
	p.producer = producer

	p.posManager = NewPositionManager(client, opt...)

	return p, nil
}

// Produce implements the Producer interface.
func (p *KafkaProducer) Produce(t *whpb.Topic, drip *whpb.Drip) (pos *whpb.Position, partition int32, err error) {
	partition, offset, err := p.sendMessage(t, drip, drip.GetKey())
	if err != nil {
		return nil, 0, errors.Trace(err)
	}

	pos = drip.Pos
	pos.KafkaOffset = offset
	log.Debugf("[producer] Produce topic %s, drip %+v, pos %+v, partition %v, offset %+v", t.GetName(), drip, pos, partition, offset)

	metaPartition := &whpb.Partition{
		TopicName:   t.GetName(),
		PartitionId: partition,
		Pos:         pos,
	}
	if err := p.posManager.MarkAsDone(metaPartition); err != nil {
		return nil, 0, errors.Trace(err)
	}

	return pos, partition, nil
}

func (p *KafkaProducer) sendMessage(t *whpb.Topic, drip *whpb.Drip, key string) (partition int32, offset int64, err error) {
	v, err := proto.Marshal(drip)
	if err != nil {
		return 0, 0, errors.Trace(err)
	}

	msg := &sarama.ProducerMessage{Topic: t.GetName(), Key: sarama.StringEncoder(key), Value: sarama.ByteEncoder(v)}
	partition, offset, err = p.producer.SendMessage(msg)
	if err != nil {
		return 0, 0, errors.Trace(err)
	}

	log.Debugf("[producer] produce drip to topic %s partition %d in offset %d", t.GetName(), partition, offset)

	return partition, offset, nil
}

// FanoutProduce implements Producer interface.
func (p *KafkaProducer) FanoutProduce(topics []*whpb.Topic, drip *whpb.Drip) (err error) {
	// only one topic and one partition will have data if got fenceid.
	log.Infof("[producer] FanoutProduce drip %+v, topics %+v", drip, topics)

	var (
		solo       bool
		partitions []int32
	)
	if drip.Data != nil && drip.FenceId != 0 {
		solo = true
	}

	placeHolderDrip := proto.Clone(drip).(*whpb.Drip)
	placeHolderDrip.Data = nil

	msgs := make([]*sarama.ProducerMessage, 0, len(topics)*10) // assumes that's 10 partition per topic. never mind, it's just preallocation.
	for i, t := range topics {
		partitions, err = p.Client.Partitions(t.GetName())
		if err != nil {
			return errors.Trace(err)
		}

		for j, partition := range partitions {
			var v []byte
			if !solo || (solo && i == 0 && j == 0) {
				v, err = proto.Marshal(drip)
			} else {
				v, err = proto.Marshal(placeHolderDrip)
			}
			if err != nil {
				return errors.Trace(err)
			}
			key := fmt.Sprintf("%s%d", trickyPartitionPrefix, partition)
			msg := &sarama.ProducerMessage{Topic: t.GetName(), Key: sarama.StringEncoder(key), Value: sarama.ByteEncoder(v)}
			msgs = append(msgs, msg)
		}
	}

	err = p.producer.SendMessages(msgs)
	if err != nil {
		return errors.Trace(err)
	}
	for _, msg := range msgs {
		pos := drip.Pos
		if pos == nil {
			pos = &whpb.Position{}
		}
		pos.KafkaOffset = msg.Offset
		metaPartition := &whpb.Partition{
			TopicName:   msg.Topic,
			PartitionId: msg.Partition,
			Pos:         pos,
		}
		err = p.posManager.MarkAsDone(metaPartition)
		if err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

// Close implements Producer interface.
func (p *KafkaProducer) Close() error {
	p.posManager.Close()
	p.producer.Close()
	return p.Client.Close()
}