package repository_ddb

import (
	"context"
	"dataservices/commons/errs"
	"dataservices/commons/metrics"
	"dataservices/commons/utils"
	"dataservices/ddb"
	"dataservices/internal/entity"
	"time"

	"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"
	tierrs "github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

func CreateClusterTable(ctx context.Context, cli *dynamodb.Client, replicaRegions []string) error {
	createTableInput := genTableSchema(userClusterTable(), types.AttributeDefinition{
		AttributeName: aws.String("cluster_id"),
		AttributeType: types.ScalarAttributeTypeN,
	}, &types.AttributeDefinition{
		AttributeName: aws.String("project_id"),
		AttributeType: types.ScalarAttributeTypeN,
	})
	return createGlobalTable(ctx, cli, createTableInput, replicaRegions)
}
func DeleteClusterTable(ctx context.Context, cli *dynamodb.Client, replicaRegions []string) error {
	return deleteGlobalTable(ctx, cli, userClusterTable(), replicaRegions)
}

func GenUserClusterInfoKey(projectId, clusterId uint64) (map[string]types.AttributeValue, error) {

	key := make(map[string]types.AttributeValue)

	clusterIdValue, err := attributevalue.Marshal(clusterId)
	if err != nil {
		return nil, err
	}
	projectIdValue, err := attributevalue.Marshal(projectId)
	if err != nil {
		return nil, err
	}
	key["cluster_id"] = clusterIdValue
	key["project_id"] = projectIdValue
	return key, nil
}

func CreateUserClusterInfoTableIfNotExist(ctx context.Context, cli *dynamodb.Client) error {
	exist, err := ddb.IsTableExist(ctx, cli, userClusterTable())
	if err != nil || exist {
		tilog.Info(ctx, "Table already exist", zap.String("table", userClusterTable()))
		return err
	}
	createTableInput := &dynamodb.CreateTableInput{
		AttributeDefinitions: []types.AttributeDefinition{{
			AttributeName: aws.String("cluster_id"),
			AttributeType: types.ScalarAttributeTypeN,
		}, {
			AttributeName: aws.String("project_id"),
			AttributeType: types.ScalarAttributeTypeN,
		}},
		KeySchema: []types.KeySchemaElement{{
			AttributeName: aws.String("cluster_id"),
			KeyType:       types.KeyTypeHash,
		}, {
			AttributeName: aws.String("project_id"),
			KeyType:       types.KeyTypeRange,
		}},
		TableName: aws.String(userClusterTable()),
		//ProvisionedThroughput: &types.ProvisionedThroughput{
		//	ReadCapacityUnits:  aws.Int64(10),
		//	WriteCapacityUnits: aws.Int64(10),
		//},
		BillingMode: types.BillingModePayPerRequest,
	}
	_, err = ddb.CreateTable(ctx, cli, createTableInput)
	return err
}

func PutUserClusterInfo(ctx context.Context, cli *dynamodb.Client, info *entity.UserClusterInfo) error {
	err := ddb.PutItem(ctx, cli, userClusterTable(), info)
	if err != nil {
		tilog.Error(ctx, "PutUserClusterInfo failed", zap.Error(err))
	}
	return err
}

func BatchPutUserClusterInfo(ctx context.Context, cli *dynamodb.Client, infos []*entity.UserClusterInfo) (int, error) {
	written, err := ddb.BatchPutItem(ctx, cli, userClusterTable(), utils.ToGeneralArr(infos...)...)
	if err != nil {
		tilog.Error(ctx, "BatchPutUserClusterInfo failed", zap.Error(err))
		tierrs.NewInternalDB("Failed to batch put cluster", err)
	}
	return written, err
}

func GetUserClusterInfo(ctx context.Context, cli *dynamodb.Client, projectId, clusterId uint64) (*entity.UserClusterInfo, error) {
	start := time.Now()
	defer func() {
		cost := time.Since(start)
		metrics.ExecutionLinkDuration.WithLabelValues("endpoint_exec", "get_cluster").Observe(cost.Seconds())
	}()
	key, err := GenUserClusterInfoKey(projectId, clusterId)
	if err != nil {
		return nil, err
	}
	var clusterInfo entity.UserClusterInfo
	err = ddb.GetItem(ctx, cli, userClusterTable(), key, &clusterInfo)
	if err != nil {
		if tierrs.IsError(err, tierrs.ErrNotFound) {
			return nil, errs.ErrUserClusterNotFound
		}
		return nil, err
	}
	return &clusterInfo, err
}

type UserClusterKey struct {
	ProjectID uint64
	ClusterID uint64
}

func ListUserClusterByIDs(ctx context.Context, cli *dynamodb.Client, clusterKeys []UserClusterKey) ([]*entity.UserClusterInfo, error) {

	if len(clusterKeys) == 0 {
		return nil, nil
	}

	keys := make([]map[string]types.AttributeValue, 0, len(clusterKeys))
	for _, clusterKey := range clusterKeys {
		key, err := GenUserClusterInfoKey(clusterKey.ProjectID, clusterKey.ClusterID)
		if err != nil {
			return nil, err
		}
		keys = append(keys, key)
	}
	items, err := ddb.BatchGetItem(ctx, cli, userClusterTable(), keys...)
	if err != nil {
		return nil, err
	}
	var clusterInfos []*entity.UserClusterInfo
	err = attributevalue.UnmarshalListOfMaps(items, &clusterInfos)
	if err != nil {
		return nil, err
	}
	return clusterInfos, err
}

func QueryUserClusterInfo(ctx context.Context, cli *dynamodb.Client, projectId, clusterId uint64) ([]entity.UserClusterInfo, error) {
	var err error
	var response *dynamodb.QueryOutput
	var clusterInfos []entity.UserClusterInfo
	projectIdEx := expression.Key("project_id").Equal(expression.Value(projectId))
	clusterIdEx := expression.Key("cluster_id").Equal(expression.Value(clusterId))
	expr, err := expression.NewBuilder().WithKeyCondition(projectIdEx.And(clusterIdEx)).Build()
	if err != nil {
		return nil, err
	}
	response, err = cli.Query(ctx, &dynamodb.QueryInput{
		TableName:                 aws.String(userClusterTable()),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		KeyConditionExpression:    expr.KeyCondition(),
	})
	if err != nil {
		return nil, err
	}
	err = attributevalue.UnmarshalListOfMaps(response.Items, &clusterInfos)
	if err != nil {
		return nil, err
	}
	return clusterInfos, err
}

func QueryUserClusterInfoWithPendingPrivateLink(ctx context.Context, cli *dynamodb.Client) ([]entity.UserClusterInfo, error) {
	conn1 := expression.Name("endpoint_id").NotEqual(expression.Value(""))
	conn2 := expression.Name("private_dns").Equal(expression.Value(""))
	expr, err := expression.NewBuilder().WithFilter(conn1.And(conn2)).Build()
	if err != nil {
		return nil, err
	}
	response, err := cli.Scan(ctx, &dynamodb.ScanInput{
		TableName:                 aws.String(userClusterTable()),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		FilterExpression:          expr.Filter(),
	})
	if err != nil {
		return nil, err
	}
	var clusterInfos []entity.UserClusterInfo
	err = attributevalue.UnmarshalListOfMaps(response.Items, &clusterInfos)
	if err != nil {
		return nil, err
	}
	return clusterInfos, err
}

func QueryUserDedicatedClusterInfoWithPage(ctx context.Context, cli *dynamodb.Client, exclusiveStartKey map[string]types.AttributeValue, limit int32) (map[string]types.AttributeValue, []entity.UserClusterInfo, error) {
	conn1 := expression.Name("is_dedicate").Equal(expression.Value(true))
	expr, err := expression.NewBuilder().WithFilter(conn1).Build()
	if err != nil {
		return nil, nil, err
	}
	request := &dynamodb.ScanInput{
		TableName:                 aws.String(userClusterTable()),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		FilterExpression:          expr.Filter(),
		ExclusiveStartKey:         exclusiveStartKey,
		Limit:                     &limit,
	}
	response, err := cli.Scan(ctx, request)
	if err != nil {
		return nil, nil, err
	}
	var clusterInfos []entity.UserClusterInfo
	err = attributevalue.UnmarshalListOfMaps(response.Items, &clusterInfos)
	if err != nil {
		return nil, nil, err
	}
	return response.LastEvaluatedKey, clusterInfos, err
}

// Use PutUserClusterInfo intead of UpdateUserClusterInfo.
// UpdateUserClusterInfo won't return error even the record not exist, it will auto create an record.
func UpdateUserClusterInfo(ctx context.Context, cli *dynamodb.Client, clusterInfo *entity.UserClusterInfo) error {
	key, err := GenUserClusterInfoKey(clusterInfo.ProjectId, clusterInfo.ClusterId)
	if err != nil {
		return err
	}
	update := expression.Set(expression.Name("private_dns"), expression.Value(clusterInfo.PrivateDNS))
	expr, err := expression.NewBuilder().WithUpdate(update).Build()
	if err != nil {
		return err
	}
	_, err = cli.UpdateItem(ctx, &dynamodb.UpdateItemInput{
		TableName:                 aws.String(userClusterTable()),
		Key:                       key,
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		UpdateExpression:          expr.Update(),
	})
	return err
}

func ScanUserClusterInfo(ctx context.Context, cli *dynamodb.Client) ([]entity.UserClusterInfo, error) {
	result, err := cli.Scan(ctx, &dynamodb.ScanInput{
		TableName: aws.String(userClusterTable()),
	})
	if err != nil {
		return nil, err
	}
	var clusterInfos []entity.UserClusterInfo
	err = attributevalue.UnmarshalListOfMaps(result.Items, &clusterInfos)
	if err != nil {
		return nil, err
	}
	return clusterInfos, nil
}

func DeleteUserClusterInfo(ctx context.Context, cli *dynamodb.Client, projectId, clusterId uint64) error {
	key, err := GenUserClusterInfoKey(projectId, clusterId)
	if err != nil {
		return err
	}
	return ddb.DeleteItem(ctx, cli, userClusterTable(), key)
}