package impl

import (
	"context"
	"encoding/json"
	"fmt"
	"regexp"
	"sort"

	"github.com/AlekSi/pointer"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"go.uber.org/zap"
)

type AWSBucketPolicy struct {
	Version   string        `json:"Version"`
	Statement []interface{} `json:"Statement"`
}

type AWSBucketPolicyStatement struct {
	Sid       string                    `json:"Sid"`
	Effect    AWSBucketPolicyEffectType `json:"Effect"`
	Principal map[string]string         `json:"Principal"`
	Action    interface{}               `json:"Action"`
	Resource  []string                  `json:"Resource"`
}

type AWSBucketPolicyEffectType string

const (
	AWSBucketPolicyEffectTypeAllow AWSBucketPolicyEffectType = "Allow"
)

const (
	DefaultBucketPolicyVersion = "2012-10-17"

	DefaultBucketResourceFormat = "arn:aws:s3:::%s"
	DefaultBucketSid            = "ManageByDataflowService"
	DefaultBucketPrincipalKey   = "AWS"

	InvalidPattern = `"Principal"\s*:\s*\{\s*"AWS"\s*:\s*"([A-Za-z0-9]+)"\s*\}`
)

func AddRolesToBucketPolicy(
	ctx context.Context, bucketName string, currentPolicy *string, iamRoleARNs []string,
) (*string, []string, error) {
	return ModifyRolesToBucketPolicy(ctx, bucketName, currentPolicy, iamRoleARNs, nil)
}

func ModifyRolesToBucketPolicy(
	ctx context.Context, bucketName string, currentPolicy *string, addIamRoleARNs []string, deleteIamRoleARNs []string,
) (*string, []string, error) {
	logger := log.GetLogger(ctx)
	if currentPolicy == nil {
		return generateAWSBucketPolicyJson(logger, bucketName, diffSlice(addIamRoleARNs, deleteIamRoleARNs))
	}
	policy := &AWSBucketPolicy{}
	err := json.Unmarshal([]byte(*currentPolicy), policy)
	if err != nil {
		logger.Warn("unmarshal bucket policy failed", zap.Error(err))
		return nil, nil, err
	}

	if len(policy.Statement) == 0 {
		return generateAWSBucketPolicyJson(logger, bucketName, diffSlice(addIamRoleARNs, deleteIamRoleARNs))
	}

	var noNeedDealStatements []interface{}
	var existedIamRoleARNs []string
	for _, item := range policy.Statement {
		valid, policyStatement := queryBucketPolicyStatement(logger, item)
		if !valid {
			continue
		}
		if policyStatement == nil {
			noNeedDealStatements = append(noNeedDealStatements, item)
			continue
		}
		existedIamRoleARNs = append(existedIamRoleARNs, policyStatement.Principal[DefaultBucketPrincipalKey])
	}

	mergeIamRoleARNs := diffSlice(sliceUnion(existedIamRoleARNs, addIamRoleARNs), deleteIamRoleARNs)
	statements := buildAWSBucketPolicyStatements(mergeIamRoleARNs, bucketName)
	var interfaceSlice []interface{}
	for _, statement := range statements {
		interfaceSlice = append(interfaceSlice, statement)
	}
	interfaceSlice = append(interfaceSlice, noNeedDealStatements...)
	policyJson, err := buildAWSBucketPolicyJson(interfaceSlice)
	if err != nil {
		logger.Warn("build bucket policy json failed", zap.Error(err))
		return nil, nil, err
	}
	return policyJson, mergeIamRoleARNs, nil
}

func generateAWSBucketPolicyJson(
	logger *zap.Logger, bucketName string, iamRoleARNs []string,
) (*string, []string, error) {
	statements := buildAWSBucketPolicyStatements(iamRoleARNs, bucketName)
	var interfaceSlice []interface{}
	for _, statement := range statements {
		interfaceSlice = append(interfaceSlice, statement)
	}
	policyJson, err := buildAWSBucketPolicyJson(interfaceSlice)
	if err != nil {
		logger.Warn("build bucket policy json failed", zap.Error(err))
		return nil, nil, err
	}
	return policyJson, iamRoleARNs, nil
}

func buildAWSBucketPolicyJson(statements []interface{}) (*string, error) {
	bucketPolicy := &AWSBucketPolicy{
		Version:   DefaultBucketPolicyVersion,
		Statement: statements,
	}
	bytes, err := json.Marshal(bucketPolicy)
	if err != nil {
		return nil, err
	}
	return pointer.To(string(bytes)), nil
}

func buildAWSBucketPolicyStatements(iamRoleARNs []string, bucketName string) []AWSBucketPolicyStatement {
	var statements []AWSBucketPolicyStatement
	for _, iamRoleARN := range iamRoleARNs {
		statements = append(statements, AWSBucketPolicyStatement{
			Sid:    DefaultBucketSid,
			Effect: AWSBucketPolicyEffectTypeAllow,
			Principal: map[string]string{
				DefaultBucketPrincipalKey: iamRoleARN,
			},
			Action: "s3:*",
			Resource: []string{
				fmt.Sprintf(DefaultBucketResourceFormat, bucketName),
				fmt.Sprintf(DefaultBucketResourceFormat, bucketName) + "/*",
			},
		})
	}
	return statements
}

func queryBucketPolicyStatement(logger *zap.Logger, value interface{}) (bool, *AWSBucketPolicyStatement) {
	marshal, err := json.Marshal(value)
	if err != nil {
		logger.Info("marshal bucket policy statement failed", zap.Any("statement", value), zap.Error(err))
		return true, nil
	}
	matched, _ := regexp.MatchString(InvalidPattern, string(marshal))
	if matched {
		logger.Info("bucket policy statement invalid", zap.Any("statement", value))
		return false, nil
	}

	var statement *AWSBucketPolicyStatement
	err = json.Unmarshal(marshal, &statement)
	if err != nil {
		logger.Info("unmarshal bucket policy statement failed", zap.Any("statement", value), zap.Error(err))
		return true, nil
	}
	if statement.Sid != DefaultBucketSid {
		logger.Info("bucket policy statement sid not match", zap.Any("sid", statement.Sid))
		return true, nil
	}
	return true, statement
}

func sliceUnion(a, b []string) []string {
	set := make(map[string]bool)
	for _, v := range a {
		set[v] = true
	}

	for _, v := range b {
		set[v] = true
	}

	result := make([]string, len(set))
	i := 0
	for k := range set {
		result[i] = k
		i++
	}
	sort.Strings(result)
	return result
}

func diffSlice(a, b []string) []string {
	if len(a) == 0 || len(b) == 0 {
		return a
	}
	var diff []string
	m := make(map[string]bool)

	for _, item := range b {
		m[item] = true
	}

	for _, item := range a {
		if _, ok := m[item]; !ok {
			diff = append(diff, item)
		}
	}
	return diff
}