package consumer

import (
	"sync"

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

// OffsetStash allows to accumulate offsets and
// mark them as processed in a bulk
type OffsetStash struct {
	offsets map[topicPartition]offsetInfo
	mu      sync.Mutex
}

// NewOffsetStash inits a blank stash
func NewOffsetStash() *OffsetStash {
	return &OffsetStash{offsets: make(map[topicPartition]offsetInfo)}
}

// MarkOffset stashes the provided message offset
func (s *OffsetStash) MarkOffset(msg *sarama.ConsumerMessage, metadata string) {
	s.MarkPartitionOffset(msg.Topic, msg.Partition, msg.Offset, metadata)
}

// MarkPartitionOffset stashes the offset for the provided topic/partition combination
func (s *OffsetStash) MarkPartitionOffset(topic string, partition int32, offset int64, metadata string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	key := topicPartition{Topic: topic, Partition: partition}
	if info := s.offsets[key]; offset >= info.Offset {
		info.Offset = offset
		info.Metadata = metadata
		s.offsets[key] = info
	}
}

// Offsets returns the latest stashed offsets by topic-partition
func (s *OffsetStash) Offsets() map[string]int64 {
	s.mu.Lock()
	defer s.mu.Unlock()

	res := make(map[string]int64, len(s.offsets))
	for tp, info := range s.offsets {
		res[tp.String()] = info.Offset
	}
	return res
}

//OffsetCommitRequest contains configuration options for requesting commit offset.
type OffsetCommitRequest struct {
	ConsumerGroup           string
	ConsumerGroupGeneration int32  // nothing to do
	ConsumerID              string // nothing to do
	RetentionTime           int64  // nothing to do

	blocks map[string]map[int32]*offsetCommitRequestBlock
}

//AddBlock to set commit offset of request block
func (r *OffsetCommitRequest) AddBlock(topic string, partitionID int32, offset int64, timestamp int64, metadata string) {
	if r.blocks == nil {
		r.blocks = make(map[string]map[int32]*offsetCommitRequestBlock)
	}

	if r.blocks[topic] == nil {
		r.blocks[topic] = make(map[int32]*offsetCommitRequestBlock)
	}

	r.blocks[topic][partitionID] = &offsetCommitRequestBlock{offset, timestamp, metadata}
}

type offsetCommitRequestBlock struct {
	offset    int64
	timestamp int64
	metadata  string
}

//OffsetCommitResponse is the value for response of commit offset
type OffsetCommitResponse struct {
	Errors map[string]map[int32]error
}

//AddError to set err for response of commit offset
func (r *OffsetCommitResponse) AddError(topic string, partition int32, err error) {
	if r.Errors == nil {
		r.Errors = make(map[string]map[int32]error)
	}
	partitions := r.Errors[topic]
	if partitions == nil {
		partitions = make(map[int32]error)
		r.Errors[topic] = partitions
	}
	partitions[partition] = err
}

//OffsetFetchRequest contains configuration options for request of fetch offset
type OffsetFetchRequest struct {
	ConsumerGroup string
	partitions    map[string][]int32
}

//AddPartition to set partition for request
func (r *OffsetFetchRequest) AddPartition(topic string, partitionID int32) {
	if r.partitions == nil {
		r.partitions = make(map[string][]int32)
	}

	r.partitions[topic] = append(r.partitions[topic], partitionID)
}

//OffsetFetchResponseBlock contains configuration options for response of fetch offset block
type OffsetFetchResponseBlock struct {
	Offset   int64
	Metadata string
	Err      error
}

//OffsetFetchResponse contains configuration options for response of fetch offset
type OffsetFetchResponse struct {
	Blocks map[string]map[int32]*OffsetFetchResponseBlock
}

//GetBlock to get fetch offset of response block
func (r *OffsetFetchResponse) GetBlock(topic string, partition int32) *OffsetFetchResponseBlock {
	if r.Blocks == nil {
		log.Warnf("[consumer] blocks == nil. topic %s, partition %d", topic, partition)
		return nil
	}

	if r.Blocks[topic] == nil {
		log.Warnf("[consumer] block is nil. blocks %+v, topic %s, partition %d", r.Blocks, topic, partition)
		return nil
	}

	return r.Blocks[topic][partition]
}

//AddBlock to set fetch offset of response block
func (r *OffsetFetchResponse) AddBlock(topic string, partitionID int32, offset int64, timestamp int64, metadata string) {
	if r.Blocks == nil {
		r.Blocks = make(map[string]map[int32]*OffsetFetchResponseBlock)
	}

	partitions := r.Blocks[topic]
	if partitions == nil {
		partitions = make(map[int32]*OffsetFetchResponseBlock)
		r.Blocks[topic] = partitions
	}

	r.Blocks[topic][partitionID] = &OffsetFetchResponseBlock{
		Offset:   offset,
		Metadata: metadata,
		Err:      nil,
	}
}

type etcdOffsetManager struct {
	client *etcdutil.Client
}

func newEtcdOffsetManager(client *etcdutil.Client) *etcdOffsetManager {
	return &etcdOffsetManager{client: client}
}

func (m *etcdOffsetManager) CommitOffset(request *OffsetCommitRequest) (*OffsetCommitResponse, error) {
	resp := new(OffsetCommitResponse)
	resp.Errors = make(map[string]map[int32]error)

	for topic, partitions := range request.blocks {
		data := make([]*whpb.Partition, 0, len(partitions))
		for partitionID, offset := range partitions {
			p := &whpb.Partition{
				TopicName:   topic,
				PartitionId: partitionID,
				Pos: &whpb.Position{
					KafkaOffset: offset.offset,
					// TODO: the rest of position information
				},
			}
			data = append(data, p)
		}

		err := m.client.SaveConsumerPartition(request.ConsumerGroup, topic, &whpb.Partitions{Partitions: data})
		if err != nil {
			log.Errorf("[consumer] save partition error %s", errors.ErrorStack(err))
			// we save partitions and offsets in batch, so we don't know which partition got error. just set -1 to identify.
			resp.AddError(topic, -1, err)
		}
	}

	return resp, nil
}

func (m *etcdOffsetManager) FetchOffset(request *OffsetFetchRequest) (*OffsetFetchResponse, error) {
	var partitions []*whpb.Partition
	var err error
	resp := new(OffsetFetchResponse)

	for topic, partitionIds := range request.partitions {
		log.Debugf("[consumer] fetch offset topic %s. partitions %+v", topic, partitionIds)
		partitions, err = m.client.LoadConsumerPartitions(topic, request.ConsumerGroup)
		if err != nil {
			log.Errorf("[consumer] load partitions err:%s, topic: %s, group:%s", err, topic, request.ConsumerGroup)
			return nil, errors.Trace(err)
		}
		partitionsMap := make(map[int32]*whpb.Partition)
		for _, p := range partitions {
			partitionsMap[p.PartitionId] = p
		}

		for _, pid := range partitionIds {
			if p, ok := partitionsMap[pid]; ok {
				log.Debugf("[consumer] fetch offset. get result for topic %s  partition %d offset %d", topic, pid, p.Pos.GetKafkaOffset())
				resp.AddBlock(topic, pid, p.Pos.GetKafkaOffset(), 0, "")
			} else {
				log.Debugf("[consumer] fetch offset. set default offset for topic %s partition %d", topic, pid)
				// set default value as sarama does.
				// &{Offset:-2 Metadata: ""}
				resp.AddBlock(topic, pid, sarama.OffsetOldest, 0, "")
			}
		}
	}
	return resp, nil
}