package logger

import (
	"context"
	"fmt"
	"os"
	"strconv"

	mlogs "github.com/caraml-dev/merlin/pkg/log"
	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"
)

type KafkaProducer interface {
	Produce(*kafka.Message, chan kafka.Event) error
	Events() chan kafka.Event
}

type KafkaAdmin interface {
	CreateTopics(context.Context, []kafka.TopicSpecification, ...kafka.CreateTopicsAdminOption) ([]kafka.TopicResult, error)
}

func getEnvOrDefault(key, fallback string) string {
	value := os.Getenv(key)
	if len(value) == 0 {
		return fallback
	}
	return value
}

func getEnvOrDefaultInt(key string, fallback int) int {
	value := os.Getenv(key)
	if len(value) == 0 {
		return fallback
	}
	if intVal, err := strconv.Atoi(value); err == nil {
		return intVal
	}
	return fallback
}

var (
	// Topic replication factor
	ReplicationFactor = getEnvOrDefaultInt("KAFKA_REPLICATION_FACTOR", 3)
	// Number of partitions for the Kafka topic
	NumPartitions = getEnvOrDefaultInt("KAFKA_NUM_PARTITIONS", 3)
	// The largest record batch size allowed by Kafka (after compression if compression is enabled)
	MaxMessageBytes = getEnvOrDefaultInt("KAFKA_MAX_MESSAGE_BYTES", 1048588)
	// The compression type for all data generated by the Producer
	CompressionType = getEnvOrDefault("KAFKA_COMPRESSION_TYPE", "snappy")
	// The maximum duration (ms) the Kafka Producer will block for to get Metadata, before timing out
	ConnectTimeoutMS = getEnvOrDefaultInt("KAFKA_CONNECT_TIMEOUT_MS", 1000)
)

type KafkaSink struct {
	logger   *zap.SugaredLogger
	producer KafkaProducer

	projectName  string
	modelName    string
	modelVersion string
}

func NewKafkaSink(
	logger *zap.SugaredLogger,
	producer KafkaProducer,
	adminClient KafkaAdmin,
	projectName string,
	modelName string,
	modelVersion string,
) (LogSink, error) {
	sink := &KafkaSink{
		logger:       logger,
		producer:     producer,
		projectName:  projectName,
		modelName:    modelName,
		modelVersion: modelVersion,
	}
	topicResults, err := adminClient.CreateTopics(context.Background(), []kafka.TopicSpecification{
		{
			Topic:             sink.topicName(),
			NumPartitions:     NumPartitions,
			ReplicationFactor: ReplicationFactor,
		},
	})
	if err != nil {
		return nil, err
	}
	for _, result := range topicResults {
		if result.Error.Code() != kafka.ErrNoError && result.Error.Code() != kafka.ErrTopicAlreadyExists {
			return nil, fmt.Errorf("failed to create topic %s: %w", result.Topic, result.Error)
		}
	}

	go sink.handleMessageDelivery()

	return sink, nil
}

func (k *KafkaSink) Sink(rawLogEntries []*LogEntry) error {
	// Format log inputs
	inferenceLogs, err := k.newLogMessages(rawLogEntries)
	if err != nil {
		return err
	}

	for _, inferenceLog := range inferenceLogs {
		keyBytes, valueBytes, err := k.buildNewKafkaMessage(inferenceLog)
		if err != nil {
			return err
		}
		topicName := k.topicName()
		err = k.producer.Produce(&kafka.Message{
			TopicPartition: kafka.TopicPartition{
				Topic:     &topicName,
				Partition: kafka.PartitionAny},
			Value: valueBytes,
			Key:   keyBytes,
		}, nil)
		if err != nil {
			k.logger.Errorf("Produce failed: %v\n", err)
		}
	}

	return nil
}

func (k *KafkaSink) topicName() string {
	return fmt.Sprintf("merlin-%s-%s-inference-log", k.projectName, k.modelName)
}

func (k *KafkaSink) handleMessageDelivery() {
	for e := range k.producer.Events() {
		switch ev := e.(type) {
		case *kafka.Message:
			if ev.TopicPartition.Error != nil {
				k.logger.Errorf("Delivery failed: %v\n", ev.TopicPartition.Error)
			}
		}
	}
}

func (k *KafkaSink) buildNewKafkaMessage(
	merlinLog *mlogs.InferenceLogMessage,
) (keyBytes []byte, valueBytes []byte, err error) {
	// Create the Kafka key
	key := &mlogs.InferenceLogKey{
		RequestId:      merlinLog.RequestId,
		EventTimestamp: merlinLog.EventTimestamp,
		ProjectName:    merlinLog.ProjectName,
		ModelName:      merlinLog.ModelName,
		ModelVersion:   merlinLog.ModelVersion,
	}

	// Marshal the key
	keyBytes, err = proto.Marshal(key)
	if err != nil {
		return nil, nil, fmt.Errorf("unable to marshal kafka key, %w", err)
	}

	// Marshal the message
	valueBytes, err = proto.Marshal(merlinLog)
	if err != nil {
		return nil, nil, fmt.Errorf("unable to marshal kafka value, %w", err)
	}

	return keyBytes, valueBytes, nil
}

func (k *KafkaSink) newLogMessages(logEntries []*LogEntry) ([]*mlogs.InferenceLogMessage, error) {
	messages := make([]*mlogs.InferenceLogMessage, 0)
	for _, logEntry := range logEntries {

		request := &mlogs.Request{}
		if logEntry.RequestPayload != nil {
			request = &mlogs.Request{
				Header: logEntry.RequestPayload.Headers,
				Body:   string(logEntry.RequestPayload.Body),
			}
		}

		response := &mlogs.Response{}
		if logEntry.ResponsePayload != nil {
			response = &mlogs.Response{
				StatusCode: int32(logEntry.ResponsePayload.StatusCode),
				Body:       string(logEntry.ResponsePayload.Body),
			}
		}

		logMessage := &mlogs.InferenceLogMessage{
			RequestId:      logEntry.RequestId,
			EventTimestamp: logEntry.EventTimestamp,
			ProjectName:    k.projectName,
			ModelName:      k.modelName,
			ModelVersion:   k.modelVersion,
			Request:        request,
			Response:       response,
		}

		messages = append(messages, logMessage)
	}

	return messages, nil
}
