/*
 * Copyright 2021 PingCAP, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"

	"golang.org/x/sync/errgroup"

	"github.com/aws/aws-sdk-go/aws/awserr"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/segmentio/kafka-go"

	"github.com/aws/aws-sdk-go/service/sqs"
)

type IQueueConsumer interface {
	Start(ctx context.Context) error
}

type SQSConsumer struct {
	sqsClient         *sqs.SQS
	queueURL          string
	visibilityTimeout int64
	mapper            IMapper
	skipErr           bool
}

type SQSConsumerOpt func(consumer *SQSConsumer) error

func WithVisibilityTimeout(timeout int64) SQSConsumerOpt {
	return func(consumer *SQSConsumer) error {
		consumer.visibilityTimeout = timeout
		return nil
	}
}

func NewSQSConsumer(cli *sqs.SQS, qURL string, opts ...SQSConsumerOpt) (*SQSConsumer, error) {
	c := &SQSConsumer{
		sqsClient:         cli,
		queueURL:          qURL,
		visibilityTimeout: 1,
	}
	for _, opt := range opts {
		if err := opt(c); err != nil {
			return nil, err
		}
	}
	return c, nil
}

func (c *SQSConsumer) Start(ctx context.Context) error {
	// todo
	if c.mapper == nil {
		return errors.New("mapper is nil")
	}
	log.Println("sqs consumer start")
	defer c.close()
	for {
		select {
		case <-ctx.Done():
			if !errors.Is(ctx.Err(), context.Canceled) {
				return ctx.Err()
			}
			return nil
		default:
			output, err := c.sqsClient.ReceiveMessage(&sqs.ReceiveMessageInput{
				AttributeNames: []*string{
					aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
				},
				MessageAttributeNames: []*string{
					aws.String(sqs.QueueAttributeNameAll),
				},
				QueueUrl:            &c.queueURL,
				MaxNumberOfMessages: aws.Int64(10),
				VisibilityTimeout:   &c.visibilityTimeout,
			})
			// for sqs,
			if err != nil {
				log.Printf("receive message err: %+v\n", err)
				if awsErr, ok := err.(awserr.Error); ok {
					if awsErr.Code() == sqs.ErrCodeOverLimit {
						continue
					}
				}
				return err
			}
			for _, msg := range output.Messages {
				data := make(map[string]interface{})
				b := []byte(*msg.Body)
				if err := json.Unmarshal(b, &data); err != nil {
					log.Printf("unmarshal %+v err: %+v\n", msg.String(), err)
					continue
				}
				if err := c.mapper.Handle(context.Background(), data); err != nil {
					log.Printf("mapper handle %+v err: %+v\n", msg.String(), err)
					if c.skipErr {
						continue
					}
					return err
				}
			}
		}
	}
}

func (c *SQSConsumer) SetMapper(m IMapper) {
	c.mapper = m
}

func (c *SQSConsumer) SkipError(skip bool) {
	c.skipErr = skip
}

func (c *SQSConsumer) close() error {
	return nil
}

// todo
type KafkaConsumer struct {
	readers []*kafka.Reader
	mapper  IMapper
	skipErr bool
}

// NewKafkaConsumer create a kafka consumer. let the caller decide use which dialer
func NewKafkaConsumer(dialer *kafka.Dialer, brokers []string, topic string, groupID string) (*KafkaConsumer, error) {
	var parts []kafka.Partition
	var err error
	for _, v := range brokers{
		parts, err = dialer.LookupPartitions(context.Background(), "tcp", v, topic)
		if err == nil {
			break
		}
	}
	if len(parts) == 0 {
		return nil, fmt.Errorf("retrieve empty partition list: err %+v\n", err)
	}

	c := &KafkaConsumer{
		readers: make([]*kafka.Reader, len(parts)),
		skipErr: false,
	}
	readerCfg := kafka.ReaderConfig{
		Brokers: brokers,
		Topic:   topic,
		GroupID: groupID,
		Dialer:  dialer,
	}
	for i := range parts {
		c.readers[i] = kafka.NewReader(readerCfg)
	}
	return c, nil
}

func (c *KafkaConsumer) Start(ctx context.Context) error {
	if c.mapper == nil {
		return errors.New("mapper is nil")
	}
	defer c.close()
	errG, ctx := errgroup.WithContext(ctx)
	for _, r := range c.readers {
		errG.Go(func() error {
			return c.startReader(ctx, r)
		})
	}
	return errG.Wait()
}

func (c *KafkaConsumer) startReader(ctx context.Context, reader *kafka.Reader) error {
	log.Println("start reader")
	for {
		select {
		case <-ctx.Done():
			if !errors.Is(ctx.Err(), context.Canceled) {
				return ctx.Err()
			}
			return nil
		default:
			m, err := reader.FetchMessage(ctx)
			// reader never return io.EOF
			if err != nil {
				if errors.Is(err, context.Canceled) {
					return nil
				}
				log.Printf("read message err: %+v\n", err)
				return err
			}
			data := make(map[string]interface{})
			// if the data is not validate json, jut log error and skip it.
			if err := json.Unmarshal(m.Value, &data); err != nil {
				log.Printf("unmarshal %+v err: %+v\n", string(m.Value), err)
				continue
			}
			// by default cannot skip
			if err := c.mapper.Handle(ctx, data); err != nil {
				log.Printf("mapper handle data: %+v err: %+v\n", data, err)
				if c.skipErr {
					continue
				}
				return err
			}
			// by default cannot skip
			if err := reader.CommitMessages(context.Background(), m); err != nil {
				log.Printf("commit message on topic %+v, partiton %+v, offset %+v error: %+v\n", m.Topic, m.Partition, m.Offset, err)
				if c.skipErr {
					continue
				}
				return err
			}
		}
	}
}

func (c *KafkaConsumer) SetMapper(m IMapper) {
	c.mapper = m
}

func (c *KafkaConsumer) SkipError(skip bool) {
	c.skipErr = skip
}

// must close the reader before exist
func (c *KafkaConsumer) close() error {
	var err error
	for _, r := range c.readers {
		log.Println("close reader")
		if closeErr := r.Close(); closeErr != nil {
			err = closeErr
		}
	}
	return err
}