package iam

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/url"

	v1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/pkg/common/errors"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/gcp"
)

// TODO refactor to move const and type definition
const (
	GServicaAccount = "gserviceaccount.com"
	SAPrefix        = "serviceAccount"
)

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

type Statement struct {
	Effect    string         `json:"Effect"`
	Principal map[string]any `json:"Principal"`
	Action    string         `json:"Action"`
	Condition map[string]any `json:"Condition,omitempty"`
}

func getAWSAssumeEntityFromCRD(assumeRolePolicyDocument string) (AssumeRolePolicy, error) {
	content, err := base64.StdEncoding.DecodeString(assumeRolePolicyDocument)
	if err != nil {
		return AssumeRolePolicy{}, err
	}

	policy := AssumeRolePolicy{}
	err = json.Unmarshal(content, &policy)
	if err != nil {
		return AssumeRolePolicy{}, err
	}

	return policy, nil
}

// getAliCloudAssumeEntityFromCRD decodes and parses Aliyun assume role policy from CRD
func getAliCloudAssumeEntityFromCRD(assumeRolePolicyDocument string) (AssumeRolePolicy, error) {
	content, err := base64.StdEncoding.DecodeString(assumeRolePolicyDocument)
	if err != nil {
		return AssumeRolePolicy{}, err
	}

	policy := AssumeRolePolicy{}
	err = json.Unmarshal(content, &policy)
	if err != nil {
		return AssumeRolePolicy{}, err
	}

	return policy, nil
}

func (iam *IAM) getAWSAssumeEntityFromAWS(ctx context.Context, roleName string) (AssumeRolePolicy, error) {
	iamRoleOutput, exists, err := iam.iamcli.GetRole(ctx, roleName)
	if err != nil {
		return AssumeRolePolicy{}, errors.New(errors.AWSClientError, "get iam role failed, err: "+err.Error())
	}

	if !exists {
		return AssumeRolePolicy{}, errors.New(errors.RoleNotFound, "role not found")
	}

	if iamRoleOutput == nil || iamRoleOutput.Role == nil {
		return AssumeRolePolicy{}, errors.New(errors.RoleNotFound, "no role in iam response")
	}

	if iamRoleOutput.Role.AssumeRolePolicyDocument == nil ||
		len(*iamRoleOutput.Role.AssumeRolePolicyDocument) == 0 {

		return AssumeRolePolicy{}, nil
	}

	awsAssumeRolePolicyDocument := *iamRoleOutput.Role.AssumeRolePolicyDocument
	awsAssumeRolePolicyDocument, err = url.QueryUnescape(awsAssumeRolePolicyDocument)
	if err != nil {
		return AssumeRolePolicy{}, err
	}
	policy := AssumeRolePolicy{}
	err = json.Unmarshal([]byte(awsAssumeRolePolicyDocument), &policy)
	if err != nil {
		return AssumeRolePolicy{}, err
	}

	return policy, nil
}

// getAliCloudAssumeEntityFromAliCloud retrieves assume role policy from Aliyun
func (iam *IAM) getAliCloudAssumeEntityFromAliCloud(ctx context.Context, roleName string) (AssumeRolePolicy, error) {
	aliRoleOutput, exists, err := iam.aliCli.GetRole(ctx, roleName)
	if err != nil {
		return AssumeRolePolicy{}, errors.New(errors.AliCloudClientError, "get ram role failed, err: "+err.Error())
	}

	if !exists {
		return AssumeRolePolicy{}, errors.New(errors.RoleNotFound, "role not found in aliyun")
	}

	if aliRoleOutput == nil || aliRoleOutput.Body == nil || aliRoleOutput.Body.Role == nil || aliRoleOutput.Body.Role.AssumeRolePolicyDocument == nil || *aliRoleOutput.Body.Role.AssumeRolePolicyDocument == "" {
		return AssumeRolePolicy{}, nil
	}

	// Aliyun policy document is already URL encoded
	decodedPolicy, err := url.QueryUnescape(*aliRoleOutput.Body.Role.AssumeRolePolicyDocument)
	if err != nil {
		return AssumeRolePolicy{}, err
	}

	policy := AssumeRolePolicy{}
	err = json.Unmarshal([]byte(decodedPolicy), &policy)
	if err != nil {
		return AssumeRolePolicy{}, err
	}

	return policy, nil
}

func (iam *IAM) updateAssumeEntityAWS(ctx context.Context, roleName string, role *v1.IAM) ([]byte, error) {
	policyFromCRD, err := getAWSAssumeEntityFromCRD(role.Spec.Aws.AssumeRolePolicyDocument)
	if err != nil {
		return nil, err
	}

	policyFromAWS, err := iam.getAWSAssumeEntityFromAWS(ctx, roleName)
	if err != nil {
		return nil, err
	}

	lastPolicyFromCRD := AssumeRolePolicy{}
	if len(role.Spec.Aws.LastAssumeRolePolicyDocument) > 0 {
		policy, err := getAWSAssumeEntityFromCRD(role.Spec.Aws.LastAssumeRolePolicyDocument)
		if err != nil {
			return nil, err
		}
		lastPolicyFromCRD = policy
	} else {
		maxOldStatementLen := len(policyFromAWS.Statement)
		if maxOldStatementLen > len(policyFromCRD.Statement) {
			maxOldStatementLen = len(policyFromCRD.Statement)
		}
		buf, _ := json.Marshal(policyFromCRD.Statement[:maxOldStatementLen])
		oldStatement := []Statement{}
		_ = json.Unmarshal(buf, &oldStatement)
		lastPolicyFromCRD.Version = "2012-10-17"
		lastPolicyFromCRD.Statement = oldStatement
	}

	newPolicy := calcAssumeRolePolicy(policyFromCRD, lastPolicyFromCRD, policyFromAWS)

	// AWS rejects an assume-role policy with no statements
	// ("MalformedPolicyDocument: Statement is empty"). An empty result means the
	// merge stripped every principal (e.g. all desired principals were removed
	// from AWS externally). Skip the update and keep the role's current trust
	// instead of failing the whole reconcile. A nil return signals the caller to
	// leave the CRD trust documents untouched.
	if len(newPolicy.Statement) == 0 {
		return nil, nil
	}

	assumeDocument, err := json.Marshal(newPolicy)
	if err != nil {
		return nil, err
	}

	err = iam.iamcli.UpdateAssumeRolePolicy(ctx, roleName, string(assumeDocument))
	if err != nil {
		return nil, err
	}

	return assumeDocument, nil
}

// updateAssumeEntityAliCloud updates assume role policy for Aliyun RAM role
func (iam *IAM) updateAssumeEntityAliCloud(ctx context.Context, roleName string, role *v1.IAM) ([]byte, error) {
	// Get policy from CRD
	policyFromCRD, err := getAliCloudAssumeEntityFromCRD(role.Spec.AliCloud.AssumeRolePolicyDocument)
	if err != nil {
		return nil, err
	}

	// Get current policy from Aliyun
	policyFromAliCloud, err := iam.getAliCloudAssumeEntityFromAliCloud(ctx, roleName)
	if err != nil {
		return nil, err
	}

	// Handle last applied policy
	lastPolicyFromCRD := AssumeRolePolicy{}
	if len(role.Spec.AliCloud.LastAssumeRolePolicyDocument) > 0 {
		policy, err := getAliCloudAssumeEntityFromCRD(role.Spec.AliCloud.LastAssumeRolePolicyDocument)
		if err != nil {
			return nil, err
		}
		lastPolicyFromCRD = policy
	} else {
		// For initial sync, use partial policy from CRD
		maxOldStatementLen := min(len(policyFromAliCloud.Statement), len(policyFromCRD.Statement))
		buf, _ := json.Marshal(policyFromCRD.Statement[:maxOldStatementLen])
		oldStatement := []Statement{}
		_ = json.Unmarshal(buf, &oldStatement)
		lastPolicyFromCRD.Version = "1" // Aliyun uses version "1"
		lastPolicyFromCRD.Statement = oldStatement
	}

	// Calculate merged policy
	// todo aliyun, if merging need?
	//newPolicy := calcAssumeRolePolicy(policyFromCRD, lastPolicyFromCRD, policyFromAliCloud)
	newPolicy := lastPolicyFromCRD
	// Aliyun requires policy document to be in specific format
	newPolicy.Version = "1"
	assumeDocument, err := json.Marshal(newPolicy)
	if err != nil {
		return nil, err
	}

	// Update assume role policy
	err = iam.aliCli.UpdateAssumeRolePolicy(ctx, roleName, string(assumeDocument))
	if err != nil {
		return nil, errors.New(errors.AliCloudClientError, "update assume role policy failed: "+err.Error())
	}

	return assumeDocument, nil
}

func (iam *IAM) getAssumeRoleAccountFromGCP(ctx context.Context, cli providers.IAMClientGCP, roleName string) ([]string, error) {
	_, exists, err := cli.GetServiceAccount(ctx, roleName)
	if err != nil {
		return nil, err
	}

	if !exists {
		return nil, errors.New(errors.RoleNotFound, "gcp service account not found")
	}

	//query AssumeRolePolicy
	policy, err := cli.GetServiceAccountIAM(ctx, roleName)
	if err != nil {
		return nil, err
	}

	if policy == nil {
		return nil, errors.New(errors.RoleNotFound, "gcp service account iam not found")
	}

	_, binding := gcp.FindBindingPolicyorCreate(policy.Bindings)
	return binding.Members, nil
}

func (iam *IAM) calcAssumeEntityGCP(ctx context.Context, cli providers.IAMClientGCP, roleName string, role *v1.IAM) ([]string, error) {
	accountsFromCRD := role.Spec.Gcp.AssumeRoleAccounts

	accountsFromGCP, err := iam.getAssumeRoleAccountFromGCP(ctx, cli, roleName)
	if err != nil {
		return nil, err
	}

	lastAccountsFromCRD := role.Spec.Gcp.LastAssumeRoleAccounts
	if len(lastAccountsFromCRD) == 0 {
		maxOldAccountsLen := len(accountsFromGCP)
		if maxOldAccountsLen > len(accountsFromCRD) {
			maxOldAccountsLen = len(accountsFromCRD)
		}
		lastAccountsFromCRD = make([]string, maxOldAccountsLen)
		copy(lastAccountsFromCRD, accountsFromCRD[:maxOldAccountsLen])
	}

	newAccount := calcAssumeRoleAccountGCP(accountsFromCRD, lastAccountsFromCRD, accountsFromGCP)
	return newAccount, nil
}

// awsPrincipalKey returns a stable, hashable map key for an AWS principal value.
// The principal ("AWS") may be a single string or a list (e.g. ["arnA","arnB"]);
// after a JSON round-trip a list decodes to []interface{}, which is not hashable
// and panics when used directly as a map key. Callers must key on this instead.
func awsPrincipalKey(principal any) string {
	if s, ok := principal.(string); ok {
		return s
	}
	b, err := json.Marshal(principal)
	if err != nil {
		return fmt.Sprintf("%v", principal)
	}
	return string(b)
}

func getAWSPolicyStatementsWithUniqueAWSPrincipals(statements []Statement) []Statement {
	policyStatementMapFromCRD := map[string]struct{}{}
	policyStatementsFromCRD := []Statement{}
	for _, s := range statements {
		awsPrincipal, ok := s.Principal["AWS"]
		if !ok {
			policyStatementsFromCRD = append(policyStatementsFromCRD, s)
			continue
		}

		key := awsPrincipalKey(awsPrincipal)
		if _, ok := policyStatementMapFromCRD[key]; ok {
			continue
		}

		policyStatementMapFromCRD[key] = struct{}{}
		policyStatementsFromCRD = append(policyStatementsFromCRD, s)
	}

	return policyStatementsFromCRD
}

type awsPrincipalStatus string

const (
	undeleted                    awsPrincipalStatus = "undeleted"
	deletedAndPrincipalsNotEmpty awsPrincipalStatus = "deletedAndPrincipalsNotEmpty"
)

func getAWSPrincipals(statementsFromCRD []Statement, awsPrincipalsFromAWS map[string]struct{}) (awsPrincipals map[string]struct{}, validAWSPrincipals map[string]awsPrincipalStatus) {
	policyStatementsFromCRD := getAWSPolicyStatementsWithUniqueAWSPrincipals(statementsFromCRD)
	awsPrincipalsFromCRD := map[string]struct{}{}
	validAWSPrincipalsFromCRD := map[string]awsPrincipalStatus{}

	for _, s := range policyStatementsFromCRD {
		if len(s.Principal) == 0 {
			continue
		}

		principal, ok := s.Principal["AWS"]
		if !ok {
			continue
		}

		key := awsPrincipalKey(principal)
		awsPrincipalsFromCRD[key] = struct{}{}

		if _, ok := awsPrincipalsFromAWS[key]; !ok {
			delete(s.Principal, "AWS")

			if len(s.Principal) > 0 {
				validAWSPrincipalsFromCRD[key] = deletedAndPrincipalsNotEmpty
			}
			continue
		}

		validAWSPrincipalsFromCRD[key] = undeleted
	}

	return awsPrincipalsFromCRD, validAWSPrincipalsFromCRD
}

// getAliCloudAssumeRolePrincipals extracts principals from Aliyun policy statements
//func getAliCloudAssumeRolePrincipals(statements []Statement) map[interface{}]struct{} {
//	principals := make(map[interface{}]struct{})
//	for _, s := range statements {
//		if principal, ok := s.Principal["Service"]; ok {
//			principals[principal] = struct{}{}
//		}
//		if principal, ok := s.Principal["RAM"]; ok {
//			principals[principal] = struct{}{}
//		}
//	}
//	return principals
//}

func calcAssumeRolePolicy(policyFromCRD, lastPolicyFromCRD, policyFromAWS AssumeRolePolicy) AssumeRolePolicy {
	awsPrincipalsFromAWS := map[string]struct{}{}
	for _, s := range policyFromAWS.Statement {
		if len(s.Principal) == 0 {
			continue
		}
		principal, ok := s.Principal["AWS"]
		if !ok {
			continue
		}

		awsPrincipalsFromAWS[awsPrincipalKey(principal)] = struct{}{}
	}

	lastAWSPrincipalsFromCRD, lastValidAWSPrincipalsFromCRD := getAWSPrincipals(lastPolicyFromCRD.Statement, awsPrincipalsFromAWS)

	policyStatementsFromCRD := getAWSPolicyStatementsWithUniqueAWSPrincipals(policyFromCRD.Statement)

	statementsFromCRD := []Statement{}
	for _, s := range policyStatementsFromCRD {
		if len(s.Principal) == 0 {
			continue
		}

		awsPrincipal, ok := s.Principal["AWS"]
		if !ok {
			statementsFromCRD = append(statementsFromCRD, s)
			continue
		}

		key := awsPrincipalKey(awsPrincipal)
		if _, ok := lastAWSPrincipalsFromCRD[key]; !ok {
			statementsFromCRD = append(statementsFromCRD, s)
			continue
		}

		status, ok := lastValidAWSPrincipalsFromCRD[key]
		if !ok {
			continue
		}

		switch status {
		case deletedAndPrincipalsNotEmpty:
			delete(s.Principal, "AWS")
			if len(s.Principal) > 0 {
				statementsFromCRD = append(statementsFromCRD, s)
			}
			continue
		case undeleted:
			statementsFromCRD = append(statementsFromCRD, s)
			continue
		default:
			panic(fmt.Sprintf("awsPrincipal in lastValidAWSPrincipalsFromCRD with unknown status: %s", status))
		}
	}

	newPolicyFromCRD := policyFromCRD
	newPolicyFromCRD.Statement = statementsFromCRD
	return newPolicyFromCRD
}

// calcAssumeRolePolicyAliCloud
//func calcAssumeRolePolicyAliCloud(policyFromCRD, lastPolicyFromCRD, policyFromAliCloud AssumeRolePolicy) AssumeRolePolicy {
//	return policyFromAliCloud
//}

func calcAssumeRoleAccountGCP(accountsFromCRD, lastAccountsFromCRD, accountsFromGCP []string) []string {
	accountsFromGCPMap := map[string]struct{}{}
	for _, a := range accountsFromGCP {
		if len(a) == 0 {
			continue
		}

		accountsFromGCPMap[a] = struct{}{}
	}

	lastAccountsFromCRDMap := map[string]struct{}{}
	lastValidAccountsFromCRDMap := map[string]struct{}{}
	for _, a := range lastAccountsFromCRD {
		if len(a) == 0 {
			continue
		}

		lastAccountsFromCRDMap[a] = struct{}{}

		if _, ok := accountsFromGCPMap[a]; !ok {
			continue
		}

		lastValidAccountsFromCRDMap[a] = struct{}{}
	}

	newAccountsFromCRD := []string{}
	for _, a := range accountsFromCRD {
		if len(a) == 0 {
			continue
		}

		if _, ok := lastAccountsFromCRDMap[a]; !ok {
			newAccountsFromCRD = append(newAccountsFromCRD, a)
			continue
		}

		if _, ok := lastValidAccountsFromCRDMap[a]; ok {
			newAccountsFromCRD = append(newAccountsFromCRD, a)
		}
	}

	return newAccountsFromCRD
}