package policy

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"path"
	"reflect"

	"code.justin.tv/systems/sandstorm/logging"

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

// path prefixes under which all generated roles and policies will be placed
const (
	IAMPathPrefix = "/sandstorm/"

	IAMRolePathSuffix      = "templated/role/"
	IAMPolicyPathSuffix    = "templated/policy/"
	IAMAuxPolicyPathSuffix = "aux_policy/"
)

const (
	wildcard               = "*"
	developmentEnvironment = "development"
)

// IAMPolicyDescription is the default value for the policy description field
const IAMPolicyDescription = "auto-generated policy document for sandstorm-agent"

// CreateRoleRequest holds parameters for creating an IAM role
type CreateRoleRequest struct {
	AllowedArns []string `json:"allowed_arns,omitempty"` //Required
	Name        string   `json:"name,omitempty"`         //Required
	SecretKeys  []string `json:"secret_keys,omitempty"`
	// List of ldapGroups that can admin this role
	Owners []string `json:"owners,omitempty"` //Required

	// allows write access to secrets
	WriteAccess bool `json:"write_access"`
}

// CreateRoleResponse holds data returned when we create a role
type CreateRoleResponse struct {
	RoleArn          string              `json:"generated_role_arn,omitempty"`
	RoleName         string              `json:"name,omitempty"`
	Policy           IAMPolicyDocument   `json:"generated_policy,omitempty"`
	AssumeRolePolicy IAMAssumeRolePolicy `json:"assume_role_policy,omitempty"`
	AllowedArns      []string            `json:"allowed_arns,omitempty"`

	// Arn for policy document that grants access to KMS/SNS/SQS for
	// sandstorm-agent
	AuxPolicyArn string `json:"aux_policy_arn,omitempty"`

	// allows write access to secrets
	WriteAccess bool `json:"write_access"`

	// list containing LDAP Groups of owners
	Owners *Owners `json:"owners,omitempty"`

	Group string `json:"group,omitempty"`
}

// PutRoleRequest holds parameters for updating an IAM role. All existing
// existing secrets and namespace prefixes will be replaced with provided
// secrets if defined. In other words, if you don't want to update secrets,
// don't include `secret_keys` in request. Values that are nil are unset
// and show not overwrite existing values.
type PutRoleRequest struct {
	RoleName    string   `json:"name,omitempty"`
	SecretKeys  []string `json:"secret_keys,omitempty"`
	AllowedArns []string `json:"allowed_arns,omitempty"`
	WriteAccess bool     `json:"write_access"`
	Owners      []string `json:"owners,omitempty"`
}

// PutRoleResponse is returned with updated policy data.
type PutRoleResponse CreateRoleResponse

// IAMRole contains information about generated role and policy
type IAMRole CreateRoleResponse

// IAMPolicy wraps role policy info with aditional golang representation of policy doc
type IAMPolicy struct {
	PolicyArn  string            `json:"policy_arn,omitempty"`
	PolicyName string            `json:"policy_name,omitempty"`
	Policy     IAMPolicyDocument `json:"policy,omitempty"`
}

// IAMPolicyGenerator manages IAM policies and roles required for
// sandstorm-agent.
type IAMPolicyGenerator struct {
	// ARN of pre-existing policy allowing sns/sqs/kms actions
	AuxPolicyArn string

	// ARN of policy similar to AuxPolicyArn but also KMS:GenerateDataKey
	RWAuxPolicyArn string

	// ARN of DynamoDB table containing secrets
	DynamoDBSecretsTableArn      string
	DynamoDBSecretsAuditTableArn string
	DynamoDBNamespaceTableArn    string

	// Table name of RoleOwner Table
	RoleOwnerTableName string

	DynamoDB dynamodbiface.DynamoDBAPI

	IAM iamiface.IAMAPI

	// path prefixes for iam policy/roles
	IAMPathPrefix string

	Logger logging.Logger

	// Account id which holds all iam resources.
	AccountID string
}

func (pg *IAMPolicyGenerator) auxPolicyArn(writeAccess bool) string {
	if writeAccess {
		return pg.RWAuxPolicyArn
	}
	return pg.AuxPolicyArn
}

// SetPathPrefixEnvironment appends environment into configured IAM path prefix.
// i.e. environment='production' sets path prefix to '/sandstorm/production/'
func (pg *IAMPolicyGenerator) SetPathPrefixEnvironment(environment string) {
	pg.IAMPathPrefix = fmt.Sprintf("%s/", path.Join(IAMPathPrefix, environment))
}

// PathPrefix returns an IAM path prefix
func (pg *IAMPolicyGenerator) PathPrefix(suffix string) string {
	s := path.Join(pg.IAMPathPrefix, suffix)
	if s[len(s)-1] != '/' {
		s = s + "/"
	}
	return s
}

// AssumeRolePolicyDocument generates policy document allowing provided arns to
// assume role
func (pg *IAMPolicyGenerator) AssumeRolePolicyDocument(arns []string) IAMAssumeRolePolicy {
	return IAMAssumeRolePolicy{
		Version: DocumentVersion,
		Statement: []IAMAssumeRoleStatement{
			IAMAssumeRoleStatement{
				Action: "sts:AssumeRole",
				Effect: "Allow",
				Principal: map[string]interface{}{
					"AWS": arns,
				},
			},
		},
	}
}

// ExternalAssumeRolePolicyDocument generates policy document allowing a role
// to assume the provided role
func (pg *IAMPolicyGenerator) ExternalAssumeRolePolicyDocument(roleArn string) IAMAssumeRolePolicy {
	return IAMAssumeRolePolicy{
		Version: DocumentVersion,
		Statement: []IAMAssumeRoleStatement{
			IAMAssumeRoleStatement{
				Action:   "sts:AssumeRole",
				Effect:   "Allow",
				Resource: []string{roleArn},
			},
		},
	}
}

// PolicyDocument generates IAM policy document for accessing provided secret keys
func (pg *IAMPolicyGenerator) PolicyDocument(secretKeys, namespaces []string, writeAccess bool, roleARN string) IAMPolicyDocument {
	statements := []IAMPolicyStatement{
		pg.NamespaceIndexPolicyStatement(namespaces, writeAccess),
		pg.SecretsPolicyStatement(secretKeys, writeAccess),
		pg.AssumeRolePolicyStatement(roleARN),
	}

	if writeAccess {
		statements = append(statements, pg.NamespacePolicyStatement(namespaces, writeAccess))
	}

	return IAMPolicyDocument{
		Version:   DocumentVersion,
		Statement: statements,
	}
}

// attachAuxPolicy attaches configured aux policy to specified role. AWS does
// not return an error if we attach a policy to a role that already has that
// policy, so this doesn't make the extra call to check if it's attached.
func (pg *IAMPolicyGenerator) attachAuxPolicy(roleName string, writeAccess bool) (err error) {
	policyArn := pg.auxPolicyArn(writeAccess)
	attachPolicyInput := &iam.AttachRolePolicyInput{
		RoleName:  aws.String(roleName),
		PolicyArn: aws.String(policyArn),
	}

	pg.Logger.Debugf("policy: attaching policy %s to role %s", policyArn, roleName)
	_, err = pg.IAM.AttachRolePolicy(attachPolicyInput)
	if err != nil {
		pg.Logger.Errorf("policy: error attaching role policy: %s", err.Error())
		return
	}
	return
}

// GetPolicyDocument retrieves a policy by ARN and returns the default version
// as an IAMPolicyDocument
func (pg *IAMPolicyGenerator) GetPolicyDocument(policyArn string) (doc IAMPolicyDocument, err error) {
	getPolicyOutput, err := pg.IAM.GetPolicy(&iam.GetPolicyInput{PolicyArn: aws.String(policyArn)})
	if err != nil {
		return
	}

	getPolicyVersionOutput, err := pg.IAM.GetPolicyVersion(&iam.GetPolicyVersionInput{
		VersionId: getPolicyOutput.Policy.DefaultVersionId,
		PolicyArn: getPolicyOutput.Policy.Arn,
	})
	if err != nil {
		return
	}

	policyDoc, err := url.QueryUnescape(aws.StringValue(getPolicyVersionOutput.PolicyVersion.Document))
	if err != nil {
		return
	}

	err = json.Unmarshal([]byte(policyDoc), &doc)
	if err != nil {
		return
	}
	return
}

// GetAttachedRolePolicy is a convenience function for getting the managed policy
// attached to a specified role
func (pg *IAMPolicyGenerator) GetAttachedRolePolicy(roleName string) (arn string, doc IAMPolicyDocument, err error) {
	listAttachedRolePoliciesInput := &iam.ListAttachedRolePoliciesInput{
		RoleName:   aws.String(roleName),
		PathPrefix: aws.String(pg.PathPrefix(IAMPolicyPathSuffix)),
	}
	rolePolicies, err := pg.IAM.ListAttachedRolePolicies(listAttachedRolePoliciesInput)
	if err != nil {
		return
	}

	// By design, shouldn't be <1 or >1.
	// When < 1, fail early
	if len(rolePolicies.AttachedPolicies) < 1 {
		err = InvalidRolePolicyError{
			msg:  "No policies attached to templated role",
			code: http.StatusConflict,
		}
		doc.IsInvalid = err
		return
	}

	arn = aws.StringValue(rolePolicies.AttachedPolicies[0].PolicyArn)
	doc, err = pg.GetPolicyDocument(arn)
	if len(rolePolicies.AttachedPolicies) > 1 {
		doc.IsInvalid = InvalidRolePolicyError{
			msg:  fmt.Sprintf("Too many policies attached to role. Policies attached: %d", len(rolePolicies.AttachedPolicies)),
			code: http.StatusConflict,
		}
	}
	return
}

// ListAttachedRolePolicies is a convenience function for getting all managed policies
// attached to a specified role
// TODO: Parallelize
func (pg *IAMPolicyGenerator) ListAttachedRolePolicies(roleName string) ([]IAMPolicy, error) {
	var rolePolicies []IAMPolicy

	listAttachedRolePoliciesInput := &iam.ListAttachedRolePoliciesInput{
		RoleName: aws.String(roleName),
	}
	attachedRolePolicies, err := pg.IAM.ListAttachedRolePolicies(listAttachedRolePoliciesInput)
	if err != nil {
		return nil, err
	}

	var currentRolePolicy IAMPolicy
	var currentPolicyDoc IAMPolicyDocument

	for _, outputPolicy := range attachedRolePolicies.AttachedPolicies {
		currentPolicyDoc, err = pg.GetPolicyDocument(*outputPolicy.PolicyArn)
		if err != nil {
			return nil, err
		}
		currentRolePolicy = IAMPolicy{
			PolicyArn:  *outputPolicy.PolicyArn,
			PolicyName: *outputPolicy.PolicyName,
			Policy:     currentPolicyDoc,
		}

		rolePolicies = append(rolePolicies, currentRolePolicy)

	}

	return rolePolicies, err
}

// GetAttachedRoleAuxPolicy returns aux policy arn and document. First searches
// to see if the arn configured for the policy generator is attached to the
// role, and if it's not attached, look for an aux policy under the aux policy
// path prefix.
func (pg *IAMPolicyGenerator) GetAttachedRoleAuxPolicy(roleName string) (arn string, doc IAMPolicyDocument, err error) {

	var auxPolicyArn *string
	auxPolicyArn, err = pg.GetAttachedRoleAuxPolicyArn(roleName)
	if err != nil {
		return
	}
	arn = aws.StringValue(auxPolicyArn)
	doc, err = pg.GetPolicyDocument(arn)
	if err != nil {
		return
	}
	return
}

// GetAttachedRoleAuxPolicyArn return aux polcy attached to a role.
func (pg *IAMPolicyGenerator) GetAttachedRoleAuxPolicyArn(roleName string) (*string, error) {

	listAttachedRolePoliciesInput := &iam.ListAttachedRolePoliciesInput{
		RoleName:   aws.String(roleName),
		PathPrefix: aws.String(pg.PathPrefix(IAMAuxPolicyPathSuffix)),
	}
	rolePolicies, err := pg.IAM.ListAttachedRolePolicies(listAttachedRolePoliciesInput)
	if err != nil {
		return nil, err
	}
	if len(rolePolicies.AttachedPolicies) != 1 {
		return nil, fmt.Errorf("More than one aux policies are attached to role: %s", roleName)
	}

	return rolePolicies.AttachedPolicies[0].PolicyArn, nil
}

// GetAttachedGroupAuxPolicyArn Attached managed policy to group
func (pg *IAMPolicyGenerator) GetAttachedGroupAuxPolicyArn(groupName string) (*string, error) {

	policieInput := pg.getGroupAuxPolicyInput(groupName)
	groupPolicies, err := pg.IAM.ListAttachedGroupPolicies(policieInput)
	if err != nil {
		return nil, err
	}

	if len(groupPolicies.AttachedPolicies) != 1 {
		return nil, fmt.Errorf("no attached aux policy to group: %s", groupName)
	}

	return groupPolicies.AttachedPolicies[0].PolicyArn, nil
}

// GetAllowedARNs takes a roleName and returns a slice of all allowed arns
func (pg *IAMPolicyGenerator) GetAllowedARNs(roleName string) ([]string, error) {
	output, err := pg.IAM.GetRole(&iam.GetRoleInput{RoleName: aws.String(roleName)})
	if err != nil {
		return nil, err
	}

	if output.Role.AssumeRolePolicyDocument == nil {
		// return some error
		return nil, InvalidRolePolicyError{msg: "No AssumeRolePolicy returned", code: 409}
	}
	u, err := url.QueryUnescape(*output.Role.AssumeRolePolicyDocument)
	if err != nil {
		return nil, err
	}

	doc := IAMAssumeRolePolicy{}
	err = json.Unmarshal([]byte(u), &doc)
	if err != nil {
		return nil, err
	}

	var allowedARNs []string
	for _, statement := range doc.Statement {
		switch arn := statement.Principal["AWS"].(type) {
		case string:
			allowedARNs = append(allowedARNs, arn)
		case []string:
			allowedARNs = append(allowedARNs, arn...)
		case []interface{}:
			for _, val := range arn {
				switch nestedStatement := val.(type) {
				case string:
					allowedARNs = append(allowedARNs, nestedStatement)
				case []string:
					allowedARNs = append(allowedARNs, nestedStatement...)
				}
			}

		default:
			pg.Logger.Warnf("GetAllowedARNs found type not handled by switch: %v\n", reflect.TypeOf(arn))
		}

	}

	return allowedARNs, err

}

// GenerateRoleARN returns the (supposed) RoleARN for a role of given name
// bases env on pg env
func (pg *IAMPolicyGenerator) GenerateRoleARN(rolename string) string {
	return fmt.Sprintf("arn:aws:iam::%s:role%s%s", pg.AccountID, pg.PathPrefix(IAMRolePathSuffix), rolename)
}

func (pg *IAMPolicyGenerator) getGroupAuxPolicyInput(groupName string) *iam.ListAttachedGroupPoliciesInput {
	listAttachedGroupAuxPolicy := &iam.ListAttachedGroupPoliciesInput{
		GroupName:  aws.String(groupName),
		PathPrefix: aws.String(pg.PathPrefix(IAMAuxPolicyPathSuffix)),
	}
	return listAttachedGroupAuxPolicy
}
