package sqs

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"time"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/aws/aws-sdk-go-v2/service/firehose"
	"github.com/aws/aws-sdk-go-v2/service/firehose/types"
	"github.com/pingcap/log"
	"github.com/prometheus/prom2json"

	"go.uber.org/zap"
)

var MessageSize = 256 * 1024 // 256K

var (
	SLOList = []string{
		"observability_service_level_dedicated_prometheus_api_available_counter",
		"observability_service_level_shared_vm_api_available_counter",
		"observability_service_level_shared_loki_api_available_counter",
		"observability_service_level_create_cluster_api_available_counter",
		"observability_service_level_get_cluster_api_available_counter",
		"observability_service_level_s3_logs_upload_delay_histogram",
		"observability_service_level_s3_metrics_upload_delay_histogram",
		"observability_service_level_prometheus_metrics_upload_delay_histogram",
	}
	SLOMetricNameMap map[string]struct{}

	streamName string
	svc        *firehose.Client
)

func InitSQSConfig() {
	SLOMetricNameMap = make(map[string]struct{}, len(SLOList))
	for _, name := range SLOList {
		SLOMetricNameMap[name] = struct{}{}
	}

	streamName = os.Getenv("AWS_DELIVERY_STREAM_NAME")

	cfg, err := config.LoadDefaultConfig(context.Background())
	if err != nil {
		panic(err)
	}
	svc = firehose.NewFromConfig(cfg)
}

type Message struct {
	Timestamp int64             `json:"timestamp"`
	Instance  string            `json:"instance"`
	Metrics   *prom2json.Family `json:"metrics"`
}

func Send(svc *firehose.Client, body *Message) {
	// Because the single message size is limited to 256K.
	// So we need to check the message size.
	messages := SplitMessage(body)

	for _, message := range messages {
		content := message
		if len(content) > 256 {
			content = content[:256]
		}
		_, err := svc.PutRecord(context.Background(), &firehose.PutRecordInput{
			DeliveryStreamName: &streamName,
			Record: &types.Record{
				Data: []byte(message),
			},
		})
		if err != nil {
			log.Error("Send record failed", zap.Error(err),
				zap.String("body", content))
		} else {
			log.Info("Send record successful", zap.String("body", content))
		}
	}
}

func SendSLAMessage(reg *prometheus.Registry) error {
	family, err := reg.Gather()
	if err != nil {
		return err
	}
	for _, fa := range family {
		metrics := prom2json.NewFamily(fa)
		if _, ok := SLOMetricNameMap[metrics.Name]; !ok {
			continue
		}
		message := &Message{
			Timestamp: time.Now().Unix(),
			Instance:  "observability-reporter",
			Metrics:   metrics,
		}
		go Send(svc, message)
	}
	return nil
}

func SplitMessage(body *Message) []string {
	var (
		d   []byte
		err error
	)
	r := make([]string, 0)
	d, err = json.Marshal(body)
	if err != nil {
		zap.Error(err)
		return nil
	}

	if len(d) > MessageSize {
		log.Info("Split the message", zap.String("instance", body.Instance))
		splitTime := (len(d) / MessageSize) + 2
		var me []*prom2json.Family
		switch body.Metrics.Type {
		case "GAUGE", "COUNTER":
			me = SplitGauge(body.Metrics, splitTime)
		default:
			zap.Error(fmt.Errorf("unknown split message type: %s", body.Metrics.Type))
			return nil
		}
		for _, metric := range me {
			metric := metric
			message := &Message{
				Timestamp: body.Timestamp,
				Instance:  body.Instance,
				Metrics:   metric,
			}
			r = append(r, SplitMessage(message)...)
		}
	} else {
		r = append(r, string(d))
	}
	return r
}

func SplitGauge(metrics *prom2json.Family, splitTime int) (r []*prom2json.Family) {
	index := len(metrics.Metrics) / splitTime
	for i := 0; i < splitTime; i++ {
		start := i * index
		end := (i + 1) * index
		if i == splitTime-1 {
			end = len(metrics.Metrics)
		}
		me := &prom2json.Family{
			Name:    metrics.Name,
			Type:    metrics.Type,
			Metrics: metrics.Metrics[start:end],
		}
		r = append(r, me)
	}
	return r
}