package cmd

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	dynamodbtypes "github.com/aws/aws-sdk-go-v2/service/dynamodb/types"

	cpaws "github.com/tidbcloud/controlplane9/internal/aws"
)

const (
	purgeAccountLockTable     = "cp9-session-cleanup-lock"
	purgeAccountLockRegion    = "us-west-2"
	purgeAccountLockLease     = 2 * time.Minute
	purgeAccountLockHeartbeat = 30 * time.Second
	purgeAccountLockRetry     = 5 * time.Second

	purgeAccountLockAccessKeyEnv    = "CP9_CLEANUP_LOCK_AWS_ACCESS_KEY_ID"
	purgeAccountLockSecretKeyEnv    = "CP9_CLEANUP_LOCK_AWS_SECRET_ACCESS_KEY"
	purgeAccountLockSessionTokenEnv = "CP9_CLEANUP_LOCK_AWS_SESSION_TOKEN"
)

type purgeAccountLockAPI interface {
	UpdateItem(ctx context.Context, input *dynamodb.UpdateItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.UpdateItemOutput, error)
	DeleteItem(ctx context.Context, input *dynamodb.DeleteItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error)
}

type purgeAccountLease struct {
	client    purgeAccountLockAPI
	accountID string
	token     string
	ctx       context.Context
	cancel    context.CancelFunc
	stop      chan struct{}
	done      chan struct{}
	err       chan error
}

func acquirePurgeAccountLease(
	ctx context.Context,
	mode cpaws.PurgeMode,
	opCfg purgeOperatorConfig,
) (*purgeAccountLease, bool, error) {
	awsCfg, err := loadPurgeAccountLeaseAWSConfig(ctx)
	if err != nil {
		return nil, false, fmt.Errorf("load cleanup lease AWS config: %w", err)
	}
	client := dynamodb.NewFromConfig(awsCfg)
	token, err := newPurgeAccountLeaseToken()
	if err != nil {
		return nil, false, err
	}
	holder := purgeAccountLeaseHolder()
	for {
		acquired, err := tryAcquirePurgeAccountLease(ctx, client, opCfg.AccountID, token, holder)
		if err != nil {
			return nil, false, err
		}
		if acquired {
			leaseCtx, cancel := context.WithCancel(ctx)
			lease := &purgeAccountLease{
				client:    client,
				accountID: opCfg.AccountID,
				token:     token,
				ctx:       leaseCtx,
				cancel:    cancel,
				stop:      make(chan struct{}),
				done:      make(chan struct{}),
				err:       make(chan error, 1),
			}
			go lease.heartbeat(leaseCtx)
			return lease, false, nil
		}
		if mode == cpaws.PurgeModeBroad {
			return nil, true, nil
		}
		timer := time.NewTimer(purgeAccountLockRetry)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, false, fmt.Errorf("wait for cleanup account lease: %w", ctx.Err())
		case <-timer.C:
		}
	}
}

func loadPurgeAccountLeaseAWSConfig(ctx context.Context) (awssdk.Config, error) {
	accessKey := strings.TrimSpace(os.Getenv(purgeAccountLockAccessKeyEnv))
	secretKey := strings.TrimSpace(os.Getenv(purgeAccountLockSecretKeyEnv))
	sessionToken := strings.TrimSpace(os.Getenv(purgeAccountLockSessionTokenEnv))
	if accessKey == "" || secretKey == "" || sessionToken == "" {
		return awssdk.Config{}, fmt.Errorf("804 cleanup coordinator credentials are required")
	}
	cfg, err := awsconfig.LoadDefaultConfig(ctx,
		awsconfig.WithRegion(purgeAccountLockRegion),
		awsconfig.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(accessKey, secretKey, sessionToken)),
	)
	if err != nil {
		return awssdk.Config{}, fmt.Errorf("load 804 cleanup coordinator credentials: %w", err)
	}
	cpaws.ApplyCleanupRetryPolicy(&cfg)
	return cfg, nil
}

func tryAcquirePurgeAccountLease(
	ctx context.Context,
	client purgeAccountLockAPI,
	accountID, token, holder string,
) (bool, error) {
	now := time.Now().UTC()
	_, err := client.UpdateItem(ctx, &dynamodb.UpdateItemInput{
		TableName: awssdk.String(purgeAccountLockTable),
		Key: map[string]dynamodbtypes.AttributeValue{
			"account_id": &dynamodbtypes.AttributeValueMemberS{Value: accountID},
		},
		ConditionExpression: awssdk.String("attribute_not_exists(#account) OR #expires < :now"),
		UpdateExpression:    awssdk.String("SET #token = :token, #holder = :holder, #expires = :expires"),
		ExpressionAttributeNames: map[string]string{
			"#account": "account_id",
			"#expires": "expires_at",
			"#holder":  "holder",
			"#token":   "lease_token",
		},
		ExpressionAttributeValues: map[string]dynamodbtypes.AttributeValue{
			":expires": &dynamodbtypes.AttributeValueMemberN{Value: fmt.Sprint(now.Add(purgeAccountLockLease).Unix())},
			":holder":  &dynamodbtypes.AttributeValueMemberS{Value: holder},
			":now":     &dynamodbtypes.AttributeValueMemberN{Value: fmt.Sprint(now.Unix())},
			":token":   &dynamodbtypes.AttributeValueMemberS{Value: token},
		},
	})
	if err == nil {
		return true, nil
	}
	if purgeAccountLockConditionalFailure(err) {
		return false, nil
	}
	return false, fmt.Errorf("acquire cleanup account lease in %s/%s: %w",
		purgeAccountLockRegion, purgeAccountLockTable, err)
}

func (l *purgeAccountLease) heartbeat(ctx context.Context) {
	defer close(l.done)
	ticker := time.NewTicker(purgeAccountLockHeartbeat)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-l.stop:
			return
		case <-ticker.C:
			if err := l.renew(ctx); err != nil {
				select {
				case l.err <- err:
				default:
				}
				l.cancel()
				return
			}
		}
	}
}

func (l *purgeAccountLease) renew(ctx context.Context) error {
	now := time.Now().UTC()
	renewCtx, cancel := context.WithTimeout(ctx, purgeAccountLockHeartbeat)
	defer cancel()
	_, err := l.client.UpdateItem(renewCtx, &dynamodb.UpdateItemInput{
		TableName: awssdk.String(purgeAccountLockTable),
		Key: map[string]dynamodbtypes.AttributeValue{
			"account_id": &dynamodbtypes.AttributeValueMemberS{Value: l.accountID},
		},
		ConditionExpression: awssdk.String("#token = :token"),
		UpdateExpression:    awssdk.String("SET #expires = :expires"),
		ExpressionAttributeNames: map[string]string{
			"#expires": "expires_at",
			"#token":   "lease_token",
		},
		ExpressionAttributeValues: map[string]dynamodbtypes.AttributeValue{
			":expires": &dynamodbtypes.AttributeValueMemberN{Value: fmt.Sprint(now.Add(purgeAccountLockLease).Unix())},
			":token":   &dynamodbtypes.AttributeValueMemberS{Value: l.token},
		},
	})
	if err != nil {
		return fmt.Errorf("renew cleanup account lease: %w", err)
	}
	return nil
}

func (l *purgeAccountLease) close() error {
	close(l.stop)
	<-l.done
	l.cancel()

	releaseCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	_, releaseErr := l.client.DeleteItem(releaseCtx, &dynamodb.DeleteItemInput{
		TableName: awssdk.String(purgeAccountLockTable),
		Key: map[string]dynamodbtypes.AttributeValue{
			"account_id": &dynamodbtypes.AttributeValueMemberS{Value: l.accountID},
		},
		ConditionExpression: awssdk.String("#token = :token"),
		ExpressionAttributeNames: map[string]string{
			"#token": "lease_token",
		},
		ExpressionAttributeValues: map[string]dynamodbtypes.AttributeValue{
			":token": &dynamodbtypes.AttributeValueMemberS{Value: l.token},
		},
	})
	if releaseErr != nil {
		releaseErr = fmt.Errorf("release cleanup account lease: %w", releaseErr)
	}
	var heartbeatErr error
	select {
	case heartbeatErr = <-l.err:
	default:
	}
	return errors.Join(heartbeatErr, releaseErr)
}

func purgeAccountLockConditionalFailure(err error) bool {
	var conditional *dynamodbtypes.ConditionalCheckFailedException
	return errors.As(err, &conditional)
}

func newPurgeAccountLeaseToken() (string, error) {
	var value [16]byte
	if _, err := rand.Read(value[:]); err != nil {
		return "", fmt.Errorf("generate cleanup account lease token: %w", err)
	}
	return hex.EncodeToString(value[:]), nil
}

func purgeAccountLeaseHolder() string {
	runID := strings.TrimSpace(os.Getenv("GITHUB_RUN_ID"))
	if runID != "" {
		return fmt.Sprintf("github:%s:%s", runID, strings.TrimSpace(os.Getenv("GITHUB_RUN_ATTEMPT")))
	}
	host, _ := os.Hostname()
	return fmt.Sprintf("%s:%d", host, os.Getpid())
}