package ddb

import (
	"context"

	"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"
)

// DdbPrivilegeRole is the struct of privilege_role table in DynamoDB
// The primary key of this table is resource_uri(partition key) and role_id(sort key)
// Example:
//
//	{
//	  "resource_uri": "https://tidbcloud.com/api/v1/organizations/1",
//	  "role_id": 1,
//	  "privilege_id": 1
//	}
type DdbPrivilegeRole struct {
	ResourceUri string `dynamodbav:"resource_uri"`
	RoleID      uint64 `dynamodbav:"role_id"`
	PrivilegeID uint64 `dynamodbav:"privilege_id"`
}

// GetKey returns the composite primary key of the ddbPrivilegeRole in a format that can be
// sent to DynamoDB.
func (pr *DdbPrivilegeRole) GetKey(ctx context.Context) (map[string]types.AttributeValue, error) {
	resourceUri, err := attributevalue.Marshal(pr.ResourceUri)
	if err != nil {
		tilog.Error(ctx, "failed to marshal resourceUri", zap.Error(err))
		return nil, err
	}
	roleID, err := attributevalue.Marshal(pr.RoleID)
	if err != nil {
		tilog.Error(ctx, "failed to marshal roleID", zap.Error(err))
		return nil, err
	}
	return map[string]types.AttributeValue{"resource_uri": resourceUri, "role_id": roleID}, nil
}

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

	pr := DdbPrivilegeRole{
		ResourceUri: resourceUri,
		RoleID:      roleID,
	}
	itemKey, err := pr.GetKey(ctx)
	if err != nil {
		return err
	}
	_, err = ddbClient.DeleteItem(ctx, &dynamodb.DeleteItemInput{
		TableName: aws.String(ddbClient.PrivilegeRoleTableName),
		Key:       itemKey,
	})
	if err != nil {
		tilog.Error(ctx, "failed to delete RolePrivilege in dynamodb", zap.Error(err))
	}
	return err
}

func CreateRolePrivilege(ctx context.Context, ddbClient *DynamoDBClient, resourceUri string, roleID uint64, privilegeID uint64) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "CreateRolePrivilege"))
		return nil
	}

	pr := DdbPrivilegeRole{
		ResourceUri: resourceUri,
		RoleID:      roleID,
		PrivilegeID: privilegeID,
	}
	item, err := attributevalue.MarshalMap(pr)
	if err != nil {
		return err
	}
	_, err = ddbClient.PutItem(ctx, &dynamodb.PutItemInput{
		TableName: aws.String(ddbClient.PrivilegeRoleTableName),
		Item:      item,
	})
	if err != nil {
		tilog.Error(ctx, "failed to create RolePrivilege in dynamodb", zap.Error(err))
	}
	return err
}

func transactWriteItems(ctx context.Context, ddbClient *DynamoDBClient, items []types.TransactWriteItem) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "transactWriteItems"))
		return nil
	}

	batchSize := 100
	for i := 0; i < len(items); i += batchSize {
		end := i + batchSize
		if end > len(items) {
			end = len(items)
		}
		_, err := ddbClient.TransactWriteItems(ctx, &dynamodb.TransactWriteItemsInput{
			TransactItems: items[i:end],
		})
		if err != nil {
			tilog.Error(ctx, "failed to transactWriteItems", zap.Error(err))
			return err
		}
	}
	return nil
}

// CreateRolePrivilegesByRoleID creates RbacRolePrivilege by roleID and privileges to dynamodb using TransactWriteItems
func CreateRolePrivilegesByRoleID(ctx context.Context, ddbClient *DynamoDBClient,
	roleID uint64, privileges []entity.RbacPrivilege,
) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "CreateRolePrivilegesByRoleID"))
		return nil
	}

	var items []types.TransactWriteItem
	for _, privilege := range privileges {
		item, err := attributevalue.MarshalMap(DdbPrivilegeRole{
			ResourceUri: privilege.ResourceUri,
			RoleID:      roleID,
			PrivilegeID: privilege.ID.Uint64(),
		})
		if err != nil {
			return err
		}
		items = append(items, types.TransactWriteItem{
			Put: &types.Put{
				TableName: aws.String(ddbClient.PrivilegeRoleTableName),
				Item:      item,
			},
		})
	}

	return transactWriteItems(ctx, ddbClient, items)
}

func DeleteRolePrivileges(ctx context.Context, ddbClient *DynamoDBClient, ddbPrivilegeRoles []DdbPrivilegeRole) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "DeleteRolePrivileges"))
		return nil
	}

	var items []types.TransactWriteItem
	for _, ddbPrivilegeRole := range ddbPrivilegeRoles {
		key, err := ddbPrivilegeRole.GetKey(ctx)
		if err != nil {
			return err
		}
		items = append(items, types.TransactWriteItem{
			Delete: &types.Delete{
				TableName: aws.String(ddbClient.PrivilegeRoleTableName),
				Key:       key,
			},
		})
	}
	return transactWriteItems(ctx, ddbClient, items)
}

// DeleteRolePrivilegesByRoleID deletes RbacRolePrivilege by roleID and privileges to dynamodb using TransactWriteItems
func DeleteRolePrivilegesByRoleID(ctx context.Context, ddbClient *DynamoDBClient,
	roleID uint64, privileges []entity.RbacPrivilege,
) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "DeleteRolePrivilegesByRoleID"))
		return nil
	}

	var items []types.TransactWriteItem
	for _, privilege := range privileges {
		pr := DdbPrivilegeRole{
			ResourceUri: privilege.ResourceUri,
			RoleID:      roleID,
		}
		key, err := pr.GetKey(ctx)
		if err != nil {
			return err
		}
		items = append(items, types.TransactWriteItem{
			Delete: &types.Delete{
				TableName: aws.String(ddbClient.PrivilegeRoleTableName),
				Key:       key,
			},
		})
	}
	return transactWriteItems(ctx, ddbClient, items)
}

// DeleteRolePrivilegesByResourceUri deletes RbacRolePrivilege by resourceUri and roleIDs to dynamodb using TransactWriteItems
func DeleteRolePrivilegesByResourceUri(ctx context.Context, ddbClient *DynamoDBClient,
	resourceUri string, roleIDs []uint64,
) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "DeleteRolePrivilegesByResourceUri"))
		return nil
	}

	var items []types.TransactWriteItem
	for _, roleID := range roleIDs {
		pr := DdbPrivilegeRole{
			ResourceUri: resourceUri,
			RoleID:      roleID,
		}
		key, err := pr.GetKey(ctx)
		if err != nil {
			return err
		}
		items = append(items, types.TransactWriteItem{
			Delete: &types.Delete{
				TableName: aws.String(ddbClient.PrivilegeRoleTableName),
				Key:       key,
			},
		})
	}
	return transactWriteItems(ctx, ddbClient, items)
}

// UpdateRbacRolePrivileges updates RbacRolePrivilege to dynamodb
// delete old rolePrivileges, create new rolePrivileges
func UpdateRbacRolePrivilegesByResourceUri(ctx context.Context, ddbClient *DynamoDBClient,
	oldResourceUri string, newResourceUri string, privilegeID uint64, roleIDs []uint64,
) error {
	if ddbClient.DisableWrite {
		tilog.Info(ctx, "dynamodb write operation is disabled", zap.String("operation", "UpdateRbacRolePrivilegesByResourceUri"))
		return nil
	}

	var deleteItems []types.TransactWriteItem
	var createItems []types.TransactWriteItem
	for _, roleID := range roleIDs {
		oldpr := DdbPrivilegeRole{
			ResourceUri: oldResourceUri,
			RoleID:      roleID,
		}
		key, err := oldpr.GetKey(ctx)
		if err != nil {
			return err
		}
		deleteItems = append(deleteItems, types.TransactWriteItem{
			Delete: &types.Delete{
				TableName: aws.String(ddbClient.PrivilegeRoleTableName),
				Key:       key,
			},
		})
		item, err := attributevalue.MarshalMap(DdbPrivilegeRole{
			ResourceUri: newResourceUri,
			RoleID:      roleID,
			PrivilegeID: privilegeID,
		})
		if err != nil {
			return err
		}
		createItems = append(createItems, types.TransactWriteItem{
			Put: &types.Put{
				TableName: aws.String(ddbClient.PrivilegeRoleTableName),
				Item:      item,
			},
		})
	}
	writeItems := append(deleteItems, createItems...)
	return transactWriteItems(ctx, ddbClient, writeItems)
}

func GetRbacRolePrivilegesByRoleID(ctx context.Context, ddbClient *DynamoDBClient, roleID uint64) ([]DdbPrivilegeRole, error) {
	var ddbPrivilegeRoles []DdbPrivilegeRole

	keyEx := expression.Key("role_id").Equal(expression.Value(roleID))
	expr, err := expression.NewBuilder().WithKeyCondition(keyEx).Build()
	if err != nil {
		tilog.Error(ctx, "failed to build expression", zap.Error(err))
		return nil, err
	}
	output, err := ddbClient.Query(ctx, &dynamodb.QueryInput{
		TableName:                 aws.String(ddbClient.PrivilegeRoleTableName),
		IndexName:                 aws.String(ddbClient.RolePrivilegeGSIName),
		KeyConditionExpression:    expr.KeyCondition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
	})
	if err != nil {
		tilog.Error(ctx, "failed to query DdbPrivilegeRole in dynamodb by roleID", zap.Error(err))
		return nil, err
	}
	err = attributevalue.UnmarshalListOfMaps(output.Items, &ddbPrivilegeRoles)
	if err != nil {
		tilog.Error(ctx, "failed to unmarshal DdbPrivilegeRole", zap.Error(err))
		return nil, err
	}
	return ddbPrivilegeRoles, nil
}