package server

import (
	"context"
	"sync"

	"ky-go-kratos/app/websocket/internal/conf"
	"ky-go-kratos/app/websocket/internal/service"
	"ky-go-kratos/pkg/kafka"

	"github.com/go-kratos/kratos/v2/log"
)

type KafkaConsumerServer struct {
	consumer       *kafka.KafkaConsumer
	messageHandler *service.MessageHandler
	wg             sync.WaitGroup
}

func NewKafkaConsumerServer(c *conf.Data, messageHandler *service.MessageHandler) (*KafkaConsumerServer, func()) {
	consumer := kafka.NewKafkaReader(c.Kafka.Brokers, c.Kafka.Topic, int(c.Kafka.Partition))
	server := &KafkaConsumerServer{
		consumer:       consumer,
		messageHandler: messageHandler,
	}

	ctx := context.Background()
	if err := server.Start(ctx); err != nil {
		panic(err)
	}

	return server, func() {
		if err := server.Stop(ctx); err != nil {
			log.Errorf("failed to stop kafka consumer: %v", err)
		}
	}
}

func (s *KafkaConsumerServer) Start(ctx context.Context) error {
	if err := s.messageHandler.Start(); err != nil {
		return err
	}

	s.wg.Add(1)
	go s.consumeMessages(ctx)
	return nil
}

func (s *KafkaConsumerServer) Stop(ctx context.Context) error {
	s.wg.Wait()
	if err := s.messageHandler.Stop(); err != nil {
		return err
	}
	return s.consumer.Close()
}

func (s *KafkaConsumerServer) consumeMessages(ctx context.Context) {
	defer s.wg.Done()

	for {
		select {
		case <-ctx.Done():
			return
		default:
			message, err := s.consumer.Reader.ReadMessage(ctx)
			log.Infof("message data: %v", message)
			if err != nil {
				if s.consumer.IsTransientNetworkError(err) {
					continue
				}
				return
			}

			if err := s.messageHandler.HandleMessage(&message); err != nil {
				log.Errorf("failed to handle message: %v", err)
			}
		}
	}
}