package ddbcacheworker

import (
	"context"
	"dataservices/commons/contextvar"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/commons/utils"
	"dataservices/ddb"
	service "dataservices/internal/service/common"
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"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"
	"go.opentelemetry.io/otel"
	"go.uber.org/zap"
)

type DynamodbStreamsSubscriber interface {
	Subscribe(ctx context.Context, eventChan chan Event)
}

func NewDynamodbStreamsSubscriber(tableName string, updateInterval time.Duration) DynamodbStreamsSubscriber {
	dynamoClient := ddb.Client()
	streamsClient := ddb.StreamClient()
	return &dynamodbStreamsSubscriber{
		dynamoClient:         dynamoClient,
		streamsClient:        streamsClient,
		tableName:            tableName,
		shardProcessingLimit: 20,
		updateInterval:       updateInterval,
	}
}

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

func (dss *dynamodbStreamsSubscriber) Subscribe(ctx context.Context, eventChan chan Event) {
	allShards := make(map[string]struct{})
	shardsChan := make(chan *dynamodbstreams.GetShardIteratorInput, dss.shardProcessingLimit)
	lock := sync.Mutex{}
	ctx = contextvar.WithLoggerField(ctx, zap.String("table", dss.tableName))
	utils.GOWithName("ddbstream_shard_updater:"+dss.tableName, func() {
		tick := time.NewTicker(dss.updateInterval)
		var errorTimes int
		for {
			select {
			case <-ctx.Done():
				logger.LogWithError(ctx, ctx.Err(), "ddbstream_shard_updater exit with context done")
				return
			case <-tick.C:
				metrics.WorkerHeartbeatGuage.WithLabelValues("ddbstream_shard_updater", dss.tableName).Set(float64(time.Now().Unix()))
				ctx, span := otel.Tracer(otelName).Start(ctx, "ReceiveUpdateChan")
				streamArn, err := dss.getLatestStreamArn(ctx)
				if err != nil {
					logger.Error(ctx, "getLatestStreamArn failed", zap.Error(err))
					eventChan <- Event{Err: err, Table: dss.tableName}
					span.End()
					continue
				}
				shards, err := dss.getShards(ctx, streamArn)
				if err != nil {
					errorTimes++
					if errorTimes%5 == 0 {
						logger.Error(ctx, "getShards failed", zap.String("streamArn", *streamArn), zap.Error(err))
					} else {
						logger.Warn(ctx, "getShards failed", zap.String("streamArn", *streamArn), zap.Error(err))
					}
					eventChan <- Event{Err: err, Table: dss.tableName}
					span.End()
					continue
				}
				errorTimes = 0
				logger.Info(ctx, "getShards success", zap.Int("size", len(shards)), zap.Int32("inProgressShard", dss.inProgressShard))
				leafShards := dss.getLeafShards(shards)
				for _, shard := range leafShards {
					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)
	utils.GOWithName("ddbstream_shard_consumer:"+dss.tableName, func() {
		time.Sleep(time.Second * 10)
		tick := time.NewTicker(dss.updateInterval)
		cacheManager := service.NewCacheManager()
		for {
			select {
			case limit <- struct{}{}:
				var shardInput *dynamodbstreams.GetShardIteratorInput
				for {
					select {
					case shardInput = <-shardsChan:
					case <-tick.C:
						metrics.WorkerHeartbeatGuage.WithLabelValues("ddbstream_shard_consumer", dss.tableName).Set(float64(time.Now().Unix()))
					}
					if shardInput != nil {
						break
					}
				}
				utils.GOWithName("ddbstream_shard_processer:"+dss.tableName+":"+*shardInput.ShardId, func() {
					ctx, span := otel.Tracer(otelName).Start(ctx, "processShard")
					atomic.AddInt32(&dss.inProgressShard, 1)
					defer func() {
						span.End()
						atomic.AddInt32(&dss.inProgressShard, -1)
					}()
					seqNumber, err := cacheManager.GetDDBStreamCheckpoint(ctx, dss.tableName, *shardInput.ShardId)
					if err != nil && err != service.ErrNil {
						logger.Error(ctx, "failed to get ddbstream checkpoint from cache", zap.String("table", dss.tableName), zap.Stringp("shard_id", shardInput.ShardId), zap.Error(err))
					} else if err == nil {
						shardInput.SequenceNumber = &seqNumber
						shardInput.ShardIteratorType = types.ShardIteratorTypeAfterSequenceNumber
						logger.Info(ctx, "ddbstream process record after last sequence number", zap.String("table", dss.tableName), zap.Stringp("shard_id", shardInput.ShardId), zap.String("sequence_number", seqNumber))
					}
					err = dss.processShard(ctx, shardInput, eventChan)
					if err != nil {
						eventChan <- Event{Err: err, Table: dss.tableName, ShardID: *shardInput.ShardId}
					}
					<-limit
				})
			case <-tick.C:
				metrics.WorkerHeartbeatGuage.WithLabelValues("ddbstream_shard_consumer", dss.tableName).Set(float64(time.Now().Unix()))
			}
		}
	})
}

func (dss *dynamodbStreamsSubscriber) processShard(
	ctx context.Context,
	input *dynamodbstreams.GetShardIteratorInput,
	eventChan chan Event,
) (err error) {
	start := time.Now()
	ctx = contextvar.WithLoggerField(ctx, zap.Stringp("shard_id", input.ShardId), zap.Stringp("sequence_number", input.SequenceNumber))
	logger.Info(ctx, "start processing shards")
	defer func() {
		logger.LogWithError(ctx, err, "complete processing shards", zap.Duration("cost", time.Since(start)))
	}()
	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 {
			logger.Info(ctx, "get records success", zap.String("table", dss.tableName), zap.String("shardId", *input.ShardId), zap.Int("size", len(records.Records)))
		}
		for _, record := range records.Records {
			eventChan <- Event{
				Table:   dss.tableName,
				Data:    record,
				ShardID: *input.ShardId,
			}
		}

		nextIterator = records.NextShardIterator
		sleepDuration := time.Second
		if nextIterator == nil {
			// Nil next itarator, shard is closed
			logger.Info(ctx, "shard is closed", zap.String("table", dss.tableName), 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
}

func (dss *dynamodbStreamsSubscriber) getLeafShards(shards []types.Shard) []*types.Shard {
	shardMap := make(map[string]*types.Shard)
	shardTree := make(map[string][]*types.Shard)
	var heads []*types.Shard
	for i := range shards {
		shard := &shards[i]
		shardMap[*shard.ShardId] = shard
		if shard.ParentShardId != nil {
			shardTree[*shard.ParentShardId] = append(shardTree[*shard.ParentShardId], shard)
		} else {
			heads = append(heads, shard)
		}
	}
	for parentId, childShards := range shardTree {
		if _, ok := shardMap[parentId]; !ok {
			heads = append(heads, childShards...)
		}
	}

	var result []*types.Shard
	for len(heads) > 0 {
		shard := heads[0]
		heads = heads[1:]
		childList := shardTree[*shard.ShardId]
		if len(childList) == 0 {
			result = append(result, shard)
			continue
		}
		heads = append(heads, childList...)
	}
	return result
}