package auth

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	"github.com/aws/aws-sdk-go-v2/aws"
	awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
	awsConfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"
	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/expression"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

// copy from sdk/rbac/helperimpl/permission_region.go

var (
	globalDynamoDBClient     *dynamodb.Client
	globalDynamoDBClientInit sync.Once
)

var DefaultRegionApiKeyPermissionHelper = NewRegionApiKeyPermissionHelper()

type RegionApiKeyPermissionHelper interface {
	PermissionPlugin
	SetAwsBasic(awsRegion string, awsNeedAssumeRole bool, awsSessionName, awsRoleArn string)
	SetDynamoDBTableName(apiKeyTableName, privilegeRoleTableName string)
}

type regionApiKeyPermissionHelper struct {
	awsRegion              string
	awsNeedAssumeRole      bool
	awsRoleSessionName     string
	awsRoleArn             string
	apiKeyTableName        string
	privilegeRoleTableName string
}

func NewRegionApiKeyPermissionHelper() RegionApiKeyPermissionHelper {

	return &regionApiKeyPermissionHelper{
		awsRegion:              "ap-northeast-2",
		awsNeedAssumeRole:      false,
		awsRoleSessionName:     "",
		awsRoleArn:             "",
		apiKeyTableName:        "account_uni_api_keys",
		privilegeRoleTableName: "account_privilege_role",
	}
}

func (h *regionApiKeyPermissionHelper) SetAwsBasic(
	awsRegion string, awsNeedAssumeRole bool, awsSessionName, awsRoleArn string) {

	if awsRegion != "" {
		h.awsRegion = awsRegion
	}

	h.awsNeedAssumeRole = awsNeedAssumeRole
	if h.awsNeedAssumeRole {
		h.awsRoleArn = awsRoleArn
		h.awsRoleSessionName = awsSessionName
	}

}

func (h *regionApiKeyPermissionHelper) SetDynamoDBTableName(apiKeyTableName, privilegeRoleTableName string) {
	if apiKeyTableName != "" {
		h.apiKeyTableName = apiKeyTableName
	}
	if privilegeRoleTableName != "" {
		h.privilegeRoleTableName = privilegeRoleTableName
	}
}

func (h *regionApiKeyPermissionHelper) getClient() *dynamodb.Client {
	globalDynamoDBClientInit.Do(func() {
		ctx := context.Background()
		buildableClient := awshttp.NewBuildableClient().WithTransportOptions(func(t *http.Transport) {
			t.MaxIdleConnsPerHost = 150
			t.MaxIdleConns = 300
		})
		cfg, err := awsConfig.LoadDefaultConfig(ctx, awsConfig.WithRegion(h.awsRegion), awsConfig.WithHTTPClient(buildableClient))
		if err != nil {
			panic(fmt.Sprintf("regionApiKeyPermissionHelper: load aws config error: %s", err))
		}

		if h.awsNeedAssumeRole {
			stsClient := sts.NewFromConfig(cfg)

			durationSeconds := int32(60 * 60)
			expiryWindow := 30 * time.Second

			_, err = stsClient.AssumeRole(ctx, &sts.AssumeRoleInput{
				RoleArn:         &h.awsRoleArn,
				RoleSessionName: &h.awsRoleSessionName,
				DurationSeconds: &durationSeconds,
			})
			if err != nil {
				panic(fmt.Sprintf("regionApiKeyPermissionHelper: assume role error: %s", err))
			}
			cfg.Credentials = aws.NewCredentialsCache(
				stscreds.NewAssumeRoleProvider(stsClient, h.awsRoleArn),
				func(options *aws.CredentialsCacheOptions) {
					options.ExpiryWindow = expiryWindow
				},
			)
		}
		globalDynamoDBClient = dynamodb.NewFromConfig(cfg)
	})
	return globalDynamoDBClient
}

func parseAppIDFromReqUri(reqUri string) (string, error) {
	// dataapi/dataapps/{app_id}/action
	// dataapi/dataapps/12345678/action
	split := strings.Split(reqUri, "/")
	l := len(split)
	for i, str := range split {
		if str == "dataapps" && i < l-1 {
			sAppId := split[i+1]
			_, err := strconv.ParseUint(sAppId, 10, 64)
			if err == nil {
				return sAppId, nil
			}
		}
	}
	return "", errors.New(fmt.Sprintf("not found valid appid from %s", reqUri))
}

func (h *regionApiKeyPermissionHelper) HasPermission(ctx context.Context, param *PermissionParam) (*PermissionRet, error) {
	if param.IdentityType != AccessIdentityTypeApiKeyAK {
		return nil, errors.New("identity type not support")
	}
	var identity = param.IdentityVal
	var resource = param.ResourceUri
	appID, err := parseAppIDFromReqUri(param.ReqUri)
	if err != nil {
		tilog.Warn(ctx, "regionApiKeyPermissionHelper, parseAppIDFromReqUri error", zap.Any("param", param), zap.Any("err", err))
		return nil, err
	}

	tilog.Info(ctx, "regionApiKeyPermissionHelper, HasPermission start", zap.Any("identity", identity), zap.Any("resource", resource), zap.Any("appID", appID))
	if identity == "" {
		return nil, errors.New("identity is blank")
	}
	if resource == "" {
		return nil, errors.New("resource is blank")
	}
	if appID == "" {
		return nil, errors.New("appID is blank")
	}
	appIDInt, err := strconv.ParseUint(appID, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("appID is not a valid uint64: %s", err)
	}

	ddb := h.getClient()
	apiKey, err := getApiKeyFromDynamodb(ddb, h.apiKeyTableName, identity)
	tilog.Info(ctx, "regionApiKeyPermissionHelper, get apiKey from dynamodb complete", zap.Any("apiKeyTableName", h.apiKeyTableName), zap.Any("identity", identity), zap.Any("apiKey", apiKey), zap.Any("err", err))
	if err != nil {
		return nil, fmt.Errorf("get apiKey from dynamodb failed: %s", err)
	}
	apiKeyRoles, ok := apiKey.AppRoleMap[appIDInt]
	if !ok {
		tilog.Info(ctx, "regionApiKeyPermissionHelper, do not found appId relate apiKeyRoles", zap.Any("appId", appIDInt), zap.Any("apiKey", apiKey))
		return nil, errors.New("there is no roles found in the apiKey for this app")
	}
	resourceRoles, err := getRolesByResourceUri(ddb, h.privilegeRoleTableName, resource)
	tilog.Info(ctx, "regionApiKeyPermissionHelper, get roles by resource_uri complete", zap.Any("resourceUri", resource), zap.Any("resourceRoles", resourceRoles), zap.Any("err", err))
	if err != nil {
		return nil, fmt.Errorf("get roles by resource uri failed: %s", err)
	}
	if hasIntersection(apiKeyRoles, resourceRoles) {
		tilog.Info(ctx, "regionApiKeyPermissionHelper, apiKeyRoles and resourceRoles hasIntersection", zap.Any("apiKeyRoles", apiKeyRoles), zap.Any("resourceRoles", resourceRoles))
		return &PermissionRet{
			HasPermission: true,
		}, nil
	}
	tilog.Info(ctx, "regionApiKeyPermissionHelper, apiKeyRoles and resourceRoles do not hasIntersection", zap.Any("apiKeyRoles", apiKeyRoles), zap.Any("resourceRoles", resourceRoles))
	return &PermissionRet{
		HasPermission: false,
	}, nil
}

type DdbApiKey struct {
	AccessKey  string              `dynamodbav:"identity"`
	PrivateKey string              `dynamodbav:"private_key"`
	AppRoleMap map[uint64][]uint64 `dynamodbav:"role_map"`
}

type DdbPrivilegeRole struct {
	ResourceUri string `dynamodbav:"resource_uri"`
	RoleID      uint64 `dynamodbav:"role_id"`
	PrivilegeID uint64 `dynamodbav:"privilege_id"`
}

func getApiKeyFromDynamodb(ddb *dynamodb.Client, apiKeyTableName, accessKey string) (DdbApiKey, error) {
	apiKey := DdbApiKey{}
	selectedKeys := map[string]string{
		"identity": fmt.Sprintf("api_key_%s", accessKey),
	}
	key, err := attributevalue.MarshalMap(selectedKeys)
	if err != nil {
		return apiKey, fmt.Errorf("marshal key error: %s", err)
	}
	data, err := ddb.GetItem(context.Background(), &dynamodb.GetItemInput{
		TableName: aws.String(apiKeyTableName),
		Key:       key,
	})
	if err != nil {
		return apiKey, fmt.Errorf("get item error: %s", err)
	}
	if data.Item == nil {
		return apiKey, fmt.Errorf("no item found")
	}
	err = attributevalue.UnmarshalMap(data.Item, &apiKey)
	if err != nil {
		return apiKey, fmt.Errorf("unmarshal item error: %s", err)
	}
	return apiKey, nil
}

func getRolesByResourceUri(ddb *dynamodb.Client, privilegeRoleTableName, resourceUri string) ([]uint64, error) {
	privilegeRoles := []DdbPrivilegeRole{}
	keyEx := expression.Key("resource_uri").Equal(expression.Value(resourceUri))
	expr, err := expression.NewBuilder().WithKeyCondition(keyEx).Build()
	if err != nil {
		return nil, fmt.Errorf("build expression error: %s", err)
	}
	data, err := ddb.Query(context.Background(), &dynamodb.QueryInput{
		TableName:                 aws.String(privilegeRoleTableName),
		KeyConditionExpression:    expr.KeyCondition(),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
	})
	if err != nil {
		return nil, fmt.Errorf("query error: %s", err)
	}
	err = attributevalue.UnmarshalListOfMaps(data.Items, &privilegeRoles)
	if err != nil {
		return nil, fmt.Errorf("unmarshal item error: %s", err)
	}
	roles := make([]uint64, 0, len(privilegeRoles))
	for _, pr := range privilegeRoles {
		roles = append(roles, pr.RoleID)
	}
	return roles, nil
}

func hasIntersection(roles1, roles2 []uint64) bool {
	if len(roles1) == 0 || len(roles2) == 0 {
		return false
	}
	m := make(map[uint64]bool)
	for _, r := range roles1 {
		m[r] = true
	}
	for _, r := range roles2 {
		if m[r] {
			return true
		}
	}
	return false
}