package repository_ddb

import (
	"context"
	"dataservices/commons/utils"
	"dataservices/ddb"
	"dataservices/internal/entity"

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

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

func DeleteAppClusterTable(ctx context.Context, cli *dynamodb.Client, replicaRegions []string) error {
	return deleteGlobalTable(ctx, cli, appClusterTable(), replicaRegions)
}

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

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

func DeleteAppCluster(ctx context.Context, cli *dynamodb.Client, dataappId, clusterId uint64) error {
	key, err := genAppClusterKey(dataappId, clusterId)
	if err != nil {
		return errs.NewBadRequest("Failed to generate app cluster key", err)
	}
	err = ddb.DeleteItem(ctx, cli, appClusterTable(), key)
	if err != nil {
		return errs.NewInternalDB("Failed to delete app cluster", err)
	}
	return nil
}

func GetAppCluster(ctx context.Context, cli *dynamodb.Client, dataappId, clusterId uint64) (*entity.AppCluster, error) {
	key, err := genAppClusterKey(dataappId, clusterId)
	if err != nil {
		return nil, errs.NewBadRequest("Failed to generate app cluster key", err)
	}
	response, err := cli.GetItem(ctx, &dynamodb.GetItemInput{
		Key:       key,
		TableName: aws.String(appClusterTable()),
	})
	if err != nil {
		return nil, errs.NewInternalDB("Failed to get app cluster", err)
	} else if len(response.Item) == 0 {
		return nil, errs.NewNotFound("App cluster not exists", nil)
	}
	var clusterInfo entity.AppCluster
	err = attributevalue.UnmarshalMap(response.Item, &clusterInfo)
	if err != nil {
		return nil, errs.NewInternalJsonMarshal("Failed to parse app cluster", err)
	}
	return &clusterInfo, nil
}

func BatchDeleteAppCluster(ctx context.Context, cli *dynamodb.Client, dataappID uint64, clusterIDs []uint64) (int, error) {
	keys := make([]map[string]types.AttributeValue, 0, len(clusterIDs))
	for _, clusterID := range clusterIDs {
		key, err := genAppClusterKey(dataappID, clusterID)
		if err != nil {
			return 0, err
		}
		keys = append(keys, key)
	}
	written, err := ddb.BatchDeleteItem(ctx, cli, appClusterTable(), keys...)
	if err != nil {
		return 0, errs.NewInternalDB("Failed to batch delete app cluster", err)
	}
	return written, nil
}

func DeleteAppClusterByAppID(ctx context.Context, cli *dynamodb.Client, dataappID uint64) error {
	clusterList, err := QueryAppClusterByAppID(ctx, cli, dataappID)
	if err != nil {
		return errs.NewInternalDB("Failed to query app's link cluster", err)
	} else if len(clusterList) == 0 {
		return nil
	}

	clusterIDs := make([]uint64, 0, len(clusterList))
	for _, cluster := range clusterList {
		clusterIDs = append(clusterIDs, cluster.ClusterID)
	}
	_, err = BatchDeleteAppCluster(ctx, cli, dataappID, clusterIDs)
	if err != nil {
		return errs.NewInternalDB("Failed to batch delete app's link cluster", err)
	}
	return nil
}

func QueryAppClusterByAppID(ctx context.Context, cli *dynamodb.Client, dataappID uint64) ([]*entity.AppCluster, error) {

	dataappExp := expression.Key("dataapp_id").Equal(expression.Value(dataappID))
	exp, err := expression.NewBuilder().WithKeyCondition(dataappExp).Build()
	if err != nil {
		return nil, err
	}
	response, err := cli.Query(ctx, &dynamodb.QueryInput{
		TableName:                 aws.String(appClusterTable()),
		ExpressionAttributeNames:  exp.Names(),
		ExpressionAttributeValues: exp.Values(),
		KeyConditionExpression:    exp.KeyCondition(),
	})
	if err != nil {
		return nil, err
	}
	var infos []*entity.AppCluster
	err = attributevalue.UnmarshalListOfMaps(response.Items, &infos)
	if err != nil {
		return nil, err
	}
	return infos, err
}

func genAppClusterKey(dataappId, clusterId uint64) (map[string]types.AttributeValue, error) {
	dataappIdValue, err := attributevalue.Marshal(dataappId)
	if err != nil {
		return nil, err
	}
	clusterIdValue, err := attributevalue.Marshal(clusterId)
	if err != nil {
		return nil, err
	}
	return map[string]types.AttributeValue{"dataapp_id": dataappIdValue, "cluster_id": clusterIdValue}, nil
}
