package iamgenerator

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"path"

	"code.justin.tv/sse/malachai/pkg/config"
	"code.justin.tv/sse/policy"
	"code.justin.tv/systems/sandstorm/manager"
	spg "code.justin.tv/systems/sandstorm/policy"
	"code.justin.tv/systems/sandstorm/queue"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/iam/iamiface"
)

const (
	ownerTeamSSE  = "team-sse"
	iamPathPrefix = "malachai/registration/"
	// version is encoded in the role and policy path prefixes
	version = "v1"
)

// IAMGenerator describes the registration iam role generator interface
type IAMGenerator interface {
	GetValidAllowedARNs(roleName string, expected []string) ([]string, error)
	RegisterServiceRole(*RegisterServiceRoleRequest) (response *RegisterServiceRoleResponse, err error)
	DeregisterServiceRole(svc *Service) (err error)
	UpdateServiceRolePolicy(req *UpdateServiceRolePolicyRequest) (err error)
	UpdateAssumeRolePolicy(svc *Service, allowedArns []string) (err error)
}

type iamGenerator struct {
	spg         *spg.IAMPolicyGenerator
	iam         iamiface.IAMAPI
	cfg         *Config
	rootUserArn string
}

// New creates an IAMGenerator
func New(spg *spg.IAMPolicyGenerator, iam iamiface.IAMAPI, cfg *Config) (ig IAMGenerator, err error) {
	if err = cfg.fillDefaults(); err != nil {
		return
	}

	igImpl := &iamGenerator{
		spg:         spg,
		iam:         iam,
		cfg:         cfg,
		rootUserArn: cfg.RootUserArn,
	}
	igImpl.spg.SetPathPrefixEnvironment(cfg.Environment)

	ig = igImpl
	return
}

func generateAssumeRolePolicy(arns []string) (policyDocument string, err error) {
	var buf bytes.Buffer
	err = json.NewEncoder(&buf).Encode(policy.IAMAssumeRolePolicy{
		Version: policy.DocumentVersion,
		Statement: []policy.IAMAssumeRoleStatement{
			{
				Action: "sts:AssumeRole",
				Effect: "Allow",
				Principal: map[string]interface{}{
					"AWS": arns,
				},
			},
		},
	})
	if err != nil {
		return
	}

	policyDocument = buf.String()
	return
}

type roleRequest struct {
	ServiceID      string
	ServiceName    string
	PolicyArn      string
	AssumeRoleArns []string
}

func (req *roleRequest) validate() (err error) {
	switch "" {
	case req.ServiceID:
		err = errors.New("ServiceID is empty")
	case req.ServiceName:
		err = errors.New("ServiceName is empty")
	case req.PolicyArn:
		err = errors.New("PolicyArn is empty")
	}

	if req.AssumeRoleArns == nil {
		req.AssumeRoleArns = []string{}
	}
	return
}

// GetIAMPathPrefix returns the path prefix iam stuff
func GetIAMPathPrefix(environment, serviceName string) string {
	return "/" + path.Join(iamPathPrefix, environment, version, serviceName) + "/"
}

func (ig *iamGenerator) getIAMAuxWildcard() string {
	return fmt.Sprintf("arn:aws:iam::%s:role/malachai/%s/services-aux/*", ig.cfg.MalachaiAccountID, ig.cfg.Environment)
}

func (ig *iamGenerator) deregisterServiceRole(svc *Service) (err error) {
	roleName := svc.ID
	attachedPolicies, err := ig.listAttachedRolePolicies(svc)
	if err != nil {
		return
	}

	for _, attachedPolicy := range attachedPolicies.AttachedPolicies {

		_, err = ig.iam.DetachRolePolicy(&iam.DetachRolePolicyInput{
			PolicyArn: attachedPolicy.PolicyArn,
			RoleName:  aws.String(roleName),
		})
		if err != nil {
			return
		}

		err = ig.deletePolicy(aws.StringValue(attachedPolicy.PolicyArn))
		if err != nil {
			return
		}
	}

	_, err = ig.iam.DeleteRole(&iam.DeleteRoleInput{
		RoleName: aws.String(roleName),
	})
	return
}

// RegisterServiceRoleRequest contains parameters for creating a role
type RegisterServiceRoleRequest struct {
	Service        *Service
	AssumeRoleArns []string
}

//UpdateServiceRolePolicyRequest contains parameter for updating policy attached to
// an existing service role
type UpdateServiceRolePolicyRequest struct {
	Service *Service
}

// RegisterServiceRoleResponse is the return values for RegisterServiceRole
type RegisterServiceRoleResponse struct {
	RoleArn           string
	SandstormRoleArn  string
	SandstormRoleName string
}

func sandstormRoleName(environment, serviceID string) (roleName string) {
	return fmt.Sprintf("malachai-%s-%s", environment, serviceID)
}

func (ig *iamGenerator) buildServicePolicyRequest(svc *Service, sandstormRoleArn, env string) *servicePolicyRequest {
	return &servicePolicyRequest{
		ServiceID:                   svc.ID,
		ServiceName:                 svc.Name,
		CapabilitiesTopicArn:        svc.CapabilitiesTopicArn,
		SandstormRoleArn:            sandstormRoleArn,
		Environment:                 env,
		ServiceRegistrationTableArn: ig.cfg.ServiceRegistrationTableArn,
		CapabilitiesTableArn:        ig.cfg.CapabilitiesTableArn,
		DiscoveryTableArn:           ig.cfg.DiscoveryTableArn,
		PubkeyInvalidationTopicArn:  ig.cfg.PubkeyInvalidationTopicArn,
		MalachaiAccountID:           ig.cfg.MalachaiAccountID,
		InventoryTableArn:           ig.cfg.InventoryTableArn,
	}
}

// RegisterServiceRole creates an IAM policy and role for a service
func (ig *iamGenerator) RegisterServiceRole(req *RegisterServiceRoleRequest) (response *RegisterServiceRoleResponse, err error) {
	var sandstormRoleResponse *spg.CreateRoleResponse
	sandstormRoleResponse, err = ig.createServiceRole(req)
	if err != nil {
		return
	}
	policyRequest := ig.buildServicePolicyRequest(req.Service, sandstormRoleResponse.RoleArn, ig.cfg.Environment)
	policyArn, err := ig.registerServicePolicy(policyRequest)
	if err != nil {
		return
	}

	roleArn, err := ig.registerServiceRole(&roleRequest{
		ServiceID:      req.Service.ID,
		ServiceName:    req.Service.Name,
		PolicyArn:      policyArn,
		AssumeRoleArns: req.AssumeRoleArns,
	})
	if err != nil {
		return
	}

	response = &RegisterServiceRoleResponse{
		RoleArn:           roleArn,
		SandstormRoleArn:  sandstormRoleResponse.RoleArn,
		SandstormRoleName: sandstormRoleResponse.RoleName,
	}
	return
}

//UpdateService updates the existing malachai polies to the newer version of policy.
// This will attach the the new policy to the role before detaching all existing policies.
func (ig *iamGenerator) UpdateServiceRolePolicy(req *UpdateServiceRolePolicyRequest) (err error) {
	svc := req.Service

	attachedPolicies, err := ig.listAttachedRolePolicies(svc)
	if err != nil {
		return
	}

	var policyArn string
	for _, p := range attachedPolicies.AttachedPolicies {
		if aws.StringValue(p.PolicyName) == svc.ID {
			policyArn = aws.StringValue(p.PolicyArn)
		}
	}

	if policyArn == "" {
		return fmt.Errorf("no service policy not attached to servcice role: %s" + svc.ID)
	}
	//Attach the new policy
	policyRequest := ig.buildServicePolicyRequest(svc, svc.SandstormRoleArn, ig.cfg.Environment)
	err = ig.updatePolicy(policyArn, policyRequest)
	if err != nil {
		return
	}

	policyVersions, err := ig.iam.ListPolicyVersions(&iam.ListPolicyVersionsInput{
		PolicyArn: aws.String(policyArn),
	})

	for _, v := range policyVersions.Versions {
		if aws.BoolValue(v.IsDefaultVersion) {
			continue
		}
		_, err = ig.iam.DeletePolicyVersion(&iam.DeletePolicyVersionInput{
			PolicyArn: aws.String(policyArn),
			VersionId: v.VersionId,
		})
		if err != nil {
			return
		}
	}
	return
}

func (ig *iamGenerator) createServiceRole(req *RegisterServiceRoleRequest) (createRoleResponse *spg.CreateRoleResponse, err error) {
	sandstormAllowedArns := append([]string{ig.rootUserArn}, req.AssumeRoleArns...)

	roleRequest := &spg.CreateRoleRequest{
		Name:        sandstormRoleName(ig.cfg.Environment, req.Service.ID),
		SecretKeys:  []string{req.Service.SandstormSecretName},
		AllowedArns: sandstormAllowedArns,
		Owners:      []string{ownerTeamSSE},
	}

	createRoleResponse, err = ig.spg.Create(
		roleRequest,
		&noOpAuthorizer{},
		manager.New(
			manager.Config{
				AWSConfig:   config.AWSConfig(ig.cfg.AWSConfigBase, ig.cfg.Region, ig.cfg.SandstormRoleArn),
				Environment: ig.cfg.Environment,
				Queue: queue.Config{
					TopicArn: ig.cfg.SandstormTopicArn,
				},
			},
		),
	)
	return
}

type servicePolicyRequest struct {
	ServiceID                            string
	ServiceName                          string
	ServiceRegistrationTableArn          string
	SandstormRoleArn                     string
	Environment                          string `default:"production"`
	CapabilitiesTableArn                 string
	CapabilitiesTopicArn                 string
	ServiceRegistrationTableByIDIndexArn string
	DiscoveryTableArn                    string
	PubkeyInvalidationTopicArn           string
	MalachaiAccountID                    string
	InventoryTableArn                    string
}

func (req *servicePolicyRequest) validate() (err error) {
	switch "" {
	case req.Environment:
		err = errors.New("Environment is empty")
	case req.MalachaiAccountID:
		err = errors.New("MalachaiAccountID is empty")
	case req.SandstormRoleArn:
		err = errors.New("SandstormRoleArn is empty")
	case req.ServiceID:
		err = errors.New("ServiceID is empty")
	case req.ServiceName:
		err = errors.New("ServiceName is empty")
	case req.ServiceRegistrationTableArn:
		err = errors.New("ServiceRegistrationTableArn is empty")
	case req.CapabilitiesTableArn:
		err = errors.New("CapabilitiesTableArn is empty")
	case req.CapabilitiesTopicArn:
		err = errors.New("CapabilitiesTopicArn is empty")
	case req.DiscoveryTableArn:
		err = errors.New("DiscoveryTableArn is empty")
	case req.PubkeyInvalidationTopicArn:
		err = errors.New("PubkeyInvalidationTopicArn is empty")
	}
	return
}

func (ig *iamGenerator) generateServicePolicy(req *servicePolicyRequest) (policyDocument string, err error) {
	var buf bytes.Buffer
	err = json.NewEncoder(&buf).Encode(policy.IAMPolicyDocument{
		Version: policy.DocumentVersion,
		Statement: []policy.IAMPolicyStatement{
			{
				Sid:      "CapabilitiesTableRead",
				Resource: req.CapabilitiesTableArn,
				Effect:   "Allow",
				Action: []string{
					"dynamodb:GetItem",
					"dynamodb:Query",
				},
			},
			{
				Sid:      "CapabilitiesSNSTopicSubscribe",
				Resource: req.CapabilitiesTopicArn,
				Effect:   "Allow",
				Action: []string{
					"sns:GetTopicAttributes",
					"sns:Subscribe",
				},
			},
			{
				Sid: "DiscoveryTableRead",
				Resource: []string{
					req.DiscoveryTableArn,
					req.DiscoveryTableArn + "/index/pubkeys-by-id",
				},
				Effect: "Allow",
				Action: []string{
					"dynamodb:GetItem",
					"dynamodb:Query",
				},
			},
			{
				Sid:      "PubkeySNSTopicSubscribe",
				Resource: req.PubkeyInvalidationTopicArn,
				Effect:   "Allow",
				Action:   "sns:Subscribe",
			},
			{
				Sid: "ServiceRegistrationTableRead",
				Resource: []string{
					req.ServiceRegistrationTableArn,
					req.ServiceRegistrationTableArn + "/index/services-by-id",
				},
				Effect: "Allow",
				Action: []string{
					"dynamodb:BatchGetItem",
					"dynamodb:GetItem",
					"dynamodb:Query",
				},
			},
			{
				Sid:      "SelfServiceSTSCallerIdentity",
				Effect:   "Allow",
				Resource: "*",
				Action: []string{
					"sts:GetCallerIdentity",
				},
			},
			{
				Sid:      "CalleeSQSAdmin",
				Resource: fmt.Sprintf("arn:aws:sqs:*:%s:%s_*", req.MalachaiAccountID, req.ServiceID),
				Effect:   "Allow",
				Action: []string{
					"sqs:AddPermission",
					"sqs:RemovePermission",
					"sqs:CreateQueue",
					"sqs:DeleteMessage",
					"sqs:DeleteQueue",
					"sqs:GetQueueAttributes",
					"sqs:SetQueueAttributes",
					"sqs:ReceiveMessage",
				},
			},
			{
				Sid:    "AllowedRoles",
				Effect: "Allow",
				Resource: []string{
					req.SandstormRoleArn,
					ig.getIAMAuxWildcard(),
				},
				Action: []string{
					"sts:AssumeRole",
				},
			},
			{
				Sid:    "InventoryTableUpdate",
				Effect: "Allow",
				Resource: []string{
					req.InventoryTableArn,
					fmt.Sprintf("%s/index/*", req.InventoryTableArn),
				},
				Action: []string{
					"dynamodb:PutItem",
					"dynamodb:UpdateItem",
				},
			},
		},
	})
	if err != nil {
		return
	}

	policyDocument = buf.String()
	return
}

func (ig *iamGenerator) registerServicePolicy(req *servicePolicyRequest) (policyArn string, err error) {
	err = req.validate()
	if err != nil {
		return
	}

	policyDocument, err := ig.generateServicePolicy(req)
	if err != nil {
		return
	}

	// PolicyName includes a timestamp, because we attach the newer policy before
	// detaching and deleting the old one. Including only serviceID will result in
	// "EntityAlreadyExists" error during UpdateService.
	input := &iam.CreatePolicyInput{
		Description:    aws.String("service policy for malachai service " + req.ServiceName),
		Path:           aws.String(GetIAMPathPrefix(ig.cfg.Environment, req.ServiceName)),
		PolicyDocument: aws.String(policyDocument),
		PolicyName:     aws.String(req.ServiceID),
	}
	output, err := ig.iam.CreatePolicy(input)
	if err != nil {
		return
	}
	policyArn = aws.StringValue(output.Policy.Arn)
	return
}

func (ig *iamGenerator) updatePolicy(policyArn string, req *servicePolicyRequest) (err error) {

	if policyArn == "" {
		return errors.New("Cannot update policy withtout policy arn")
	}

	err = req.validate()
	if err != nil {
		return
	}

	policyDocument, err := ig.generateServicePolicy(req)
	if err != nil {
		return
	}

	input := &iam.CreatePolicyVersionInput{
		PolicyArn:      aws.String(policyArn),
		PolicyDocument: aws.String(policyDocument),
		SetAsDefault:   aws.Bool(true),
	}

	_, err = ig.iam.CreatePolicyVersion(input)
	return
}

func (ig *iamGenerator) UpdateAssumeRolePolicy(svc *Service, allowedArns []string) (err error) {
	assumeRolePolicy, err := generateAssumeRolePolicy(allowedArns)
	if err != nil {
		return
	}

	assumeRolePolicyInput := &iam.UpdateAssumeRolePolicyInput{
		PolicyDocument: aws.String(assumeRolePolicy),
		RoleName:       aws.String(svc.ID),
	}

	_, err = ig.iam.UpdateAssumeRolePolicy(assumeRolePolicyInput)
	if err != nil {
		if isInvalidPrincipalsError(err) {
			deletedPrincipals, pErr := ig.getDeletedPrincipals(svc)
			if pErr == nil && len(deletedPrincipals) > 0 {
				return InvalidPrincipalsError{
					Principals: deletedPrincipals,
				}
			}
		}
		return err
	}
	return
}

// GetValidAllowedARNs returns all allowed arns that can be validated.
// This will weed out any arns that have been deleted.
func (ig *iamGenerator) GetValidAllowedARNs(roleName string, expected []string) ([]string, error) {
	trusted, err := ig.getTrustedAWSPrincipals(roleName)
	if err != nil {
		return nil, err
	}
	return newStringSet(trusted...).Intersection(expected).ToSlice(), nil
}

func (ig *iamGenerator) getDeletedPrincipals(svc *Service) ([]string, error) {
	trusted, err := ig.getTrustedAWSPrincipals(svc.ID)
	if err != nil {
		return nil, err
	}

	deleted := make(stringSet)
	deleted.AddAll(svc.AllowedArns)
	deleted.RemoveAll(trusted)

	return deleted.ToSlice(), nil
}

func (ig *iamGenerator) getTrustedAWSPrincipals(roleName string) ([]string, error) {
	role, err := ig.iam.GetRole(&iam.GetRoleInput{
		RoleName: aws.String(roleName),
	})
	if err != nil {
		return nil, err
	}

	rawDocument, err := url.QueryUnescape(aws.StringValue(role.Role.AssumeRolePolicyDocument))
	if err != nil {
		return nil, err
	}

	var trustRelationship policyDocument
	if err := json.Unmarshal([]byte(rawDocument), &trustRelationship); err != nil {
		return nil, err
	}

	principals, err := trustRelationship.AWSPrincipalsAllowedTo("sts:AssumeRole")
	if err != nil {
		return nil, err
	}

	return principals, nil
}

func (ig *iamGenerator) registerServiceRole(req *roleRequest) (roleArn string, err error) {
	err = req.validate()
	if err != nil {
		return
	}

	assumeRolePolicy, err := generateAssumeRolePolicy(req.AssumeRoleArns)
	if err != nil {
		return
	}

	input := &iam.CreateRoleInput{
		Description:              aws.String("generated role for a malachai service"),
		Path:                     aws.String(GetIAMPathPrefix(ig.cfg.Environment, req.ServiceName)),
		RoleName:                 aws.String(req.ServiceID),
		AssumeRolePolicyDocument: aws.String(assumeRolePolicy),
	}
	output, err := ig.iam.CreateRole(input)
	if err != nil {
		return
	}
	_, err = ig.iam.AttachRolePolicy(&iam.AttachRolePolicyInput{
		PolicyArn: aws.String(req.PolicyArn),
		RoleName:  output.Role.RoleName,
	})
	if err != nil {
		return
	}

	roleArn = aws.StringValue(output.Role.Arn)
	return
}

// DeregisterServiceRole deregisters a service's role
func (ig *iamGenerator) DeregisterServiceRole(svc *Service) (err error) {
	if svc.SandstormRoleName != "" {
		err = ig.spg.Delete(svc.SandstormRoleName, &noOpAuthorizer{})
		if err != nil {
			return
		}
	}

	err = ig.deregisterServiceRole(svc)
	return
}

func (ig *iamGenerator) deletePolicy(policyArn string) (err error) {
	if policyArn == "" {
		return fmt.Errorf("please provide policy arn to delete")
	}

	policyVersions, err := ig.iam.ListPolicyVersions(&iam.ListPolicyVersionsInput{
		PolicyArn: aws.String(policyArn),
	})
	if err != nil {
		return fmt.Errorf("failed to retrieve policy versions for %s, err: %s", policyArn, err.Error())
	}

	err = ig.deletePolicyVersions(policyArn, policyVersions)
	if err != nil {
		return
	}
	_, err = ig.iam.DeletePolicy(&iam.DeletePolicyInput{
		PolicyArn: aws.String(policyArn),
	})

	return
}

func (ig *iamGenerator) deletePolicyVersions(policyArn string, lpvo *iam.ListPolicyVersionsOutput) (err error) {

	for _, v := range lpvo.Versions {
		//Cant delete the default version.
		if aws.BoolValue(v.IsDefaultVersion) {
			continue
		}
		_, err = ig.deletePolicyVersion(policyArn, v.VersionId)
		if err != nil {
			return
		}
	}
	return
}

func (ig *iamGenerator) deletePolicyVersion(policyArn string, versionID *string) (out *iam.DeletePolicyVersionOutput, err error) {
	out, err = ig.iam.DeletePolicyVersion(&iam.DeletePolicyVersionInput{
		PolicyArn: aws.String(policyArn),
		VersionId: versionID,
	})
	return
}

func (ig *iamGenerator) listAttachedRolePolicies(svc *Service) (attachedPolicies *iam.ListAttachedRolePoliciesOutput, err error) {

	attachedPolicies, err = ig.iam.ListAttachedRolePolicies(&iam.ListAttachedRolePoliciesInput{
		PathPrefix: aws.String(GetIAMPathPrefix(ig.cfg.Environment, svc.Name)),
		RoleName:   aws.String(svc.ID),
	})
	return
}
