package firehose

import (
	"context"
	"fmt"
	"sort"
	"time"

	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	aws_firehose "github.com/aws/aws-sdk-go-v2/service/firehose"
	"github.com/aws/aws-sdk-go-v2/service/firehose/types"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/pingcap/log"
	"github.com/prometheus/prom2json"

	"github.com/tidbcloud/sli-reporter/slprober"

	"go.uber.org/zap"

	"github.com/prometheus/client_golang/prometheus"
)

var (
	HashSize    = 79
	MessageSize = 512*1024 - HashSize // 512K
)

type Sender struct {
	deliveryStreamCli  *aws_firehose.Client
	deliveryStreamName string
	reporterMode       string
}

func NewSender(
	reporterMode string,
	deliveryStreamName string,
	deliveryStreamRoleArn string,
) *Sender {
	cfg, err := config.LoadDefaultConfig(context.Background())
	if err != nil {
		panic(err)
	}
	if deliveryStreamRoleArn != "" {
		stsCli := sts.NewFromConfig(cfg)
		cred := aws.NewCredentialsCache(
			stscreds.NewAssumeRoleProvider(stsCli, deliveryStreamRoleArn))
		cfg.Credentials = cred
	}
	svc := aws_firehose.NewFromConfig(cfg)

	return &Sender{
		reporterMode:       reporterMode,
		deliveryStreamCli:  svc,
		deliveryStreamName: deliveryStreamName,
	}
}

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

type MessageWithSignature struct {
	Message
	Signature string `json:"signature"`
}

func (s *Sender) Send(body *Message) {
	messages := SplitMessage(body)

	for _, message := range messages {
		content := message
		if len(content) > 512 {
			content = content[:512]
		}
		signatureMessage, err := GenerateSignatureHashForSLIData(message)
		if err != nil {
			log.Error("generate signature hash for message failed", zap.Error(err),
				zap.String("body", content))
			return
		}

		_, err = s.deliveryStreamCli.PutRecord(context.Background(), &aws_firehose.PutRecordInput{
			DeliveryStreamName: &s.deliveryStreamName,
			Record: &types.Record{
				Data: signatureMessage,
			},
		})
		if err != nil {
			log.Error("Send record failed", zap.Error(err),
				zap.String("body", content))
			failedToSendMessage()
		} else {
			log.Info("Send record successful", zap.Int("len", len(signatureMessage)), zap.String("body", content))
		}
	}
}

func (s *Sender) SendMetricsFromRegistry(reg *prometheus.Registry) error {
	family, err := reg.Gather()
	if err != nil {
		return err
	}
	for _, fa := range family {
		metrics := prom2json.NewFamily(fa)
		message := &Message{
			Timestamp: time.Now().Unix(),
			Instance:  "dbaas-reporter",
			Metrics:   metrics,
		}
		go s.Send(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 := range splitTime {
		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
}

func GenerateSignatureHashForSLIData(sliData string) ([]byte, error) {
	var parsedData interface{}
	err := json.Unmarshal([]byte(sliData), &parsedData)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal message: %w", err)
	}

	normalizedJSON, err := normalizeJSON(parsedData)
	if err != nil {
		return nil, fmt.Errorf("failed to normalize JSON: %w", err)
	}

	dataHash := sha256.Sum256([]byte(normalizedJSON))
	dataHashHex := hex.EncodeToString(dataHash[:]) // Convert to a 64-character hexadecimal string
	secretKey, err := IRKE(dataHashHex)
	if err != nil {
		return nil, fmt.Errorf("failed to generate secret key: %w", err)
	}
	hmacHash := hmac.New(sha256.New, []byte(secretKey))
	_, err = hmacHash.Write([]byte(dataHashHex))
	if err != nil {
		return nil, fmt.Errorf("failed to calculate HMAC: %w", err)
	}
	signatureHash := hex.EncodeToString(hmacHash.Sum(nil))

	message := Message{}
	err = json.Unmarshal([]byte(sliData), &message)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal message: %w", err)
	}

	messageWithSignature := MessageWithSignature{
		message,
		signatureHash,
	}

	d, err := json.Marshal(messageWithSignature)
	if err != nil {
		return nil, err
	}

	return d, nil
}

// IRKE implements the Interleaved-Rotational Key Extraction algorithm
// Input: 64-character hexadecimal SHA-256 hash string
// Output: 32-character hexadecimal key string
func IRKE(hash string) (string, error) {
	// Verify input length is 64 characters
	if len(hash) != 64 {
		return "", fmt.Errorf("input must be a 64-character hexadecimal SHA-256 hash")
	}

	// Step 1: Split hash into A and B (32 characters each, 128 bits)
	A := hash[:32] // First 32 characters
	B := hash[32:] // Last 32 characters

	// Convert A and B from hex strings to byte arrays (16 bytes each, 128 bits)
	aBytes, err := hex.DecodeString(A)
	if err != nil {
		return "", fmt.Errorf("segment A is not valid hexadecimal: %w", err)
	}
	bBytes, err := hex.DecodeString(B)
	if err != nil {
		return "", fmt.Errorf("segment B is not valid hexadecimal: %w", err)
	}

	// Step 2: XOR A and B to get K1 (128 bits)
	k1 := make([]byte, 16)
	for i := 0; i < 16; i++ {
		k1[i] = aBytes[i] ^ bBytes[i]
	}

	// Step 3: Dynamic rotation to get K2
	// 3.1 Extract last two characters of H (H63H64)
	lastTwoChars := hash[62:]
	// 3.2 Convert to integer r (0-255)
	lastByte, err := hex.DecodeString(lastTwoChars)
	if err != nil || len(lastByte) != 1 {
		return "", fmt.Errorf("failed to extract rotation amount")
	}
	r := int(lastByte[0])
	// 3.3 Take rotation amount modulo 128 (since K1 is 128 bits)
	rMod := r % 128
	// 3.4 Perform left cyclic rotation on K1 by rMod bits to get K2
	k2 := rotateLeft128(k1, rMod)

	// Step 4: Interleaving and final XOR to get Kfinal
	// 4.1 Extract C1 (H17..32, i.e., indices 16-31, 16 characters) and C2 (H49..64, i.e., indices 48-63, 16 characters)
	C1 := hash[16:32]
	C2 := hash[48:64]
	// 4.2 Concatenate C1 and C2 to get P (32 characters, 128 bits)
	P := C1 + C2
	pBytes, err := hex.DecodeString(P)
	if err != nil {
		return "", fmt.Errorf("segment P is not valid hexadecimal: %w", err)
	}
	// 4.3 XOR K2 with P to get Kfinal
	kFinal := make([]byte, 16)
	for i := 0; i < 16; i++ {
		kFinal[i] = k2[i] ^ pBytes[i]
	}

	// Step 5: Convert to 32-character hexadecimal string
	kSecret := hex.EncodeToString(kFinal)

	return kSecret, nil
}

// rotateLeft128 performs left cyclic rotation on 128-bit (16-byte) data
// bits: number of bits to rotate (0-127)
func rotateLeft128(data []byte, bits int) []byte {
	if len(data) != 16 {
		panic("rotateLeft128: input must be 16 bytes (128 bits)")
	}
	if bits <= 0 || bits >= 128 {
		return data // Rotating 0 or 128 bits is equivalent to no rotation
	}

	// Convert 16-byte array to 128-bit slice (each element is 0 or 1)
	bitsSlice := make([]uint8, 128)
	for i := 0; i < 16; i++ {
		for j := 0; j < 8; j++ {
			// Extract the j-th bit of each byte (high bit first)
			bitsSlice[i*8+j] = (data[i] >> (7 - j)) & 0x01
		}
	}

	// Left cyclic rotation: move first 'bits' bits to the end
	var rotated = append(bitsSlice[bits:], bitsSlice[:bits]...)

	// Convert bit slice back to 16-byte array
	result := make([]byte, 16)
	for i := 0; i < 16; i++ {
		var b byte
		for j := 0; j < 8; j++ {
			b = (b << 1) | rotated[i*8+j]
		}
		result[i] = b
	}

	return result
}

// normalizeJSON Normalize JSON data: sort fields, remove unnecessary spaces
func normalizeJSON(data interface{}) (string, error) {
	normalized := normalizeValue(data)

	// use compact JSON format (no spaces)
	bytes, err := json.Marshal(normalized)
	if err != nil {
		return "", err
	}
	return string(bytes), nil
}

// normalizeValue
func normalizeValue(value interface{}) interface{} {
	switch v := value.(type) {
	case map[string]interface{}:
		normalizedMap := make(map[string]interface{})

		keys := make([]string, 0, len(v))
		for k := range v {
			keys = append(keys, k)
		}
		sort.Strings(keys)

		for _, k := range keys {
			normalizedMap[k] = normalizeValue(v[k])
		}
		return normalizedMap

	case []interface{}:
		normalizedArray := make([]interface{}, len(v))
		for i, item := range v {
			normalizedArray[i] = normalizeValue(item)
		}
		return normalizedArray

	default:
		return v
	}
}

func failedToSendMessage() {
	slprober.SendMessageFailedTotal.WithLabelValues().Inc()
}