package impl

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

	"github.com/AlekSi/pointer"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/iam/iamiface"

	"github.com/tidbcloud/dataflow-service/app/models/backup"
	"github.com/tidbcloud/dataflow-service/app/models/bucket"
	restoreaccess "github.com/tidbcloud/dataflow-service/app/models/classic_restore_access"
	"github.com/tidbcloud/dataflow-service/app/models/domain"
	"github.com/tidbcloud/dataflow-service/app/repo"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
	dfsaws "github.com/tidbcloud/dataflow-service/pkg/aws"
)

const (
	roleNamePrefix          = "tidb-classic-restore-"
	rolePath                = "/tidb-cloud/classic-restore/"
	policyName              = "tidb-classic-restore-source-read"
	permissionsBoundaryName = "tidb-cloud-classic-restore-permissions-boundary"
	expiresAtTag            = "tidb-cloud-expires-at"
	requestHashTag          = "tidb-cloud-request-hash"
	roleTTL                 = 7 * 24 * time.Hour
)

type model struct {
	backupModel backup.SnapshotBackupModel
	bucketModel bucket.BucketModelV2
	accessRepo  repo.ClassicRestoreAccessRepo
	iamProxy    dfsaws.AWSIAMClientProxy
}

func NewModel(
	backupModel backup.SnapshotBackupModel,
	bucketModel bucket.BucketModelV2,
	accessRepo repo.ClassicRestoreAccessRepo,
	iamProxy dfsaws.AWSIAMClientProxy,
) restoreaccess.Model {
	return &model{backupModel: backupModel, bucketModel: bucketModel, accessRepo: accessRepo, iamProxy: iamProxy}
}

func (m *model) Prepare(ctx context.Context, req *restoreaccess.PrepareReq) (*restoreaccess.PrepareResp, error) {
	if err := validateRoleARN(req.TrustedRoleARN); err != nil {
		return nil, err
	}
	bucketEntity, prefix, err := m.resolveSource(ctx, req.BackupID, req.BucketID)
	if err != nil {
		return nil, err
	}
	if bucketEntity.AwsKmsArn == "" {
		return nil, fmt.Errorf("backup bucket %d does not use an AWS KMS key", req.BucketID)
	}
	roleName, externalID := accessIdentifiers(req.RestoreID, req.BucketID)
	trustPolicy, err := marshalPolicy(map[string]any{
		"Version": "2012-10-17",
		"Statement": []any{map[string]any{
			"Effect": "Allow", "Principal": map[string]string{"AWS": req.TrustedRoleARN},
			"Action":    "sts:AssumeRole",
			"Condition": map[string]any{"StringEquals": map[string]string{"sts:ExternalId": externalID}},
		}},
	})
	if err != nil {
		return nil, err
	}
	permissionsPolicy, err := sourcePermissionsPolicy(bucketEntity.Name, prefix, bucketEntity.AwsKmsArn)
	if err != nil {
		return nil, err
	}
	accountID := pointer.Get(bucketEntity.CloudAccountID)
	partition, err := awsPartition(bucketEntity.AwsKmsArn)
	if err != nil {
		return nil, err
	}
	boundaryARN := iamResourceARN(partition, accountID, "policy/"+permissionsBoundaryName)
	roleARN := iamResourceARN(partition, accountID, "role"+rolePath+roleName)
	requestHash, err := accessRequestHash(
		req, accountID, bucketEntity.Name, prefix, bucketEntity.AwsKmsArn, trustPolicy, permissionsPolicy, boundaryARN,
	)
	if err != nil {
		return nil, err
	}
	expiresAt := time.Now().UTC().Add(roleTTL)
	desiredAccess := &entity.ClassicRestoreAccess{
		RestoreID: req.RestoreID, BucketID: entity.ID(req.BucketID), BackupID: entity.ID(req.BackupID),
		AWSAccountID: accountID, RoleName: roleName, RoleARN: roleARN, RequestHash: requestHash, ExpiresAt: expiresAt,
	}
	access, rerr := m.accessRepo.GetOrCreate(ctx, desiredAccess)
	if rerr != nil {
		return nil, fmt.Errorf("persist classic restore access: %w", rerr)
	}
	if access.RequestHash != desiredAccess.RequestHash || access.BackupID != desiredAccess.BackupID ||
		access.AWSAccountID != desiredAccess.AWSAccountID || access.RoleName != desiredAccess.RoleName ||
		access.RoleARN != desiredAccess.RoleARN {
		return nil, fmt.Errorf("classic restore access parameters do not match the existing request")
	}
	expiresAt = access.ExpiresAt
	roleARN = access.RoleARN
	client, err := m.iamProxy.GetClient(ctx, accountID)
	if err != nil {
		return nil, fmt.Errorf("get IAM client: %w", err)
	}
	actualRoleARN, err := ensureRole(
		ctx, client, roleName, trustPolicy, permissionsPolicy, boundaryARN, requestHash, expiresAt,
	)
	if err != nil {
		return nil, err
	}
	if actualRoleARN != roleARN {
		return nil, fmt.Errorf("temporary restore role ARN does not match the persisted access record")
	}
	if err = client.WaitUntilRoleExistsWithContext(ctx, &iam.GetRoleInput{RoleName: aws.String(roleName)}); err != nil {
		return nil, fmt.Errorf("wait for temporary restore role: %w", err)
	}
	_, err = client.PutRolePolicyWithContext(ctx, &iam.PutRolePolicyInput{
		RoleName: aws.String(roleName), PolicyName: aws.String(policyName), PolicyDocument: aws.String(permissionsPolicy),
	})
	if err != nil {
		return nil, fmt.Errorf("put temporary restore role policy: %w", err)
	}
	return &restoreaccess.PrepareResp{SourceRoleARN: roleARN, ExternalID: externalID}, nil
}

func (m *model) Cleanup(ctx context.Context, req *restoreaccess.CleanupReq) error {
	access, rerr := m.accessRepo.Get(ctx, req.RestoreID, req.BucketID)
	if rerr != nil {
		return fmt.Errorf("get classic restore access: %w", rerr)
	}
	if access == nil {
		return nil
	}
	client, err := m.iamProxy.GetClient(ctx, access.AWSAccountID)
	if err != nil {
		return fmt.Errorf("get IAM client: %w", err)
	}
	if err := deleteRole(ctx, client, access.RoleName); err != nil {
		return err
	}
	if rerr := m.accessRepo.Delete(ctx, access.ID); rerr != nil {
		return fmt.Errorf("delete classic restore access record: %w", rerr)
	}
	return nil
}

func (m *model) ReconcileExpired(ctx context.Context) error {
	accesses, rerr := m.accessRepo.ListExpired(ctx, time.Now().UTC())
	if rerr != nil {
		return fmt.Errorf("list expired classic restore accesses: %w", rerr)
	}
	var result error
	for i := range accesses {
		access := &accesses[i]
		client, err := m.iamProxy.GetClient(ctx, access.AWSAccountID)
		if err != nil {
			result = errors.Join(result, fmt.Errorf("get IAM client for account %s: %w", access.AWSAccountID, err))
			continue
		}
		if err := deleteRole(ctx, client, access.RoleName); err != nil {
			result = errors.Join(result, err)
			continue
		}
		if rerr := m.accessRepo.Delete(ctx, access.ID); rerr != nil {
			result = errors.Join(result, fmt.Errorf("delete expired classic restore access record: %w", rerr))
		}
	}
	return result
}

func (m *model) resolveSource(ctx context.Context, backupID, bucketID uint64) (*domainBucket, string, error) {
	backupResp, berr := m.backupModel.Query(ctx, &domain.QueryBackupReq{Filter: &domain.ListBackupFilter{BackupIDs: []uint64{backupID}}})
	if berr != nil {
		return nil, "", fmt.Errorf("get backup: %w", berr)
	}
	if backupResp == nil || backupResp.BackupEntity == nil {
		return nil, "", fmt.Errorf("backup %d not found", backupID)
	}
	if backupResp.BackupEntity.BucketSettingSnapshot == nil {
		return nil, "", fmt.Errorf("backup %d has no bucket snapshot", backupID)
	}
	belongs := false
	for _, setting := range backupResp.BackupEntity.BucketSettingSnapshot.Buckets {
		if setting.BucketId == bucketID {
			belongs = true
			break
		}
	}
	if !belongs {
		return nil, "", fmt.Errorf("bucket %d does not belong to backup %d", bucketID, backupID)
	}
	u, err := url.Parse(backupResp.BackupEntity.Storage)
	if err != nil || !strings.EqualFold(u.Scheme, "s3") {
		return nil, "", fmt.Errorf("backup %d has invalid S3 storage", backupID)
	}
	bucketResp, berr := m.bucketModel.GetBucket(ctx, &domain.GetBucketReq{
		BucketID: bucketID, InfraCloudProvider: pointer.To(domain.InfraProviderType(domain.CloudProviderTypeAWS)),
	})
	if berr != nil {
		return nil, "", fmt.Errorf("get backup bucket: %w", berr)
	}
	if bucketResp == nil || bucketResp.Bucket == nil || bucketResp.Bucket.CloudAccountID == nil ||
		*bucketResp.Bucket.CloudAccountID == "" {
		return nil, "", fmt.Errorf("backup bucket %d has no AWS account", bucketID)
	}
	return &domainBucket{
		Name: bucketResp.Bucket.Name, AwsKmsArn: bucketResp.Bucket.AwsKmsArn, CloudAccountID: bucketResp.Bucket.CloudAccountID,
	}, strings.TrimPrefix(u.Path, "/"), nil
}

type domainBucket struct {
	Name           string
	AwsKmsArn      string
	CloudAccountID *string
}

func ensureRole(
	ctx context.Context,
	client iamiface.IAMAPI,
	roleName, trustPolicy, permissionsPolicy, boundaryARN, requestHash string,
	expiresAt time.Time,
) (string, error) {
	created, err := client.CreateRoleWithContext(ctx, &iam.CreateRoleInput{
		RoleName: aws.String(roleName), Path: aws.String(rolePath), AssumeRolePolicyDocument: aws.String(trustPolicy),
		Description:         aws.String("Temporary source access for a TiDB Cloud classic restore"),
		PermissionsBoundary: aws.String(boundaryARN),
		Tags: []*iam.Tag{
			{Key: aws.String("tidb-cloud-purpose"), Value: aws.String("classic-restore")},
			{Key: aws.String(expiresAtTag), Value: aws.String(expiresAt.Format(time.RFC3339))},
			{Key: aws.String(requestHashTag), Value: aws.String(requestHash)},
		},
	})
	if err == nil {
		return aws.StringValue(created.Role.Arn), nil
	}
	if awsErr, ok := err.(awserr.Error); !ok || awsErr.Code() != iam.ErrCodeEntityAlreadyExistsException {
		return "", fmt.Errorf("create temporary restore role: %w", err)
	}
	existing, err := client.GetRoleWithContext(ctx, &iam.GetRoleInput{RoleName: aws.String(roleName)})
	if err != nil {
		return "", fmt.Errorf("get temporary restore role: %w", err)
	}
	if err := validateExistingRole(existing.Role, trustPolicy, boundaryARN); err != nil {
		return "", err
	}
	existingPolicy, err := client.GetRolePolicyWithContext(ctx, &iam.GetRolePolicyInput{
		RoleName: aws.String(roleName), PolicyName: aws.String(policyName),
	})
	if err != nil && !isNoSuchEntity(err) {
		return "", fmt.Errorf("get temporary restore role policy: %w", err)
	}
	if err == nil {
		matches, compareErr := policiesEqual(permissionsPolicy, aws.StringValue(existingPolicy.PolicyDocument))
		if compareErr != nil {
			return "", fmt.Errorf("compare temporary restore role policy: %w", compareErr)
		}
		if !matches {
			return "", fmt.Errorf("temporary restore role policy does not match the existing request")
		}
	}
	if _, err = client.TagRoleWithContext(ctx, &iam.TagRoleInput{RoleName: aws.String(roleName), Tags: []*iam.Tag{
		{Key: aws.String(expiresAtTag), Value: aws.String(expiresAt.Format(time.RFC3339))},
		{Key: aws.String(requestHashTag), Value: aws.String(requestHash)},
	}}); err != nil {
		return "", fmt.Errorf("refresh temporary restore role expiry: %w", err)
	}
	return aws.StringValue(existing.Role.Arn), nil
}

func deleteRole(ctx context.Context, client iamiface.IAMAPI, roleName string) error {
	if _, err := client.DeleteRolePolicyWithContext(ctx, &iam.DeleteRolePolicyInput{
		RoleName: aws.String(roleName), PolicyName: aws.String(policyName),
	}); err != nil && !isNoSuchEntity(err) {
		return fmt.Errorf("delete temporary restore role policy: %w", err)
	}
	if _, err := client.DeleteRoleWithContext(ctx, &iam.DeleteRoleInput{RoleName: aws.String(roleName)}); err != nil && !isNoSuchEntity(err) {
		return fmt.Errorf("delete temporary restore role: %w", err)
	}
	return nil
}

func sourcePermissionsPolicy(bucketName, prefix, kmsARN string) (string, error) {
	objectARN := fmt.Sprintf("arn:aws:s3:::%s/*", bucketName)
	prefixValues := []string{"*"}
	if prefix != "" {
		objectARN = fmt.Sprintf("arn:aws:s3:::%s/%s/*", bucketName, strings.TrimSuffix(prefix, "/"))
		prefixValues = []string{prefix, strings.TrimSuffix(prefix, "/") + "/*"}
	}
	return marshalPolicy(map[string]any{
		"Version": "2012-10-17",
		"Statement": []any{
			map[string]any{"Effect": "Allow", "Action": []string{"s3:GetObject", "s3:GetObjectVersion"}, "Resource": objectARN},
			map[string]any{"Effect": "Allow", "Action": "s3:ListBucket", "Resource": fmt.Sprintf("arn:aws:s3:::%s", bucketName), "Condition": map[string]any{"StringLike": map[string]any{"s3:prefix": prefixValues}}},
			map[string]any{"Effect": "Allow", "Action": "s3:GetBucketLocation", "Resource": fmt.Sprintf("arn:aws:s3:::%s", bucketName)},
			map[string]any{"Effect": "Allow", "Action": []string{"kms:Decrypt", "kms:DescribeKey"}, "Resource": kmsARN},
		},
	})
}

func accessRequestHash(
	req *restoreaccess.PrepareReq,
	accountID, bucketName, prefix, kmsARN, trustPolicy, permissionsPolicy, boundaryARN string,
) (string, error) {
	payload, err := json.Marshal(struct {
		BackupID          uint64 `json:"backup_id"`
		BucketID          uint64 `json:"bucket_id"`
		RestoreID         string `json:"restore_id"`
		TrustedRoleARN    string `json:"trusted_role_arn"`
		AWSAccountID      string `json:"aws_account_id"`
		BucketName        string `json:"bucket_name"`
		Prefix            string `json:"prefix"`
		KMSARN            string `json:"kms_arn"`
		TrustPolicy       string `json:"trust_policy"`
		PermissionsPolicy string `json:"permissions_policy"`
		BoundaryARN       string `json:"boundary_arn"`
	}{
		BackupID: req.BackupID, BucketID: req.BucketID, RestoreID: req.RestoreID, TrustedRoleARN: req.TrustedRoleARN,
		AWSAccountID: accountID, BucketName: bucketName, Prefix: prefix, KMSARN: kmsARN,
		TrustPolicy: trustPolicy, PermissionsPolicy: permissionsPolicy, BoundaryARN: boundaryARN,
	})
	if err != nil {
		return "", fmt.Errorf("marshal classic restore access fingerprint: %w", err)
	}
	sum := sha256.Sum256(payload)
	return hex.EncodeToString(sum[:]), nil
}

func awsPartition(resourceARN string) (string, error) {
	parsed, err := arn.Parse(resourceARN)
	if err != nil || parsed.Partition == "" || parsed.Service != "kms" || parsed.AccountID == "" {
		return "", fmt.Errorf("invalid AWS KMS ARN")
	}
	return parsed.Partition, nil
}

func iamResourceARN(partition, accountID, resource string) string {
	return arn.ARN{Partition: partition, Service: "iam", AccountID: accountID, Resource: resource}.String()
}

func validateExistingRole(role *iam.Role, trustPolicy, boundaryARN string) error {
	if role == nil {
		return fmt.Errorf("temporary restore role response is empty")
	}
	if aws.StringValue(role.Path) != rolePath {
		return fmt.Errorf("temporary restore role path does not match the existing request")
	}
	matches, err := policiesEqual(trustPolicy, aws.StringValue(role.AssumeRolePolicyDocument))
	if err != nil {
		return fmt.Errorf("compare temporary restore role trust policy: %w", err)
	}
	if !matches {
		return fmt.Errorf("temporary restore role trust policy does not match the existing request")
	}
	if role.PermissionsBoundary == nil ||
		aws.StringValue(role.PermissionsBoundary.PermissionsBoundaryArn) != boundaryARN {
		return fmt.Errorf("temporary restore role permissions boundary does not match the required boundary")
	}
	return nil
}

func policiesEqual(expected, urlEncodedActual string) (bool, error) {
	actual, err := url.QueryUnescape(urlEncodedActual)
	if err != nil {
		return false, err
	}
	normalize := func(document string) ([]byte, error) {
		var value any
		if err := json.Unmarshal([]byte(document), &value); err != nil {
			return nil, err
		}
		return json.Marshal(value)
	}
	expectedJSON, err := normalize(expected)
	if err != nil {
		return false, err
	}
	actualJSON, err := normalize(actual)
	if err != nil {
		return false, err
	}
	return string(expectedJSON) == string(actualJSON), nil
}

func accessIdentifiers(restoreID string, bucketID uint64) (string, string) {
	sum := sha256.Sum256([]byte(fmt.Sprintf("%s:%d", restoreID, bucketID)))
	token := hex.EncodeToString(sum[:16])
	return roleNamePrefix + token, "tidb-cloud-classic-restore-" + token
}

func validateRoleARN(value string) error {
	parsed, err := arn.Parse(value)
	if err != nil || parsed.Service != "iam" || parsed.AccountID == "" || !strings.HasPrefix(parsed.Resource, "role/") {
		return fmt.Errorf("invalid trusted IAM role ARN")
	}
	return nil
}

func marshalPolicy(policy any) (string, error) {
	bytes, err := json.Marshal(policy)
	if err != nil {
		return "", fmt.Errorf("marshal IAM policy: %w", err)
	}
	return string(bytes), nil
}

func isNoSuchEntity(err error) bool {
	awsErr, ok := err.(awserr.Error)
	return ok && awsErr.Code() == iam.ErrCodeNoSuchEntityException
}