package registration

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"

	"code.justin.tv/sse/malachai/pkg/rorolepicker"
	"code.justin.tv/sse/policy"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"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/expression"
	"github.com/aws/aws-sdk-go/service/iam"
)

// SyncReadOnlyRole syncs a specific read-only role, referenced by role name
func (reg *registrar) SyncReadOnlyRole(ctx context.Context, readOnlyRoleName string) error {
	services, err := reg.queryServicesWithReadOnlyRole(ctx, readOnlyRoleName)
	if err != nil {
		return err
	}

	ids := make([]string, 1, 200)
	// always whitelist the s2s account
	ids[0] = reg.res.MalachaiAccountID

	for _, svc := range services {
		accountIDs, err := svc.allowedAccountIDs()
		if err != nil {
			return err
		}

		ids = append(ids, accountIDs...)
	}

	ids = uniqueAccountIDs(ids)

	err = reg.ensureRole(ctx, readOnlyRoleName, ids)
	return err
}

// modifies accountIDs and reuses the underlying array
func uniqueAccountIDs(accountIDs []string) []string {
	m := map[string]struct{}{}
	for _, id := range accountIDs {
		m[id] = struct{}{}
	}
	out := accountIDs[:0]
	for k := range m {
		out = append(out, k)
	}
	return out
}

func (reg *registrar) queryServicesWithReadOnlyRole(ctx context.Context, roleName string) (services []*Service, err error) {
	if roleName == "" {
		return nil, errors.New("read only role name is required")
	}

	proj := expression.NamesList(expression.Name("service_id"), expression.Name("service_name"), expression.Name("allowed_arns"))
	keyCond := expression.KeyEqual(expression.Key("ro_role_shard_name"), expression.Value(roleName))
	expression, err := expression.NewBuilder().WithProjection(proj).WithKeyCondition(keyCond).Build()
	if err != nil {
		return nil, err
	}
	var pageErr error
	err = reg.db.QueryPagesWithContext(ctx, &dynamodb.QueryInput{
		ProjectionExpression:      expression.Projection(),
		TableName:                 aws.String(reg.res.ServicesTable),
		KeyConditionExpression:    expression.KeyCondition(),
		IndexName:                 aws.String("services-by-ro-role"),
		ExpressionAttributeNames:  expression.Names(),
		ExpressionAttributeValues: expression.Values(),
	}, func(out *dynamodb.QueryOutput, more bool) bool {
		if services == nil {
			services = make([]*Service, 0, aws.Int64Value(out.ScannedCount))
		}
		for _, item := range out.Items {
			svc := &Service{}
			pageErr = dynamodbattribute.UnmarshalMap(item, svc)
			if pageErr != nil {
				return false
			}
			svc.AllowedArns, pageErr = reg.ig.GetValidAllowedARNs(svc.ID, svc.AllowedArns)
			if pageErr != nil {
				return false
			}
			services = append(services, svc)
		}
		return len(out.LastEvaluatedKey) != 0
	})
	if err != nil {
		return nil, err
	}
	if pageErr != nil {
		return nil, pageErr
	}

	return services, nil
}

func (reg *registrar) createRole(ctx context.Context, roleName string, accountIDs []string) error {
	arp, err := reg.prepareAssumeRolePolicy(ctx, accountIDs)
	if err != nil {
		return err
	}

	_, err = reg.iam.CreateRoleWithContext(ctx, &iam.CreateRoleInput{
		AssumeRolePolicyDocument: aws.String(arp),
		Path:                     aws.String(rorolepicker.RolePathPrefix),
		Description:              aws.String("sharded role providing read-only permissions for the s2s service registration table"),
		RoleName:                 aws.String(roleName),
	})
	return err
}

func (reg *registrar) checkIfRoleExists(ctx context.Context, roleName string) (bool, error) {
	_, err := reg.iam.GetRoleWithContext(ctx, &iam.GetRoleInput{
		RoleName: aws.String(roleName),
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == iam.ErrCodeNoSuchEntityException {
				return false, nil
			}
		}
		return false, err
	}
	return true, nil
}

func (reg *registrar) ensureAssumeRolePolicy(ctx context.Context, roleName string, accountIDs []string) error {
	arp, err := reg.prepareAssumeRolePolicy(ctx, accountIDs)
	if err != nil {
		return err
	}

	_, err = reg.iam.UpdateAssumeRolePolicyWithContext(ctx, &iam.UpdateAssumeRolePolicyInput{
		RoleName:       aws.String(roleName),
		PolicyDocument: aws.String(arp),
	})
	return err
}

func (reg *registrar) prepareAssumeRolePolicy(ctx context.Context, accountIDs []string) (string, error) {
	arns := make([]string, len(accountIDs))
	for p, id := range accountIDs {
		arns[p] = fmt.Sprintf("arn:aws:iam::%s:root", id)
	}

	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(policy.IAMAssumeRolePolicy{
		Version: policy.DocumentVersion,
		Statement: []policy.IAMAssumeRoleStatement{
			{
				Sid:    "AccountIDWhitelist",
				Effect: "Allow",
				Action: "sts:AssumeRole",
				Principal: map[string]interface{}{
					"AWS": arns,
				},
			},
		},
	})
	return buf.String(), err
}

func (reg *registrar) ensureRoleInlinePolicy(ctx context.Context, roleName string) error {
	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(policy.IAMPolicyDocument{
		Version: policy.DocumentVersion,
		Statement: []policy.IAMPolicyStatement{
			{
				Sid:    "ServicesReadOnly",
				Effect: "Allow",
				Action: []string{
					"dynamodb:GetItem",
					"dynamodb:Query",
					"dynamodb:Scan",
				},
				Resource: []string{
					reg.res.ServicesTableArn,
					reg.res.ServicesTableArn + "/index/services-by-id",
					reg.res.InventoryTableArn,
					reg.res.InventoryTableArn + "/index/instance-id",
				},
			},
		},
	})
	if err != nil {
		return err
	}

	_, err = reg.iam.PutRolePolicyWithContext(ctx, &iam.PutRolePolicyInput{
		PolicyDocument: aws.String(buf.String()),
		PolicyName:     aws.String("ServicesReadOnly"),
		RoleName:       aws.String(roleName),
	})
	return err
}

func (reg *registrar) ensureRole(ctx context.Context, roleName string, accountIDs []string) error {
	exists, err := reg.checkIfRoleExists(ctx, roleName)
	if err != nil {
		return err
	}
	if !exists {
		err = reg.createRole(ctx, roleName, accountIDs)
		if err != nil {
			return err
		}
	}

	err = reg.ensureRoleInlinePolicy(ctx, roleName)
	if err != nil {
		return err
	}

	return reg.ensureAssumeRolePolicy(ctx, roleName, accountIDs)
}
