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

const (
	endpointIndexDataappEndpoint        = "endpoint_idx_datapp_endpoint"
	endpointPublishIndexDataappEndpoint = "endpointpublish_idx_datapp_endpoint"
)

func CreateEndpointTable(ctx context.Context, cli *dynamodb.Client, replicaRegions []string) error {
	createTableInput := genTableSchema(endpointsTable(), types.AttributeDefinition{
		AttributeName: aws.String("id"),
		AttributeType: types.ScalarAttributeTypeN,
	}, nil)
	createTableInput.AttributeDefinitions = append(createTableInput.AttributeDefinitions, types.AttributeDefinition{
		AttributeName: aws.String("dataapp_id"),
		AttributeType: types.ScalarAttributeTypeN,
	}, types.AttributeDefinition{
		AttributeName: aws.String("endpoint"),
		AttributeType: types.ScalarAttributeTypeS,
	})
	createTableInput.GlobalSecondaryIndexes = append(createTableInput.GlobalSecondaryIndexes, types.GlobalSecondaryIndex{
		IndexName: aws.String(endpointIndexDataappEndpoint),
		KeySchema: []types.KeySchemaElement{
			{
				AttributeName: aws.String("dataapp_id"),
				KeyType:       types.KeyTypeHash,
			},
			{
				AttributeName: aws.String("endpoint"),
				KeyType:       types.KeyTypeRange,
			},
		},
		Projection: &types.Projection{
			ProjectionType: types.ProjectionTypeAll,
		},
	})
	return createGlobalTable(ctx, cli, createTableInput, replicaRegions)
}
func DeleteEndpointTable(ctx context.Context, cli *dynamodb.Client, replicaRegions []string) error {
	return deleteGlobalTable(ctx, cli, endpointsTable(), replicaRegions)
}

func CreateEndpointPublishTable(ctx context.Context, cli *dynamodb.Client, replicaRegions []string) error {
	createTableInput := genTableSchema(endpointPublishTable(), types.AttributeDefinition{
		AttributeName: aws.String("endpoint_id"),
		AttributeType: types.ScalarAttributeTypeN,
	}, nil)
	createTableInput.AttributeDefinitions = append(createTableInput.AttributeDefinitions, types.AttributeDefinition{
		AttributeName: aws.String("dataapp_id"),
		AttributeType: types.ScalarAttributeTypeN,
	}, types.AttributeDefinition{
		AttributeName: aws.String("endpoint"),
		AttributeType: types.ScalarAttributeTypeS,
	})
	createTableInput.GlobalSecondaryIndexes = append(createTableInput.GlobalSecondaryIndexes, types.GlobalSecondaryIndex{
		IndexName: aws.String(endpointPublishIndexDataappEndpoint),
		KeySchema: []types.KeySchemaElement{
			{
				AttributeName: aws.String("dataapp_id"),
				KeyType:       types.KeyTypeHash,
			},
			{
				AttributeName: aws.String("endpoint"),
				KeyType:       types.KeyTypeRange,
			},
		},
		Projection: &types.Projection{
			ProjectionType: types.ProjectionTypeAll,
		},
	})
	return createGlobalTable(ctx, cli, createTableInput, replicaRegions)
}
func DeleteEndpointPublishTable(ctx context.Context, cli *dynamodb.Client, replicaRegions []string) error {
	return deleteGlobalTable(ctx, cli, endpointPublishTable(), replicaRegions)
}

func CreateEndpointRegionTable(ctx context.Context, cli *dynamodb.Client) error {
	createTableInput := &dynamodb.CreateTableInput{
		AttributeDefinitions: []types.AttributeDefinition{{
			AttributeName: aws.String("id"),
			AttributeType: types.ScalarAttributeTypeN,
		}},
		KeySchema: []types.KeySchemaElement{{
			AttributeName: aws.String("id"),
			KeyType:       types.KeyTypeHash,
		}},
		TableName: aws.String(endpointsTable()),
		ProvisionedThroughput: &types.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(10),
			WriteCapacityUnits: aws.Int64(10),
		},
	}
	_, err := ddb.CreateTableForGlobal(ctx, cli, createTableInput)
	return err
}

func CreateEndpointPublishRegionTable(ctx context.Context, cli *dynamodb.Client) error {
	createTableInput := &dynamodb.CreateTableInput{
		AttributeDefinitions: []types.AttributeDefinition{{
			AttributeName: aws.String("endpoint_id"),
			AttributeType: types.ScalarAttributeTypeN,
		}},
		KeySchema: []types.KeySchemaElement{{
			AttributeName: aws.String("endpoint_id"),
			KeyType:       types.KeyTypeHash,
		}},
		TableName: aws.String(endpointPublishTable()),
		ProvisionedThroughput: &types.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(10),
			WriteCapacityUnits: aws.Int64(10),
		},
	}
	_, err := ddb.CreateTableForGlobal(ctx, cli, createTableInput)
	return err
}

func CreateEndpointGlobalTable(ctx context.Context, cli *dynamodb.Client, regions []string) error {
	return ddb.CreateGlobalTable(ctx, cli, endpointsTable(), regions)
}

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

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

func GetEndpoint(ctx context.Context, cli *dynamodb.Client, endpointId uint64) (*entity.Endpoint, error) {
	key, err := getEndpointDdbKey(endpointId)
	if err != nil {
		return nil, err
	}
	var endpoint entity.Endpoint
	if err := ddb.GetItem(ctx, cli, endpointsTable(), key, &endpoint); err != nil {
		return nil, err
	}
	return &endpoint, nil
}

func GetEndpointByName(ctx context.Context, cli *dynamodb.Client, dataAppID uint64, endpoint string) ([]*entity.Endpoint, error) {
	conn1 := expression.Key("dataapp_id").Equal(expression.Value(dataAppID))
	conn2 := expression.Key("endpoint").Equal(expression.Value(endpoint))
	expr, err := expression.NewBuilder().WithKeyCondition(conn1.And(conn2)).Build()
	if err != nil {
		return nil, err
	}
	response, err := cli.Query(ctx, &dynamodb.QueryInput{
		TableName:                 aws.String(endpointsTable()),
		IndexName:                 aws.String(endpointIndexDataappEndpoint),
		KeyConditionExpression:    expr.KeyCondition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
	})
	if err != nil {
		return nil, err
	}
	endpoints := []*entity.Endpoint{}
	err = attributevalue.UnmarshalListOfMaps(response.Items, &endpoints)
	if err != nil {
		return nil, err
	}
	return endpoints, nil
}

func GetPublishedEndpoint(ctx context.Context, cli *dynamodb.Client, endpointId uint64) (*entity.EndpointPublish, error) {
	key, err := getPublishedEndpointDdbKey(endpointId)
	if err != nil {
		return nil, err
	}
	var endpoint entity.EndpointPublish
	if err := ddb.GetItem(ctx, cli, endpointPublishTable(), key, &endpoint); err != nil {
		return nil, err
	}
	return &endpoint, nil
}

func GetAllDraftEndpointsByAppID(ctx context.Context, cli *dynamodb.Client, dataAppID uint64) ([]*entity.Endpoint, error) {
	conn1 := expression.Key("dataapp_id").Equal(expression.Value(dataAppID))
	expr, err := expression.NewBuilder().WithKeyCondition(conn1).Build()
	if err != nil {
		return nil, err
	}
	response, err := cli.Query(ctx, &dynamodb.QueryInput{
		TableName:                 aws.String(endpointsTable()),
		IndexName:                 aws.String(endpointIndexDataappEndpoint),
		KeyConditionExpression:    expr.KeyCondition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
	})
	if err != nil {
		return nil, err
	}
	endpoints := []*entity.Endpoint{}
	err = attributevalue.UnmarshalListOfMaps(response.Items, &endpoints)
	if err != nil {
		return nil, err
	}
	return endpoints, nil
}

func GetAllPublishedEndpointsByAppID(ctx context.Context, cli *dynamodb.Client, dataAppID uint64) ([]*entity.EndpointPublish, error) {
	conn1 := expression.Key("dataapp_id").Equal(expression.Value(dataAppID))
	expr, err := expression.NewBuilder().WithKeyCondition(conn1).Build()
	if err != nil {
		return nil, err
	}
	response, err := cli.Query(ctx, &dynamodb.QueryInput{
		TableName:                 aws.String(endpointPublishTable()),
		IndexName:                 aws.String(endpointPublishIndexDataappEndpoint),
		KeyConditionExpression:    expr.KeyCondition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
	})
	if err != nil {
		return nil, err
	}
	endpoints := []*entity.EndpointPublish{}
	err = attributevalue.UnmarshalListOfMaps(response.Items, &endpoints)
	if err != nil {
		return nil, err
	}
	return endpoints, nil
}

func GetPublishedEndpointByName(ctx context.Context, cli *dynamodb.Client, dataAppID uint64, endpoint string) ([]*entity.EndpointPublish, error) {
	conn1 := expression.Key("dataapp_id").Equal(expression.Value(dataAppID))
	conn2 := expression.Key("endpoint").Equal(expression.Value(endpoint))
	expr, err := expression.NewBuilder().WithKeyCondition(conn1.And(conn2)).Build()
	if err != nil {
		return nil, err
	}
	response, err := cli.Query(ctx, &dynamodb.QueryInput{
		TableName:                 aws.String(endpointPublishTable()),
		IndexName:                 aws.String(endpointPublishIndexDataappEndpoint),
		KeyConditionExpression:    expr.KeyCondition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
	})
	if err != nil {
		return nil, err
	}
	endpoints := []*entity.EndpointPublish{}
	err = attributevalue.UnmarshalListOfMaps(response.Items, &endpoints)
	if err != nil {
		return nil, err
	}
	return endpoints, nil
}

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

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

func DeleteEndpoint(ctx context.Context, cli *dynamodb.Client, endpointId uint64) error {
	key, err := getEndpointDdbKey(endpointId)
	if err != nil {
		return err
	}
	return ddb.DeleteItem(ctx, cli, endpointsTable(), key)
}

func BatchDeleteEndpoint(ctx context.Context, cli *dynamodb.Client, endpointIds []uint64) (int, error) {
	keys := make([]map[string]types.AttributeValue, 0, len(endpointIds))
	for _, endpointId := range endpointIds {
		key, err := getEndpointDdbKey(endpointId)
		if err != nil {
			return 0, err
		}
		keys = append(keys, key)
	}
	return ddb.BatchDeleteItem(ctx, cli, endpointsTable(), keys...)
}

func DeletePublishedEndpoint(ctx context.Context, cli *dynamodb.Client, endpointId uint64) error {
	key, err := getPublishedEndpointDdbKey(endpointId)
	if err != nil {
		return err
	}
	return ddb.DeleteItem(ctx, cli, endpointPublishTable(), key)
}

func BatchDeletePublishedEndpoint(ctx context.Context, cli *dynamodb.Client, endpointIds []uint64) (int, error) {
	keys := make([]map[string]types.AttributeValue, 0, len(endpointIds))
	for _, endpointId := range endpointIds {
		key, err := getPublishedEndpointDdbKey(endpointId)
		if err != nil {
			return 0, err
		}
		keys = append(keys, key)
	}
	return ddb.BatchDeleteItem(ctx, cli, endpointPublishTable(), keys...)
}

func DeleteEndpointByAppID(ctx context.Context, cli *dynamodb.Client, dataappID uint64) error {
	endpointList, err := QueryEndpointByAppID(ctx, cli, dataappID)
	if err != nil {
		return tierrs.NewInternalDB("Failed to query endpoint by appid", err)
	} else if len(endpointList) == 0 {
		return nil
	}

	endpointIDs := make([]uint64, 0, len(endpointList))
	for _, endpoint := range endpointList {
		endpointIDs = append(endpointIDs, endpoint.ID)
	}
	_, err = BatchDeleteEndpoint(ctx, cli, endpointIDs)
	if err != nil {
		return tierrs.NewInternalDB("Failed to batch delete endpoint", err)
	}
	return nil
}

func DeletePublishedEndpointByAppID(ctx context.Context, cli *dynamodb.Client, dataappID uint64) error {
	endpointList, err := QueryPublishedEndpointByAppID(ctx, cli, dataappID)
	if err != nil {
		return tierrs.NewInternalDB("Failed to query deployed endpoint by appid", err)
	} else if len(endpointList) == 0 {
		return nil
	}

	endpointIDs := make([]uint64, 0, len(endpointList))
	for _, endpoint := range endpointList {
		endpointIDs = append(endpointIDs, endpoint.EndpointId)
	}
	_, err = BatchDeletePublishedEndpoint(ctx, cli, endpointIDs)
	if err != nil {
		return tierrs.NewInternalDB("Failed to batch delete deployed endpoint", err)
	}
	return nil
}

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

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

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

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

func getEndpointDdbKey(id uint64) (map[string]types.AttributeValue, error) {
	idValue, err := attributevalue.Marshal(id)
	if err != nil {
		return nil, err
	}
	return map[string]types.AttributeValue{"id": idValue}, nil
}

func getPublishedEndpointDdbKey(endpointID uint64) (map[string]types.AttributeValue, error) {
	idValue, err := attributevalue.Marshal(endpointID)
	if err != nil {
		return nil, err
	}
	return map[string]types.AttributeValue{"endpoint_id": idValue}, nil
}