//nolint:gocritic // Role lifecycle boundaries deliberately receive immutable activation contracts by value.
package byoc

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/url"
	"slices"
	"strings"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/iam"
	iamtypes "github.com/aws/aws-sdk-go-v2/service/iam/types"
)

const customerDataplaneManagerRoleName = "tidbcloud-dataplane-manager-assumed-role"

func customerDataplaneManagerRoleARN(accountID string) string {
	return "arn:aws:iam::" + accountID + ":role/" + customerDataplaneManagerRoleName
}

func byocControlPlaneRoleName(session string) string {
	// Session names are validated by the top-level cp9 contract as DNS labels
	// of at most 32 characters, so this exact prefix always fits IAM's 64-byte
	// role-name limit without truncating or hashing away the ownership boundary.
	return session + "-byoc-dataplane-manager"
}

func byocControlPlaneRoleARN(input Input) string {
	return fmt.Sprintf("arn:aws:iam::%s:role/%s", input.ControlPlaneAccountID, byocControlPlaneRoleName(input.Session))
}

func roleNameFromARN(arn string) string {
	marker := ":role/"
	_, after, ok := strings.Cut(arn, marker)
	if !ok {
		return ""
	}
	name := after
	if slash := strings.LastIndex(name, "/"); slash >= 0 {
		name = name[slash+1:]
	}
	return name
}

func (d *AWSDriver) applyControlPlaneRole(ctx context.Context, input Input) (*StageResult, error) {
	roleName := byocControlPlaneRoleName(input.Session)
	roleARN := byocControlPlaneRoleARN(input)
	remoteRuntimeRoleName := roleNameFromARN(input.RemoteRuntimeRoleARN)
	if remoteRuntimeRoleName == "" {
		return nil, fmt.Errorf("remote runtime role ARN is invalid: %s", input.RemoteRuntimeRoleARN)
	}
	trust := marshalPolicyDocument(map[string]any{
		"Version": "2012-10-17",
		"Statement": []any{map[string]any{
			"Sid":       "RemoteRuntimeAssumeBYOCDataplaneManager",
			"Effect":    "Allow",
			"Principal": map[string]any{"AWS": input.RemoteRuntimeRoleARN},
			"Action":    "sts:AssumeRole",
		}},
	})
	role, err := d.readIAMRole(ctx, input, roleName)
	if err != nil && !isAWSNotFound(err) {
		return nil, fmt.Errorf("observe control-plane BYOC DPM role: %w", err)
	}
	if role == nil {
		clients, clientErr := d.clients(ctx, input.ControlPlaneProfile, input.Region)
		if clientErr != nil {
			return nil, clientErr
		}
		output, createErr := clients.iam.CreateRole(ctx, &iam.CreateRoleInput{
			RoleName:                 awssdk.String(roleName),
			Description:              awssdk.String("CP9 session-isolated BYOC dataplane-manager hop"),
			AssumeRolePolicyDocument: awssdk.String(trust),
			Tags: []iamtypes.Tag{
				{Key: awssdk.String("cp9-owner"), Value: awssdk.String(input.Owner)},
				{Key: awssdk.String("cp9-session"), Value: awssdk.String(input.Session)},
				{Key: awssdk.String("cp9-persona"), Value: awssdk.String("byoc")},
			},
		})
		if createErr != nil {
			return nil, fmt.Errorf("create control-plane BYOC DPM role: %w", createErr)
		}
		if output.Role == nil {
			return nil, fmt.Errorf("create control-plane BYOC DPM role returned no role")
		}
		role = iamRoleFromSDK(*output.Role)
	} else if err := validateOwnedIAMRole(role, roleName, roleARN, input); err != nil {
		return nil, err
	}
	if role.ARN != "" && role.ARN != roleARN {
		return nil, fmt.Errorf("control-plane BYOC DPM role ARN %q, want %q", role.ARN, roleARN)
	}
	clients, err := d.clients(ctx, input.ControlPlaneProfile, input.Region)
	if err != nil {
		return nil, err
	}
	if _, err := clients.iam.UpdateAssumeRolePolicy(ctx, &iam.UpdateAssumeRolePolicyInput{
		RoleName: awssdk.String(roleName), PolicyDocument: awssdk.String(trust),
	}); err != nil {
		return nil, fmt.Errorf("reconcile control-plane BYOC DPM role trust: %w", err)
	}

	customerRoleARN := customerDataplaneManagerRoleARN(input.CustomerAccountID)
	if _, err := clients.iam.PutRolePolicy(ctx, &iam.PutRolePolicyInput{
		RoleName: awssdk.String(roleName), PolicyName: awssdk.String("cp9-byoc-assume-customer"),
		PolicyDocument: awssdk.String(marshalAssumeRolePolicy(customerRoleARN, "AssumeCustomerDataplaneManager")),
	}); err != nil {
		return nil, fmt.Errorf("allow control-plane BYOC DPM role to assume customer role: %w", err)
	}
	controlPlaneNetworkPolicy := marshalPolicyDocument(map[string]any{
		"Version": "2012-10-17",
		"Statement": []any{map[string]any{
			"Sid":    "ManageBYOCControlPlanePrivateLink",
			"Effect": "Allow",
			"Action": []any{
				"ec2:CreateVpcEndpoint",
				"ec2:CreateTags",
				"ec2:DeleteTags",
				"ec2:DeleteVpcEndpoints",
				"ec2:Describe*",
				"ec2:ModifyVpcEndpoint",
			},
			"Resource": "*",
		}},
	})
	if _, err := clients.iam.PutRolePolicy(ctx, &iam.PutRolePolicyInput{
		RoleName: awssdk.String(roleName), PolicyName: awssdk.String("cp9-byoc-control-plane-network"),
		PolicyDocument: awssdk.String(controlPlaneNetworkPolicy),
	}); err != nil {
		return nil, fmt.Errorf("allow control-plane BYOC DPM role to manage PrivateLink endpoint: %w", err)
	}
	remotePolicyName := remoteBYOCAssumePolicyName(input.Session)
	if _, err := clients.iam.PutRolePolicy(ctx, &iam.PutRolePolicyInput{
		RoleName: awssdk.String(remoteRuntimeRoleName), PolicyName: awssdk.String(remotePolicyName),
		PolicyDocument: awssdk.String(marshalAssumeRolePolicy(roleARN, "AssumeSessionBYOCDataplaneManager")),
	}); err != nil {
		return nil, fmt.Errorf("allow remote runtime to assume control-plane BYOC DPM role: %w", err)
	}
	return &StageResult{Artifacts: map[string]string{
		"role_name":                   roleName,
		"role_arn":                    roleARN,
		"customer_role_arn":           customerRoleARN,
		"remote_runtime_role_name":    remoteRuntimeRoleName,
		"remote_inline_policy_name":   remotePolicyName,
		"customer_inline_policy_name": "cp9-byoc-assume-customer",
		"network_inline_policy_name":  "cp9-byoc-control-plane-network",
	}}, nil
}

type iamTag struct {
	Key   string
	Value string
}

type iamRole struct {
	RoleName                 string
	ARN                      string
	Path                     string
	MaxSessionDuration       int
	AssumeRolePolicyDocument any
	PermissionsBoundary      any
	Tags                     []iamTag
}

func (d *AWSDriver) readIAMRole(ctx context.Context, input Input, roleName string) (*iamRole, error) {
	return d.readIAMRoleWithProfile(ctx, input, input.ControlPlaneProfile, roleName)
}

func (d *AWSDriver) readIAMRoleWithProfile(
	ctx context.Context,
	input Input,
	profile string,
	roleName string,
) (*iamRole, error) {
	clients, err := d.clients(ctx, profile, input.Region)
	if err != nil {
		return nil, err
	}
	output, err := clients.iam.GetRole(ctx, &iam.GetRoleInput{RoleName: awssdk.String(roleName)})
	if err != nil {
		return nil, err
	}
	if output.Role == nil {
		return nil, fmt.Errorf("IAM role %s returned no role", roleName)
	}
	return iamRoleFromSDK(*output.Role), nil
}

func iamRoleFromSDK(role iamtypes.Role) *iamRole {
	tags := make([]iamTag, 0, len(role.Tags))
	for _, tag := range role.Tags {
		tags = append(tags, iamTag{Key: awssdk.ToString(tag.Key), Value: awssdk.ToString(tag.Value)})
	}
	var permissionsBoundary any
	if role.PermissionsBoundary != nil {
		permissionsBoundary = role.PermissionsBoundary
	}
	return &iamRole{
		RoleName:                 awssdk.ToString(role.RoleName),
		ARN:                      awssdk.ToString(role.Arn),
		Path:                     awssdk.ToString(role.Path),
		MaxSessionDuration:       int(awssdk.ToInt32(role.MaxSessionDuration)),
		AssumeRolePolicyDocument: decodeIAMPolicyDocument(role.AssumeRolePolicyDocument),
		PermissionsBoundary:      permissionsBoundary,
		Tags:                     tags,
	}
}

func decodeIAMPolicyDocument(document *string) any {
	decoded, err := url.QueryUnescape(awssdk.ToString(document))
	if err != nil {
		return nil
	}
	var policy any
	if json.Unmarshal([]byte(decoded), &policy) != nil {
		return nil
	}
	return policy
}

func validateOwnedIAMRole(role *iamRole, roleName, roleARN string, input Input) error {
	if role == nil || role.RoleName != roleName || role.ARN != roleARN {
		return fmt.Errorf("existing control-plane BYOC DPM role %s has an unexpected identity", roleName)
	}
	tags := make(map[string]string, len(role.Tags))
	for _, tag := range role.Tags {
		tags[tag.Key] = tag.Value
	}
	if tags["cp9-session"] != input.Session || tags["cp9-owner"] != input.Owner || tags["cp9-persona"] != "byoc" {
		return fmt.Errorf("existing control-plane BYOC DPM role %s is not owned by this session", roleName)
	}
	return nil
}

func marshalAssumeRolePolicy(resource, sid string) string {
	return marshalPolicyDocument(map[string]any{
		"Version": "2012-10-17",
		"Statement": []any{map[string]any{
			"Sid":      sid,
			"Effect":   "Allow",
			"Action":   "sts:AssumeRole",
			"Resource": resource,
		}},
	})
}

func marshalPolicyDocument(value any) string {
	data, _ := json.Marshal(value)
	return string(data)
}

func remoteBYOCAssumePolicyName(session string) string {
	digest := sha256.Sum256([]byte(session))
	return "cp9-byoc-assume-" + hex.EncodeToString(digest[:6])
}

func (d *AWSDriver) revalidateControlPlaneRole(ctx context.Context, input Input, record *StageRecord) (*StageResult, error) {
	if record == nil || record.Artifacts["role_arn"] != byocControlPlaneRoleARN(input) {
		return nil, ErrNeedsReconcile
	}
	roleName := byocControlPlaneRoleName(input.Session)
	role, err := d.readIAMRole(ctx, input, roleName)
	if err != nil {
		if isAWSNotFound(err) {
			return nil, ErrNeedsReconcile
		}
		return nil, err
	}
	clients, err := d.clients(ctx, input.ControlPlaneProfile, input.Region)
	if err != nil {
		return nil, err
	}
	if err := validateOwnedIAMRole(role, roleName, byocControlPlaneRoleARN(input), input); err != nil {
		return nil, err
	}
	if !trustPolicyExactlyAllowsPrincipal(role.AssumeRolePolicyDocument, input.RemoteRuntimeRoleARN) {
		return nil, ErrNeedsReconcile
	}
	checks := []struct {
		roleName   string
		policyName string
		resource   string
	}{
		{roleName: roleName, policyName: "cp9-byoc-assume-customer", resource: customerDataplaneManagerRoleARN(input.CustomerAccountID)},
		{
			roleName:   roleNameFromARN(input.RemoteRuntimeRoleARN),
			policyName: remoteBYOCAssumePolicyName(input.Session),
			resource:   byocControlPlaneRoleARN(input),
		},
	}
	for _, check := range checks {
		output, err := clients.iam.GetRolePolicy(ctx, &iam.GetRolePolicyInput{
			RoleName: awssdk.String(check.roleName), PolicyName: awssdk.String(check.policyName),
		})
		if err != nil {
			if isAWSNotFound(err) {
				return nil, ErrNeedsReconcile
			}
			return nil, err
		}
		if !policyDocumentAllowsResource(decodeIAMPolicyDocument(output.PolicyDocument), check.resource) {
			return nil, ErrNeedsReconcile
		}
	}
	networkOutput, err := clients.iam.GetRolePolicy(ctx, &iam.GetRolePolicyInput{
		RoleName: awssdk.String(roleName), PolicyName: awssdk.String("cp9-byoc-control-plane-network"),
	})
	if err != nil {
		if isAWSNotFound(err) {
			return nil, ErrNeedsReconcile
		}
		return nil, err
	}
	if !policyDocumentHasActions(
		decodeIAMPolicyDocument(networkOutput.PolicyDocument),
		"ec2:CreateVpcEndpoint",
		"ec2:CreateTags",
		"ec2:DeleteTags",
		"ec2:DeleteVpcEndpoints",
		"ec2:Describe*",
		"ec2:ModifyVpcEndpoint",
	) {
		return nil, ErrNeedsReconcile
	}
	return &StageResult{Artifacts: cloneArtifacts(record.Artifacts)}, nil
}

func policyDocumentAllowsResource(document any, resource string) bool {
	return policyDocumentExactlyMatches(document, []string{"sts:AssumeRole"}, []string{resource})
}

func policyDocumentHasActions(document any, actions ...string) bool {
	return policyDocumentExactlyMatches(document, actions, []string{"*"})
}

func policyDocumentExactlyMatches(document any, actions, resources []string) bool {
	if encoded, ok := document.(string); ok {
		if json.Unmarshal([]byte(encoded), &document) != nil {
			return false
		}
	}
	documentMap, ok := document.(map[string]any)
	if !ok {
		return false
	}
	statements := policyStatements(documentMap["Statement"])
	if len(statements) != 1 || statements[0]["Effect"] != "Allow" {
		return false
	}
	statement := statements[0]
	for _, forbidden := range []string{"Condition", "NotAction", "NotResource", "Principal", "NotPrincipal"} {
		if _, exists := statement[forbidden]; exists {
			return false
		}
	}
	return equalPolicyValues(policyStringValues(statement["Action"]), actions) &&
		equalPolicyValues(policyStringValues(statement["Resource"]), resources)
}

func policyStatements(value any) []map[string]any {
	switch typed := value.(type) {
	case map[string]any:
		return []map[string]any{typed}
	case []any:
		result := make([]map[string]any, 0, len(typed))
		for _, item := range typed {
			statement, ok := item.(map[string]any)
			if !ok {
				return nil
			}
			result = append(result, statement)
		}
		return result
	default:
		return nil
	}
}

func policyStringValues(value any) []string {
	switch typed := value.(type) {
	case string:
		return []string{typed}
	case []any:
		result := make([]string, 0, len(typed))
		for _, item := range typed {
			text, ok := item.(string)
			if !ok {
				return nil
			}
			result = append(result, text)
		}
		return result
	default:
		return nil
	}
}

func equalPolicyValues(actual, expected []string) bool {
	actual = append([]string(nil), actual...)
	expected = append([]string(nil), expected...)
	slices.Sort(actual)
	slices.Sort(expected)
	return slices.Equal(actual, expected)
}

func trustPolicyExactlyAllowsPrincipal(document any, principal string) bool {
	if encoded, ok := document.(string); ok {
		if json.Unmarshal([]byte(encoded), &document) != nil {
			return false
		}
	}
	documentMap, ok := document.(map[string]any)
	if !ok {
		return false
	}
	statements := policyStatements(documentMap["Statement"])
	if len(statements) != 1 || statements[0]["Effect"] != "Allow" {
		return false
	}
	statement := statements[0]
	if !equalPolicyValues(policyStringValues(statement["Action"]), []string{"sts:AssumeRole"}) {
		return false
	}
	for _, forbidden := range []string{"Condition", "NotAction", "Resource", "NotResource", "NotPrincipal"} {
		if _, exists := statement[forbidden]; exists {
			return false
		}
	}
	principalMap, ok := statement["Principal"].(map[string]any)
	if !ok || len(principalMap) != 1 {
		return false
	}
	return equalPolicyValues(policyStringValues(principalMap["AWS"]), []string{principal})
}

// customerDPMTrustAllowsControlPlane accepts additive account-specific trust
// statements while rejecting wildcard or unconditional cross-account access.
// Shared retained setup roles may serve multiple control planes, so every
// statement must remain restricted to one account's dataplane-manager roles.
func customerDPMTrustAllowsControlPlane(document any, controlPlaneAccountID string) bool {
	if encoded, ok := document.(string); ok {
		if json.Unmarshal([]byte(encoded), &document) != nil {
			return false
		}
	}
	documentMap, ok := document.(map[string]any)
	if !ok {
		return false
	}
	statements := policyStatements(documentMap["Statement"])
	if len(statements) == 0 {
		return false
	}
	found := false
	for _, statement := range statements {
		accountID, ok := customerDPMTrustStatementAccount(statement)
		if !ok {
			return false
		}
		found = found || accountID == controlPlaneAccountID
	}
	return found
}

func customerDPMTrustExactlyAllowsControlPlane(document any, controlPlaneAccountID string) bool {
	if encoded, ok := document.(string); ok {
		if json.Unmarshal([]byte(encoded), &document) != nil {
			return false
		}
	}
	documentMap, ok := document.(map[string]any)
	if !ok {
		return false
	}
	statements := policyStatements(documentMap["Statement"])
	if len(statements) != 1 {
		return false
	}
	accountID, ok := customerDPMTrustStatementAccount(statements[0])
	return ok && accountID == controlPlaneAccountID
}

func customerDPMTrustStatementAccount(statement map[string]any) (string, bool) {
	for key := range statement {
		switch key {
		case "Sid", "Effect", "Principal", "Action", "Condition":
		default:
			return "", false
		}
	}
	if statement["Effect"] != "Allow" ||
		!equalPolicyValues(policyStringValues(statement["Action"]), []string{"sts:AssumeRole"}) {
		return "", false
	}
	principal, ok := statement["Principal"].(map[string]any)
	if !ok || len(principal) != 1 {
		return "", false
	}
	principals := policyStringValues(principal["AWS"])
	if len(principals) != 1 {
		return "", false
	}
	const prefix = "arn:aws:iam::"
	const suffix = ":root"
	principalARN := principals[0]
	if !strings.HasPrefix(principalARN, prefix) || !strings.HasSuffix(principalARN, suffix) {
		return "", false
	}
	accountID := strings.TrimSuffix(strings.TrimPrefix(principalARN, prefix), suffix)
	if len(accountID) != 12 || strings.Trim(accountID, "0123456789") != "" {
		return "", false
	}
	condition, ok := statement["Condition"].(map[string]any)
	if !ok || len(condition) != 1 {
		return "", false
	}
	stringLike, ok := condition["StringLike"].(map[string]any)
	return accountID, ok && len(stringLike) == 1 && equalPolicyValues(
		policyStringValues(stringLike["aws:PrincipalArn"]),
		[]string{"arn:aws:iam::" + accountID + ":role/*-dataplane-manager"},
	)
}