package ddbtoredis

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"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"
	streamsattributevalue "github.com/aws/aws-sdk-go-v2/feature/dynamodbstreams/attributevalue"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	"github.com/aws/aws-sdk-go-v2/service/dynamodbstreams"
	"github.com/aws/aws-sdk-go-v2/service/dynamodbstreams/types"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
)

type KeyDynamodb struct {
	PublicKey           string `json:"PublicKey"`
	EncryptedPrivateKey string `json:"EncryptedPrivateKey"`
	RateLimitRpm        int    `json:"RateLimitRpm"`
	ExpiredTs           int    `json:"expired_ts"`
}

// the dynamodb stored structure
type OriginRecord struct {
	AccessKey    string `dynamodbav:"identity"`
	PrivateKey   string `dynamodbav:"private_key"`
	RateLimitRpm int    `dynamodbav:"rate_limit_rpm"`
	ExpiredTs    int    `dynamodbav:"expired_ts"`
}

func (r *OriginRecord) ToRecord() (*KeyDynamodb, error) {
	if !strings.HasPrefix(r.AccessKey, "api_key_") {
		return nil, fmt.Errorf("invalid access key: %s", r.AccessKey)
	}
	publicKey := strings.TrimPrefix(r.AccessKey, "api_key_")
	rateLimitRpm := r.RateLimitRpm
	// if dynamodb has not rate limit, set it to 100 rpm
	if rateLimitRpm == 0 {
		rateLimitRpm = 100
	}
	return &KeyDynamodb{
		PublicKey:           publicKey,
		EncryptedPrivateKey: r.PrivateKey,
		RateLimitRpm:        rateLimitRpm,
		ExpiredTs:           r.ExpiredTs,
	}, nil
}

type DynamodbReader interface {
	ScanRecords(ctx context.Context) ([]KeyDynamodb, error)
}

type dynamodbReader struct {
	client    dynamodb.Client
	tableName string
}

func loadAwsConfig(region, assumeRole string) (*aws.Config, error) {
	ctx := context.TODO()
	conf, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return nil, err
	}
	if assumeRole != "" {
		tilog.Info(ctx, "create assume role config", zap.String("roleArn", assumeRole))
		stsClient := sts.NewFromConfig(conf)
		_, err = stsClient.AssumeRole(ctx, &sts.AssumeRoleInput{
			RoleArn:         &assumeRole,
			RoleSessionName: aws.String("ddbtoredis"),
			DurationSeconds: aws.Int32(15 * 60),
		})
		if err != nil {
			return nil, err
		}
		newConf := conf.Copy()
		newConf.Credentials = aws.NewCredentialsCache(
			stscreds.NewAssumeRoleProvider(stsClient, assumeRole),
		)
		return &newConf, nil
	}
	return &conf, nil
}

func NewDynamodbReader(region, tableName, assumeRole string) (*dynamodbReader, error) {
	conf, err := loadAwsConfig(region, assumeRole)
	if err != nil {
		return nil, err
	}
	client := dynamodb.NewFromConfig(*conf)
	return &dynamodbReader{
		client:    *client,
		tableName: tableName,
	}, nil
}

// `ScanRecords` list objects in dynamodb
func (dr *dynamodbReader) ScanRecords(ctx context.Context) ([]KeyDynamodb, error) {
	projExpr := expression.NamesList(
		expression.Name("identity"),
		expression.Name("private_key"),
		expression.Name("rate_limit_rpm"),
		expression.Name("expired_ts"),
	)
	expr, err := expression.NewBuilder().WithProjection(projExpr).Build()
	if err != nil {
		return nil, err
	}

	paginator := dynamodb.NewScanPaginator(&dr.client, &dynamodb.ScanInput{
		TableName:                 aws.String(dr.tableName),
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		ProjectionExpression:      expr.Projection(),
	})
	var records []OriginRecord
	for paginator.HasMorePages() {
		var pageRecords []OriginRecord
		resp, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		err = attributevalue.UnmarshalListOfMaps(resp.Items, &pageRecords)
		if err != nil {
			return nil, err
		}
		records = append(records, pageRecords...)
	}

	result := make([]KeyDynamodb, 0, len(records))
	for _, r := range records {
		record, err := r.ToRecord()
		if err != nil {
			tilog.Error(ctx, "convert record failed", zap.Error(err))
			continue
		}
		result = append(result, *record)
	}
	tilog.Info(ctx, "ScanRecords success", zap.Int("size", len(records)))

	return result, nil
}

type DynamodbStreamsSubscriber interface {
	Subscribe(eventChan chan Event, errorChan chan error)
}

func NewDynamodbStreamsSubscriber(region, tableName, assumeRole string) (DynamodbStreamsSubscriber, error) {
	conf, err := loadAwsConfig(region, assumeRole)
	if err != nil {
		return nil, err
	}
	dynamoClient := dynamodb.NewFromConfig(*conf)
	streamsClient := dynamodbstreams.NewFromConfig(*conf)
	return &dynamodbStreamsSubscriber{
		dynamoClient:         dynamoClient,
		streamsClient:        streamsClient,
		tableName:            tableName,
		shardProcessingLimit: 20,
		updateInterval:       time.Minute,
	}, nil
}

type dynamodbStreamsSubscriber struct {
	dynamoClient         *dynamodb.Client
	streamsClient        *dynamodbstreams.Client
	tableName            string
	shardProcessingLimit int
	updateInterval       time.Duration
	inProgressShard      int32
}

func (dss *dynamodbStreamsSubscriber) Subscribe(eventChan chan Event, errorChan chan error) {
	needUpdateChan := make(chan struct{}, 1)
	needUpdateChan <- struct{}{}

	allShards := make(map[string]struct{})
	shardsChan := make(chan *dynamodbstreams.GetShardIteratorInput, dss.shardProcessingLimit)
	lock := sync.Mutex{}

	go func() {
		tick := time.NewTicker(dss.updateInterval)
		for {
			<-tick.C
			needUpdateChan <- struct{}{}
		}
	}()

	go func() {
		for {
			<-needUpdateChan
			ctx, span := otel.Tracer(otelName).Start(context.TODO(), "ReceiveUpdateChan")
			streamArn, err := dss.getLatestStreamArn(ctx)
			if err != nil {
				tilog.Error(ctx, "getLatestStreamArn failed", zap.Error(err))
				errorChan <- err
				return
			}
			shards, err := dss.getShards(ctx, streamArn)
			if err != nil {
				tilog.Error(ctx, "getShards failed", zap.String("streamArn", *streamArn), zap.Error(err))
				errorChan <- err
				return
			}
			tilog.Info(ctx, "getShards success", zap.Int("size", len(shards)), zap.Int32("inProgressShard", dss.inProgressShard))
			for _, shard := range shards {
				lock.Lock()
				if _, ok := allShards[*shard.ShardId]; !ok {
					allShards[*shard.ShardId] = struct{}{}
					shardsChan <- &dynamodbstreams.GetShardIteratorInput{
						StreamArn:         streamArn,
						ShardId:           shard.ShardId,
						ShardIteratorType: types.ShardIteratorTypeLatest,
					}
				}
				lock.Unlock()
			}
			span.End()
		}
	}()

	limit := make(chan struct{}, dss.shardProcessingLimit)
	go func() {
		time.Sleep(time.Second * 10)
		for shardInput := range shardsChan {
			limit <- struct{}{}
			go func(shardInput *dynamodbstreams.GetShardIteratorInput) {
				ctx, span := otel.Tracer(otelName).Start(context.TODO(), "processShard")
				atomic.AddInt32(&dss.inProgressShard, 1)
				defer func() {
					span.End()
					atomic.AddInt32(&dss.inProgressShard, -1)
				}()
				tilog.Info(ctx, "start to process shard", zap.String("shardId", *shardInput.ShardId))
				err := dss.processShard(ctx, shardInput, eventChan)
				if err != nil {
					tilog.Error(ctx, "failed to process shard", zap.String("shardId", *shardInput.ShardId), zap.Error(err))
					errorChan <- err
				}
				tilog.Info(ctx, "process shard success", zap.String("shardId", *shardInput.ShardId))
				<-limit
			}(shardInput)
		}
	}()
}

func (dss *dynamodbStreamsSubscriber) processShard(
	ctx context.Context,
	input *dynamodbstreams.GetShardIteratorInput,
	eventChan chan Event,
) error {
	iter, err := dss.streamsClient.GetShardIterator(ctx, input)
	if err != nil {
		return err
	}
	nextIterator := iter.ShardIterator
	for nextIterator != nil {
		records, err := dss.streamsClient.GetRecords(ctx, &dynamodbstreams.GetRecordsInput{
			ShardIterator: nextIterator,
		})
		if err != nil {
			return err
		}
		if len(records.Records) > 0 {
			tilog.Info(ctx, "get records success", zap.String("shardId", *input.ShardId), zap.Int("size", len(records.Records)))
		}
		for _, record := range records.Records {
			var newRecord OriginRecord
			var oldRecord OriginRecord
			if err = streamsattributevalue.UnmarshalMap(record.Dynamodb.NewImage, &newRecord); err != nil {
				return err
			}
			if err = streamsattributevalue.UnmarshalMap(record.Dynamodb.OldImage, &oldRecord); err != nil {
				return err
			}

			var eventType EventType
			var key *KeyDynamodb
			var err error
			if oldRecord.AccessKey == "" {
				eventType = Add
				key, err = newRecord.ToRecord()
				if err != nil {
					tilog.Error(ctx, "convert add record failed", zap.Error(err))
					continue
				}
			} else if newRecord.AccessKey == "" {
				eventType = Delete
				key, err = oldRecord.ToRecord()
				if err != nil {
					tilog.Error(ctx, "convert delete record failed", zap.Error(err))
					continue
				}
			} else {
				eventType = Modify
				key, err = newRecord.ToRecord()
				if err != nil {
					tilog.Error(ctx, "convert modify record failed", zap.Error(err))
					continue
				}
			}
			spanCtx := trace.SpanContextFromContext(ctx)
			eventChan <- Event{
				Type:    eventType,
				Key:     *key,
				TraceID: spanCtx.TraceID(),
			}
		}

		nextIterator = records.NextShardIterator
		sleepDuration := time.Second
		if nextIterator == nil {
			// Nil next itarator, shard is closed
			tilog.Info(ctx, "shard is closed", zap.String("shardId", *input.ShardId))
			sleepDuration = time.Millisecond * 10
		} else if len(records.Records) == 0 {
			// Empty set, but shard is not closed -> sleep a little
			sleepDuration = time.Second * 10
		}
		time.Sleep(sleepDuration)
	}
	return nil
}

func (dss *dynamodbStreamsSubscriber) getLatestStreamArn(ctx context.Context) (*string, error) {
	tableInfo, err := dss.dynamoClient.DescribeTable(
		ctx,
		&dynamodb.DescribeTableInput{TableName: &dss.tableName},
	)
	if err != nil {
		return nil, err
	}
	if tableInfo.Table.LatestStreamArn == nil {
		return nil, errors.New("empty table stream arn")
	}
	return tableInfo.Table.LatestStreamArn, nil
}

func (dss *dynamodbStreamsSubscriber) getShards(ctx context.Context, streamArn *string) ([]types.Shard, error) {
	streamOutput, err := dss.streamsClient.DescribeStream(
		ctx,
		&dynamodbstreams.DescribeStreamInput{StreamArn: streamArn},
	)
	if err != nil {
		return nil, err
	}
	if len(streamOutput.StreamDescription.Shards) == 0 {
		return nil, nil
	}
	return streamOutput.StreamDescription.Shards, nil
}