package kafka

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"git.hlsq.asia/alex/service-common/log"
	"git.hlsq.asia/alex/service-common/utils"
	"github.com/IBM/sarama"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	otelcodes "go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	"reflect"
	"sync"
	"time"
)

func NewConsumer() *Consumer {
	ctx, cancel := context.WithCancel(context.Background())
	return &Consumer{
		ctx:    ctx,
		cancel: cancel,
	}
}

type Consumer struct {
	ctx    context.Context
	cancel context.CancelFunc
}

func (c *Consumer) Consume(t []Topic) {
	for {
		topicArr := make([]string, 0)
		handlerMap := make(map[string]*sync.Pool)
		for _, v := range t {
			topicArr = append(topicArr, v.Name())
			typ := reflect.ValueOf(v).Elem().Type()
			handlerMap[v.Name()] = &sync.Pool{New: func() interface{} {
				return reflect.New(typ).Interface()
			}}
		}
		err := client.consumer.Consume(c.ctx, topicArr, &handler{
			handler: handlerMap,
		})
		if errors.Is(err, context.Canceled) || errors.Is(c.ctx.Err(), context.Canceled) {
			return
		}
		time.Sleep(time.Second)
	}
}

func (c *Consumer) Stop() error {
	c.cancel()
	return client.consumer.Close()
}

type Handler func(context.Context, []byte) error

type handler struct {
	handler map[string]*sync.Pool
}

func (h *handler) Setup(_ sarama.ConsumerGroupSession) error {
	return nil
}

func (h *handler) Cleanup(sess sarama.ConsumerGroupSession) error {
	sess.Commit()
	return nil
}

func (h *handler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	for {
		select {
		case message, ok := <-claim.Messages():
			if ok {
				ctx := NewCarrier().ExtractConsumer(message.Headers)
				_, span := otel.Tracer("common.db.kafka").Start(ctx, "kafka.consume")

				pool := h.handler[message.Topic]
				if pool == nil {
					span.SetStatus(otelcodes.Error, "handler not found")
					span.End()
					return utils.ErrorsWrap(errors.New("handler not found"))
				}
				span.SetAttributes(attribute.String("kafka.topic", message.Topic))
				span.SetAttributes(attribute.String("kafka.message", string(message.Value)))

				newCb := pool.Get().(Topic)
				if err := json.Unmarshal(message.Value, newCb); err != nil {
					releaseTopicObject(pool, newCb)
					span.SetStatus(otelcodes.Error, "handler json.Unmarshal error")
					span.End()
					return utils.ErrorsWrap(err, "handler json.Unmarshal error")
				}

				if err := func() (err error) {
					defer func() {
						if e := recover(); e != nil {
							span.SetStatus(otelcodes.Error, fmt.Sprintf("%v", e))
							span.End()
							err = fmt.Errorf("kafka handler panic: %v", e)
						}
					}()
					if err = newCb.OnMessage(ctx); err != nil {
						if stack, ok := err.(interface{ StackTrace() string }); ok {
							span.AddEvent("Stack Trace", trace.WithAttributes(
								attribute.String("stack.trace", fmt.Sprintf("%v", stack.StackTrace())),
							))
						}
						span.SetStatus(otelcodes.Error, err.Error())
						span.End()
						log.Errorf("kafka handler error: %v", err)
						err = utils.ErrorsWrap(err, "kafka handler error")
					}
					return
				}(); err != nil {
					releaseTopicObject(pool, newCb)
					return err
				}

				sess.MarkMessage(message, "")
				releaseTopicObject(pool, newCb)
				span.End()
			}
		case <-sess.Context().Done():
			log.Warnf("kafka session done, reason: %v", sess.Context().Err())
			sess.Commit()
			return sess.Context().Err()
		}
	}
}

// 回收Topic对象
func releaseTopicObject(pool *sync.Pool, obj Topic) {
	v := reflect.ValueOf(obj).Elem()
	v.Set(reflect.Zero(v.Type()))
	pool.Put(obj)
}

func consumerError() {
	for err := range client.consumer.Errors() {
		log.Errorf("kafka consumer error: %v", err)
	}
}