package role

import (
	"context"
	"errors"
	"fmt"
	"sort"

	"code.justin.tv/beefcake/server/internal/awsutil"
	"code.justin.tv/beefcake/server/internal/config"
	"code.justin.tv/beefcake/server/internal/perm"
	"code.justin.tv/beefcake/server/internal/stringutil"
	"code.justin.tv/beefcake/server/rpc/beefcake"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/aws/aws-sdk-go/service/dynamodb/expression"
	"github.com/gofrs/uuid"
)

// ErrCannotAddLegacyPermission returned when a user is adding a legacy
// permission using AddPermission
var (
	ErrCannotAddLegacyPermission = errors.New("cannot add legacy permission using AddPermission")
)

// ErrLegacyPermissionDoesNotExist is return when attaching a permission before
// it exists.
type ErrLegacyPermissionDoesNotExist struct {
	LegacyPermissionID string
}

func (e ErrLegacyPermissionDoesNotExist) Error() string {
	return fmt.Sprintf("legacy permission does not exist: %s", e.LegacyPermissionID)
}

// ErrRoleDoesNotExist is returned when an action on a non existant role is
// performed
type ErrRoleDoesNotExist struct {
	RoleID string
}

func (e ErrRoleDoesNotExist) Error() string {
	return fmt.Sprintf("role does not exist: %s", e.RoleID)
}

// ErrMissingRoleParameters is returned when parameters are missing
type ErrMissingRoleParameters struct {
	Parameters []string
}

func (e ErrMissingRoleParameters) Error() string {
	return fmt.Sprintf("missing role parameters: %v", e.Parameters)
}

// DynamoDB Attributes
const (
	IDAttribute                = "ID"
	NameAttribute              = "Name"
	LegacyPermissionsAttribute = "LegacyPermissions"
	PermissionsAttribute       = "Permissions"
	UserMembershipsAttribute   = "UserMemberships"
)

// Roles ...
type Roles struct {
	Config   *config.Config
	DynamoDB dynamodbiface.DynamoDBAPI
}

// All returns all roles
func (c *Roles) All(ctx context.Context) ([]*Role, error) {
	var innerErr error
	roles := make([]*Role, 0, 100)
	if err := c.DynamoDB.ScanPagesWithContext(ctx,
		&dynamodb.ScanInput{
			Limit:     aws.Int64(100),
			TableName: aws.String(c.Config.RolesTableName.Get()),
		},
		func(out *dynamodb.ScanOutput, last bool) (cont bool) {
			for _, item := range out.Items {
				var role Role
				if innerErr = dynamodbattribute.UnmarshalMap(item, &role); innerErr != nil {
					return false
				}
				roles = append(roles, &role)
			}
			return !last
		}); err != nil {
		return nil, err
	}

	if innerErr != nil {
		return nil, innerErr
	}

	sort.Sort(sortableRoles(roles))

	return roles, nil
}

// Create a role
func (c *Roles) Create(ctx context.Context, roleName string) (*Role, error) {
	if len(roleName) == 0 {
		return nil, ErrMissingRoleParameters{Parameters: []string{"Name"}}
	}

	roleID, err := uuid.NewV4()
	if err != nil {
		return nil, err
	}

	role := &Role{ID: roleID.String(), Name: roleName}

	item, err := dynamodbattribute.MarshalMap(role)
	if err != nil {
		return nil, err
	}

	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(c.Config.RolesHashKey.Get()).AttributeNotExists()).
		Build()
	if err != nil {
		return nil, err
	}

	item[PermissionsAttribute] = &dynamodb.AttributeValue{
		M: map[string]*dynamodb.AttributeValue{},
	}
	item[UserMembershipsAttribute] = &dynamodb.AttributeValue{
		M: map[string]*dynamodb.AttributeValue{},
	}

	_, err = c.DynamoDB.PutItemWithContext(ctx, &dynamodb.PutItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Item:                      item,
		TableName:                 aws.String(c.Config.RolesTableName.Get()),
	})
	if err != nil {
		return nil, err
	}

	return role, nil
}

// Get a role with its permissions
func (c *Roles) Get(ctx context.Context, roleID string) (*Role, error) {
	res, err := c.DynamoDB.GetItemWithContext(ctx, &dynamodb.GetItemInput{
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName: aws.String(c.Config.RolesTableName.Get()),
	})
	if err != nil {
		return nil, err
	}

	r := new(Role)
	if err := dynamodbattribute.UnmarshalMap(res.Item, r); err != nil {
		return nil, err
	}

	return r, nil
}

// AddPermission adds a permission on a role.
func (c *Roles) AddPermission(
	ctx context.Context,
	roleID string,
	permission *beefcake.Permission,
) error {
	if permission.GetLegacy() != nil {
		return ErrCannotAddLegacyPermission
	}

	rolePermID, err := uuid.NewV4()
	if err != nil {
		return err
	}

	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(c.Config.RolesHashKey.Get()).AttributeExists()).
		WithUpdate(expression.Set(
			expression.Name(c.rolePermissionPath(rolePermID.String())),
			expression.Value(perm.Permission(*permission)),
		)).
		Build()
	if err != nil {
		return err
	}

	_, err = c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	})
	return err
}

// UpdateOptions ...
type UpdateOptions struct {
	Name                *string
	LegacyPermissionIDs *[]string
}

// Update a role. All values are expected to be set.
func (c *Roles) Update(ctx context.Context, roleID string, options UpdateOptions) error {
	if options.Name == nil && options.LegacyPermissionIDs == nil {
		return nil
	}

	var update expression.UpdateBuilder
	if options.Name != nil {
		update = update.Set(expression.Name(NameAttribute), expression.Value(*options.Name))
	}
	if options.LegacyPermissionIDs != nil {
		legacyPermissionIDs := *options.LegacyPermissionIDs
		if len(legacyPermissionIDs) > 0 {
			update = update.Set(
				expression.Name(LegacyPermissionsAttribute),
				expression.Value(stringutil.NewSet(legacyPermissionIDs...)),
			)
		} else {
			update = update.Remove(expression.Name(LegacyPermissionsAttribute))
		}
	}

	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(c.Config.RolesHashKey.Get()).AttributeExists()).
		WithUpdate(update).
		Build()
	if err != nil {
		return err
	}

	if _, err := c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	}); err != nil {
		if awsutil.IsStatusCode(err, dynamodb.ErrCodeConditionalCheckFailedException) {
			return ErrRoleDoesNotExist{RoleID: roleID}
		}
		return err
	}

	return nil
}

// AddLegacyPermission adds a legacy permission on a role.
func (c *Roles) AddLegacyPermission(
	ctx context.Context,
	roleID string,
	legacyPermissionID string,
) error {
	legacyPermissionExistsCheck, err := c.legacyPermissionExistsCheck(legacyPermissionID)
	if err != nil {
		return err
	}

	update, err := c.updateLegacyPermissions(roleID, []string{legacyPermissionID})
	if err != nil {
		return err
	}

	set, err := c.setLegacyPermissions(roleID, []string{legacyPermissionID})
	if err != nil {
		return err
	}

	if _, err := c.DynamoDB.TransactWriteItemsWithContext(ctx, &dynamodb.TransactWriteItemsInput{
		TransactItems: []*dynamodb.TransactWriteItem{
			legacyPermissionExistsCheck,
			set,
		},
	}); err != nil {
		if awsutil.IsStatusCode(err, dynamodb.ErrCodeTransactionCanceledException) {
			_, err := c.DynamoDB.TransactWriteItemsWithContext(ctx, &dynamodb.TransactWriteItemsInput{
				TransactItems: []*dynamodb.TransactWriteItem{
					legacyPermissionExistsCheck,
					update,
				},
			})

			if awsutil.IsStatusCode(err, dynamodb.ErrCodeTransactionCanceledException) {
				return ErrLegacyPermissionDoesNotExist{legacyPermissionID}
			}

			return err
		}

		return err
	}
	return nil
}

// RemoveLegacyPermission adds a legacy permission on a role.
func (c *Roles) RemoveLegacyPermission(
	ctx context.Context,
	roleID string,
	legacyPermissionID string,
) error {
	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(LegacyPermissionsAttribute).AttributeExists()).
		WithUpdate(expression.Delete(
			expression.Name(LegacyPermissionsAttribute),
			expression.Value(stringutil.NewSet(legacyPermissionID)),
		)).
		Build()
	if err != nil {
		return err
	}

	if _, err := c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	}); err != nil {
		if awsutil.IsStatusCode(err, dynamodb.ErrCodeConditionalCheckFailedException) {
			return nil
		}
		return err
	}
	return err
}

// AddLegacyPermissionDetails adds a legacy permission details. Used for
// syncing.
func (c *Roles) AddLegacyPermissionDetails(
	ctx context.Context,
	roleID string,
	details *beefcake.Permission_LegacyPermission,
) error {
	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(c.Config.RolesHashKey.Get()).AttributeExists()).
		WithUpdate(expression.Set(
			expression.Name(c.rolePermissionPath(c.legacyPermissionID(details.Id))),
			expression.Value(perm.Permission(beefcake.Permission{
				Value: &beefcake.Permission_Legacy{Legacy: details},
			})),
		)).
		Build()
	if err != nil {
		return err
	}

	if _, err := c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	}); err != nil {
		if awsutil.IsStatusCode(err, dynamodb.ErrCodeConditionalCheckFailedException) {
			return nil
		}
		return err
	}

	return nil
}

// RemoveLegacyPermissionDetails removes a legacy permission from the
// permissions map. Should only be called from sync lambda.
func (c *Roles) RemoveLegacyPermissionDetails(
	ctx context.Context,
	roleID string,
	legacyPermissionID string,
) error {

	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(c.Config.RolesHashKey.Get()).AttributeExists()).
		WithUpdate(expression.Remove(expression.Name(
			c.rolePermissionPath(c.legacyPermissionID(legacyPermissionID)),
		))).
		Build()
	if err != nil {
		return err
	}

	if _, err := c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	}); err != nil {
		if awsutil.IsStatusCode(err, dynamodb.ErrCodeConditionalCheckFailedException) {
			return nil
		}
		return err
	}
	return nil
}

func (c *Roles) setLegacyPermissions(roleID string, legacyPermissionIDs []string) (*dynamodb.TransactWriteItem, error) {
	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(LegacyPermissionsAttribute).AttributeNotExists()).
		WithUpdate(expression.Set(
			expression.Name(LegacyPermissionsAttribute),
			expression.Value(stringutil.NewSet(legacyPermissionIDs...)),
		)).
		Build()
	if err != nil {
		return nil, err
	}

	return &dynamodb.TransactWriteItem{
		Update: &dynamodb.Update{
			ConditionExpression:       expr.Condition(),
			ExpressionAttributeNames:  expr.Names(),
			ExpressionAttributeValues: expr.Values(),
			Key: map[string]*dynamodb.AttributeValue{
				c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
			},
			TableName:        aws.String(c.Config.RolesTableName.Get()),
			UpdateExpression: expr.Update(),
		},
	}, nil
}

func (c *Roles) updateLegacyPermissions(roleID string, legacyPermissionIDs []string) (*dynamodb.TransactWriteItem, error) {
	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(LegacyPermissionsAttribute).AttributeExists()).
		WithUpdate(expression.Add(
			expression.Name(LegacyPermissionsAttribute),
			expression.Value(stringutil.NewSet(legacyPermissionIDs...)),
		)).
		Build()
	if err != nil {
		return nil, err
	}

	return &dynamodb.TransactWriteItem{
		Update: &dynamodb.Update{
			ConditionExpression:       expr.Condition(),
			ExpressionAttributeNames:  expr.Names(),
			ExpressionAttributeValues: expr.Values(),
			Key: map[string]*dynamodb.AttributeValue{
				c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
			},
			TableName:        aws.String(c.Config.RolesTableName.Get()),
			UpdateExpression: expr.Update(),
		},
	}, nil
}

func (c *Roles) legacyPermissionExistsCheck(legacyPermissionID string) (*dynamodb.TransactWriteItem, error) {
	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(c.Config.LegacyPermissionsHashKey.Get()).AttributeExists()).
		Build()
	if err != nil {
		return nil, err
	}

	return &dynamodb.TransactWriteItem{
		ConditionCheck: &dynamodb.ConditionCheck{
			ConditionExpression:       expr.Condition(),
			ExpressionAttributeNames:  expr.Names(),
			ExpressionAttributeValues: expr.Values(),
			Key: map[string]*dynamodb.AttributeValue{
				c.Config.LegacyPermissionsHashKey.Get(): {S: aws.String(legacyPermissionID)},
			},
			TableName: aws.String(c.Config.LegacyPermissionsTableName.Get()),
		},
	}, nil
}

// Override sets role attributes, overriding existing ones if not existing.
func (c *Roles) Override(
	ctx context.Context,
	roleID string,
	name string,
	permissions []*beefcake.Permission,
	memberships []*beefcake.Role_UserMembership,
) error {
	legacyPerms := stringutil.NewSet()
	aps := make([]*beefcake.AttachedPermission, 0, len(permissions))
	for _, p := range permissions {
		if legacyPerm := p.GetLegacy(); legacyPerm != nil {
			legacyPerms.Add(legacyPerm.GetId())
		} else {
			rolePermID, err := uuid.NewV4()
			if err != nil {
				return err
			}

			aps = append(aps, &beefcake.AttachedPermission{
				Id:         rolePermID.String(),
				Permission: p,
			})
		}
	}

	update := expression.
		Set(
			expression.Name(NameAttribute),
			expression.Value(name),
		).
		Set(
			expression.Name(PermissionsAttribute),
			expression.Value(perm.AttachedPermissions(aps)),
		).
		Set(
			expression.Name(UserMembershipsAttribute),
			expression.Value(UserMemberships(memberships)),
		)

	if len(legacyPerms) > 0 {
		update = update.Set(expression.Name(LegacyPermissionsAttribute), expression.Value(legacyPerms))
	}

	expr, err := expression.NewBuilder().
		WithUpdate(update).
		Build()
	if err != nil {
		return err
	}

	_, err = c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	})
	return err
}

// RemovePermission removes a permission on a role.
func (c *Roles) RemovePermission(
	ctx context.Context,
	roleID,
	rolePermissionID string,
) error {
	if err := c.removePermissionAndLegacyAttachment(ctx, roleID, rolePermissionID); err != nil {
		if !awsutil.IsStatusCode(err, dynamodb.ErrCodeConditionalCheckFailedException) {
			return err
		}
	} else {
		return nil
	}

	return c.removePermissionOnly(ctx, roleID, rolePermissionID)
}

func (c *Roles) removePermissionAndLegacyAttachment(
	ctx context.Context,
	roleID,
	rolePermissionID string,
) error {
	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(LegacyPermissionsAttribute).AttributeExists()).
		WithUpdate(expression.
			Remove(expression.Name(c.rolePermissionPath(rolePermissionID))).
			Delete(
				expression.Name(LegacyPermissionsAttribute),
				expression.Value(stringutil.NewSet(rolePermissionID)),
			),
		).
		Build()
	if err != nil {
		return err
	}

	_, err = c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	})
	return err
}

func (c *Roles) removePermissionOnly(
	ctx context.Context,
	roleID,
	rolePermissionID string,
) error {
	expr, err := expression.NewBuilder().
		WithCondition(expression.Name(LegacyPermissionsAttribute).AttributeNotExists()).
		WithUpdate(expression.
			Remove(expression.Name(c.rolePermissionPath(rolePermissionID)))).
		Build()
	if err != nil {
		return err
	}

	_, err = c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ConditionExpression:       expr.Condition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	})
	return err
}

// AddUserMembership adds a user to a role
func (c *Roles) AddUserMembership(ctx context.Context, roleID string, um *beefcake.Role_UserMembership) error {
	// Add to string set
	expr, err := expression.NewBuilder().
		WithUpdate(expression.Set(
			expression.Name(c.userMembershipPath(um.GetUserId())),
			expression.Value(UserMembership(*um)),
		)).
		Build()
	if err != nil {
		return err
	}

	_, err = c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	})
	return err
}

// RemoveUserMembership removes a user to a role
func (c *Roles) RemoveUserMembership(ctx context.Context, roleID, userID string) error {
	// Add to string set
	expr, err := expression.NewBuilder().
		WithUpdate(expression.Remove(
			expression.Name(c.userMembershipPath(userID)),
		)).
		Build()
	if err != nil {
		return err
	}

	_, err = c.DynamoDB.UpdateItemWithContext(ctx, &dynamodb.UpdateItemInput{
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName:        aws.String(c.Config.RolesTableName.Get()),
		UpdateExpression: expr.Update(),
	})
	return err
}

// Delete a role
func (c *Roles) Delete(ctx context.Context, roleID string) error {
	_, err := c.DynamoDB.DeleteItemWithContext(ctx, &dynamodb.DeleteItemInput{
		Key: map[string]*dynamodb.AttributeValue{
			c.Config.RolesHashKey.Get(): {S: aws.String(roleID)},
		},
		TableName: aws.String(c.Config.RolesTableName.Get()),
	})
	return err
}

// rolePermissionPath returns the dynamodb path to a role permission id
func (c *Roles) rolePermissionPath(rolePermissionID string) string {
	return PermissionsAttribute + "." + rolePermissionID
}

func (c *Roles) legacyPermissionID(id string) string {
	return "legacy:" + id
}

// rolePermissionPath returns the dynamodb path to a role permission id
func (c *Roles) userMembershipPath(userID string) string {
	return UserMembershipsAttribute + "." + userID
}
