package ddb

import (
	"context"
	"errors"
	"fmt"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"
	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/expression"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
	"github.com/tidbcloud/account/internel/dao/entity"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

// DdbApiKey is the struct of api_key table in DynamoDB
// The primary key of this table is identity(partition key)
// Example:
//
//	{
//	  "identity": "api_key_kvxFZir",
//	  "private_key": "182938920180dhiahda",
//	  "role_map": {1: [1, 2]}  // app_id: [role_id]
//	}
type DdbApiKey struct {
	AccessKey    string              `dynamodbav:"identity"`
	PrivateKey   string              `dynamodbav:"private_key"`
	AppRoleMap   map[uint64][]uint64 `dynamodbav:"role_map"`
	RateLimitRpm uint32              `dynamodbav:"rate_limit_rpm"`
}

// GetKey returns the composite primary key of the apiKey in a format that can be
// sent to DynamoDB.
func (apiKey *DdbApiKey) GetKey(ctx context.Context) (map[string]types.AttributeValue, error) {
	accessKey, err := attributevalue.Marshal(apiKey.AccessKey)
	if err != nil {
		tilog.Error(ctx, "failed to marshal accessKey", zap.Error(err))
		return nil, err
	}
	return map[string]types.AttributeValue{"identity": accessKey}, nil
}

func FormatAccessKey(accessKey string) string {
	return fmt.Sprintf("api_key_%s", accessKey)
}

func CreateApiKey(
	ctx context.Context, ddbClient *DynamoDBClient,
	accessKey string, privateKey string, rateLimitRpm uint32, roleBindings []entity.RoleBinding,
) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "CreateApiKey"))
		return nil
	}

	appRoleMap := BuildAppRoleMap(roleBindings)
	ak := DdbApiKey{
		AccessKey:    FormatAccessKey(accessKey),
		PrivateKey:   privateKey,
		AppRoleMap:   appRoleMap,
		RateLimitRpm: rateLimitRpm,
	}
	item, err := attributevalue.MarshalMap(ak)
	if err != nil {
		tilog.Error(ctx, "failed to marshal apiKey", zap.Error(err))
		return err
	}
	_, err = ddbClient.PutItem(ctx, &dynamodb.PutItemInput{
		TableName: aws.String(ddbClient.ApiKeyTableName),
		Item:      item,
	})
	if err != nil {
		tilog.Error(ctx, "failed to create apiKey in dynamodb", zap.Error(err))
	}
	return err
}

// DeleteApiKey deletes one record from the apiKey table by accessKey
func DeleteApiKey(ctx context.Context, ddbClient *DynamoDBClient, accessKey string) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "DeleteApiKey"))
		return nil
	}

	ak := DdbApiKey{
		AccessKey: FormatAccessKey(accessKey),
	}
	itemKey, err := ak.GetKey(ctx)
	if err != nil {
		return err
	}
	_, err = ddbClient.DeleteItem(ctx, &dynamodb.DeleteItemInput{
		TableName: aws.String(ddbClient.ApiKeyTableName),
		Key:       itemKey,
	})
	if err != nil {
		tilog.Error(ctx, "failed to delete apiKey in dynamodb", zap.Error(err))
	}
	return err
}

// UpdateApiKey updates the privateKey,rateLimitRpm of the apiKey by accessKey
func UpdateApiKey(ctx context.Context, ddbClient *DynamoDBClient, accessKey string, updateDdbMap map[string]any) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "UpdateApiKey"))
		return nil
	}
	if len(updateDdbMap) == 0 {
		tilog.Info(ctx, "privateKey and rateLimitRpm no change，update ddb ignore", zap.Any("accessKey", accessKey), zap.Any("updateDdbMap", updateDdbMap))
		return nil
	}

	var update expression.UpdateBuilder

	i := 0
	for k, v := range updateDdbMap {
		if i == 0 {
			update = expression.Set(expression.Name(k), expression.Value(v))
		} else {
			update.Set(expression.Name(k), expression.Value(v))
		}
		tilog.Info(ctx, "DDB UpdateBuilder "+fmt.Sprintf("%d", i), zap.Any("key", k), zap.Any("value", v))
		i++
	}
	tilog.Info(ctx, "create/update api key in ddb", zap.Any("accessKey", accessKey), zap.Any("updateDdbMap", updateDdbMap))

	ak := DdbApiKey{
		AccessKey: FormatAccessKey(accessKey),
	}
	itemKey, err := ak.GetKey(ctx)
	if err != nil {
		return err
	}
	expr, err := expression.NewBuilder().WithUpdate(update).Build()
	if err != nil {
		return err
	}
	_, err = ddbClient.UpdateItem(ctx, &dynamodb.UpdateItemInput{
		TableName:                 aws.String(ddbClient.ApiKeyTableName),
		Key:                       itemKey,
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		UpdateExpression:          expr.Update(),
	})
	if err != nil {
		tilog.Error(ctx, "failed to update apiKey's private_key in dynamodb", zap.Error(err))
	}
	return err
}

// appRoleMap[scopeId]roleId
func BuildAppRoleMap(roleBindings []entity.RoleBinding) map[uint64][]uint64 {
	appRoleMap := make(map[uint64][]uint64)
	for _, roleBinding := range roleBindings {
		_key := roleBinding.ScopeID.Uint64()
		appRoleMap[_key] = append(appRoleMap[_key], roleBinding.RoleID.Uint64())
	}
	return appRoleMap
}

// UpdateApiKeyRoleMap updates the role_map of the apiKey by accessKey
func UpdateApiKeyRoleMap(ctx context.Context, ddbClient *DynamoDBClient, accessKey string, roleBindings []entity.RoleBinding) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "UpdateApiKeyRoleMap"))
		return nil
	}

	appRoleMap := BuildAppRoleMap(roleBindings)
	updateDdbMap := map[string]any{
		"role_map": appRoleMap,
	}
	tilog.Info(ctx, "UpdateApiKeyRoleMap", zap.Any("accessKey", accessKey), zap.Any("appRoleMap", appRoleMap))
	return UpdateApiKey(ctx, ddbClient, accessKey, updateDdbMap)
	//update := expression.Set(expression.Name("role_map"), expression.Value(appRoleMap))
	//expr, err := expression.NewBuilder().WithUpdate(update).Build()
	//if err != nil {
	//	return err
	//}
	//_, err = ddbClient.UpdateItem(ctx, &dynamodb.UpdateItemInput{
	//	TableName:                 aws.String(ddbClient.ApiKeyTableName),
	//	Key:                       itemKey,
	//	ExpressionAttributeNames:  expr.Names(),
	//	ExpressionAttributeValues: expr.Values(),
	//	UpdateExpression:          expr.Update(),
	//})
	//if err != nil {
	//	tilog.Error(ctx, "failed to update apiKey's role_map in dynamodb", zap.Error(err))
	//}
	//return err
}

// GetApiKey gets one record from the apiKey table by accessKey
func GetApiKey(ctx context.Context, ddbClient *DynamoDBClient, accessKey string) (*DdbApiKey, error) {
	ak := DdbApiKey{
		AccessKey: FormatAccessKey(accessKey),
	}
	itemKey, err := ak.GetKey(ctx)
	if err != nil {
		return nil, err
	}
	output, err := ddbClient.GetItem(ctx, &dynamodb.GetItemInput{
		TableName: aws.String(ddbClient.ApiKeyTableName),
		Key:       itemKey,
	})
	if err != nil {
		tilog.Error(ctx, "failed to get apiKey in dynamodb", zap.Error(err))
		return nil, err
	}
	if output.Item == nil {
		return nil, nil
	}
	var apiKey DdbApiKey
	err = attributevalue.UnmarshalMap(output.Item, &apiKey)
	if err != nil {
		tilog.Warn(ctx, "failed to unmarshal apiKey", zap.Error(err))
		return nil, err
	}
	return &apiKey, nil
}

// BatchGetApiKeys gets the apiKeys by accessKeys, max input size is 100
func BatchGetApiKeys(ctx context.Context, ddbClient *DynamoDBClient, accessKeys []string) ([]DdbApiKey, error) {
	var ddbApiKeys []DdbApiKey
	if len(accessKeys) > 100 {
		return nil, errors.New("max input size is 100")
	}
	// dynamodb batchGetItem can't handle duplicate keys
	accessKeys = removeDuplicateStr(accessKeys)
	keys := make([]map[string]types.AttributeValue, 0, len(accessKeys))
	for _, accessKey := range accessKeys {
		keys = append(keys, map[string]types.AttributeValue{
			"identity": &types.AttributeValueMemberS{Value: FormatAccessKey(accessKey)},
		})
	}

	output, err := ddbClient.BatchGetItem(ctx, &dynamodb.BatchGetItemInput{
		RequestItems: map[string]types.KeysAndAttributes{
			ddbClient.ApiKeyTableName: {
				Keys: keys,
			},
		},
	})
	if err != nil {
		tilog.Error(ctx, "failed to get apiKeys from dynamodb", zap.Error(err))
		return nil, err
	}
	err = attributevalue.UnmarshalListOfMaps(output.Responses[ddbClient.ApiKeyTableName], &ddbApiKeys)
	if err != nil {
		tilog.Error(ctx, "failed to unmarshal apiKeys from dynamodb", zap.Error(err))
		return nil, err
	}
	return ddbApiKeys, nil
}

func removeDuplicateStr(strSlice []string) []string {
	allKeys := make(map[string]bool)
	list := []string{}
	for _, item := range strSlice {
		if _, exists := allKeys[item]; !exists {
			allKeys[item] = true
			list = append(list, item)
		}
	}
	return list
}