package consumer

import (
	"sort"
	"strconv"
	"sync"
	"time"

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

type partitionConsumer struct {
	pcm   sarama.PartitionConsumer
	state partitionState
	mu    sync.Mutex

	closed      bool
	dying, dead chan none

	ec        ether.Client
	topic     string
	partition int32
	group     string
}

func newPartitionConsumer(manager sarama.Consumer, topic string, partition int32, info offsetInfo, defaultOffset int64, ec ether.Client, group string) (*partitionConsumer, error) {
	pcm, err := manager.ConsumePartition(topic, partition, info.NextOffset(defaultOffset))

	// Resume from default offset, if requested offset is out-of-range
	if err == sarama.ErrOffsetOutOfRange {
		info.Offset = -1
		pcm, err = manager.ConsumePartition(topic, partition, defaultOffset)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}

	return &partitionConsumer{
		pcm:   pcm,
		state: partitionState{Info: info},

		dying:     make(chan none),
		dead:      make(chan none),
		ec:        ec,
		topic:     topic,
		partition: partition,
		group:     group,
	}, nil
}

func (c *partitionConsumer) Loop(messages chan<- Message, errCh chan<- error) {
	defer close(c.dead)

	for {
		select {
		case msg, ok := <-c.pcm.Messages():
			if !ok {
				return
			}
			// hook
			err := c.hijack(messages, msg)
			if err != nil {
				log.Errorf("[consumer] hijack failed: %v", errors.ErrorStack(err))
				close(c.dying)
				return
			}
			select {
			case <-c.dying:
				return
			default:
			}
		case err, ok := <-c.pcm.Errors():
			if !ok {
				return
			}
			select {
			case errCh <- err:
			case <-c.dying:
				return
			}
		case <-c.dying:
			return
		}
	}
}

func (c *partitionConsumer) hijack(messages chan<- Message, msg *sarama.ConsumerMessage) error {
	drip := &whpb.Drip{}
	err := proto.Unmarshal(msg.Value, drip)
	if err != nil {
		return errors.Trace(err)
	}

	if drip.GetFenceId() != 0 {
		log.Debugf("[consumer] group %s, topic %s meet fence %d at partition %d offset %d", c.group, c.topic, drip.GetFenceId(), c.partition, msg.Offset)
		jobID, err := strconv.ParseUint(c.group, 10, 64)
		if err != nil {
			return errors.Trace(err)
		}
		if drip.GetData() != nil {
			err = c.waitOnFenceMidway(drip.GetFenceId(), jobID, &whpb.Topic{Name: msg.Topic}, []int32{c.partition})
			if err != nil {
				return errors.Trace(err)
			}
			syncMsg := Message{
				Message: msg,
				Sync:    true,
				Done:    make(chan struct{}),
			}
			messages <- syncMsg
			<-syncMsg.Done
		} else {
			// message that contains no data should mark as consumed, not just discard it.
			c.MarkOffset(msg.Offset+1, "empty fence")
		}
		err = c.waitOnFenceFinal(drip.GetFenceId(), jobID, &whpb.Topic{Name: msg.Topic}, []int32{c.partition})
		if err != nil {
			return errors.Trace(err)
		}
	} else {
		if drip.GetData() != nil {
			messages <- Message{
				Message: msg,
			}
		}
	}
	return nil
}

func (c *partitionConsumer) waitOnFenceMidway(fenceID uint64, jobID uint64, topic *whpb.Topic, pids []int32) error {
	fenceClient, err := c.ec.WaitOnFenceMidway(fenceID, jobID, topic, pids)
	if err != nil {
		return errors.Trace(err)
	}
	log.Debugf("[consumer] group %s topic %s partition %d is going to wait midway fence %d", c.group, c.topic, c.partition, fenceID)
	for {
		resp, err := fenceClient.Recv()
		if err != nil {
			return errors.Trace(err)
		}
		if err := resp.GetHeader().GetError(); err != nil {
			return errors.Errorf("code %v message %v", err.GetErrcode(), err.GetMessage())
		}
		if resp.GetIsSynced() {
			log.Infof("[consumer] group %s topic %s partition %d, fence %d is synced midway", c.group, c.topic, c.partition, fenceID)
			break
		}
	}
	return nil
}

func (c *partitionConsumer) waitOnFenceFinal(fenceID uint64, jobID uint64, topic *whpb.Topic, pids []int32) error {
	fenceClient, err := c.ec.WaitOnFenceFinal(fenceID, jobID, topic, pids)
	if err != nil {
		return errors.Trace(err)
	}
	log.Debugf("[consumer] group %s topic %s partition %d is going to wait final fence %d", c.group, c.topic, c.partition, fenceID)
	for {
		resp, err := fenceClient.Recv()
		if err != nil {
			return errors.Trace(err)
		}
		if err := resp.GetHeader().GetError(); err != nil {
			return errors.Errorf("code %v message %v", err.GetErrcode(), err.GetMessage())
		}
		if resp.GetIsSynced() {
			log.Infof("[consumer] group %s topic %s partition %d, fence %d is synced final", c.group, c.topic, c.partition, fenceID)
			break
		}
	}
	return nil
}

func (c *partitionConsumer) Close() error {
	if c.closed {
		return nil
	}

	err := c.pcm.Close()
	c.closed = true
	close(c.dying)
	<-c.dead

	return errors.Trace(err)
}

func (c *partitionConsumer) State() partitionState {
	if c == nil {
		return partitionState{}
	}

	c.mu.Lock()
	state := c.state
	c.mu.Unlock()

	return state
}

func (c *partitionConsumer) MarkCommitted(offset int64) {
	if c == nil {
		return
	}

	c.mu.Lock()
	if offset == c.state.Info.Offset {
		c.state.Dirty = false
	}
	c.mu.Unlock()
}

func (c *partitionConsumer) MarkOffset(offset int64, metadata string) {
	if c == nil {
		return
	}

	log.Debugf("mark offset group %v, topic %v, partition %v, offset %d, metadata %s, state.info.offset %d", c.group, c.topic, c.partition, offset, metadata, c.state.Info.Offset)

	c.mu.Lock()
	if offset > c.state.Info.Offset {
		c.state.Info.Offset = offset
		c.state.Info.Metadata = metadata
		c.state.Dirty = true
		c.state.LastCommit = time.Now()
	}
	c.mu.Unlock()
}

// --------------------------------------------------------------------

type partitionState struct {
	Info       offsetInfo
	Dirty      bool
	LastCommit time.Time
}

// --------------------------------------------------------------------

type partitionMap struct {
	data map[topicPartition]*partitionConsumer
	mu   sync.RWMutex
}

func newPartitionMap() *partitionMap {
	return &partitionMap{
		data: make(map[topicPartition]*partitionConsumer),
	}
}

func (m *partitionMap) IsSubscribedTo(topic string) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for tp := range m.data {
		if tp.Topic == topic {
			return true
		}
	}
	return false
}

func (m *partitionMap) Fetch(topic string, partition int32) *partitionConsumer {
	m.mu.RLock()
	pc, _ := m.data[topicPartition{topic, partition}]
	m.mu.RUnlock()
	return pc
}

func (m *partitionMap) Store(topic string, partition int32, pc *partitionConsumer) {
	m.mu.Lock()
	m.data[topicPartition{topic, partition}] = pc
	m.mu.Unlock()
}

func (m *partitionMap) Snapshot() map[topicPartition]partitionState {
	m.mu.RLock()
	defer m.mu.RUnlock()

	snap := make(map[topicPartition]partitionState, len(m.data))
	for tp, pc := range m.data {
		snap[tp] = pc.State()
	}
	return snap
}

func (m *partitionMap) Stop() {
	m.mu.RLock()
	defer m.mu.RUnlock()

	var wg sync.WaitGroup
	for tp := range m.data {
		wg.Add(1)
		go func(p *partitionConsumer) {
			_ = p.Close()
			wg.Done()
		}(m.data[tp])
	}
	wg.Wait()
}

func (m *partitionMap) Clear() {
	m.mu.Lock()
	for tp := range m.data {
		delete(m.data, tp)
	}
	m.mu.Unlock()
}

func (m *partitionMap) Info() map[string][]int32 {
	info := make(map[string][]int32)
	m.mu.RLock()
	for tp := range m.data {
		info[tp.Topic] = append(info[tp.Topic], tp.Partition)
	}
	m.mu.RUnlock()

	for topic := range info {
		sort.Sort(int32Slice(info[topic]))
	}
	return info
}