package ddbcacheworker

import (
	"context"
	"dataservices/commons/contextvar"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/commons/utils"
	service "dataservices/internal/service/common"
	"dataservices/worker"
	"fmt"
	"math/big"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/dynamodbstreams/types"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
)

type EventType string

const (
	Add    EventType = "ADD"
	Modify EventType = "MODIFY"
	Delete EventType = "DELETE"
)

const otelName = "syncer"

type Event struct {
	Type    EventType
	Table   string
	ShardID string
	Data    types.Record
	Err     error
}

type Table func() string
type Handler func(ctx context.Context, record Event, buffer service.CacheManager) error

type registerItem struct {
	table   Table
	handler Handler
}

var registerBox []registerItem

func RegisterHandler(table Table, handler Handler) {
	registerBox = append(registerBox, registerItem{
		table:   table,
		handler: handler,
	})
}

type Syncer struct {
	ctx                    context.Context
	cacheManager           service.CacheManager
	ddbStreamsSubscribers  map[string]DynamodbStreamsSubscriber
	eventChanMap           map[string]chan Event
	ddbTableHandlerMap     map[string]Handler
	workerNumPreTable      int
	heartbeatInterval      time.Duration
	checkpointChan         chan Event
	checkpointSaveInterval time.Duration
}

func NewSyncer(ctx context.Context) worker.Worker {
	syncer := &Syncer{
		ctx:                    ctx,
		cacheManager:           service.NewCacheManager(),
		ddbStreamsSubscribers:  make(map[string]DynamodbStreamsSubscriber),
		eventChanMap:           make(map[string]chan Event),
		ddbTableHandlerMap:     make(map[string]Handler),
		workerNumPreTable:      1,
		heartbeatInterval:      time.Second * 30,
		checkpointChan:         make(chan Event, 1000),
		checkpointSaveInterval: time.Minute,
	}
	for _, registerItem := range registerBox {
		table := registerItem.table()
		syncer.ddbTableHandlerMap[table] = registerItem.handler
		syncer.eventChanMap[table] = make(chan Event, 100)
		syncer.ddbStreamsSubscribers[table] = NewDynamodbStreamsSubscriber(table, time.Minute)
	}
	return syncer
}

func (s *Syncer) Name() string {
	return "ddbcache_worker"
}

func (s *Syncer) IsEnable() bool {
	return true
}

func (s *Syncer) RunningMsInterval() time.Duration {
	return time.Minute
}

func (s *Syncer) Handle(tracer trace.Tracer) {
	ctx, span := tracer.Start(s.ctx, s.Name())
	defer span.End()

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	ctx = contextvar.WithLoggerField(ctx, zap.String("handler_id", utils.RandStr(32)))

	// Preempting locks to prevent multi process parallelism
	cacheKey := "ddbcache_worker"
	if err := s.cacheManager.Lock(ctx, cacheKey, time.Minute); err != nil {
		if err == service.ErrLockOccupied {
			logger.Warn(ctx, "there are already ddbcache worker executing tasks, will try again latter", zap.Error(err))
		} else {
			logger.Error(ctx, "ddbcache worker lock failed, will try again latter", zap.Error(err))
		}
		return
	}
	defer s.cacheManager.Unlock(ctx, cacheKey)
	logger.Info(ctx, "ddbcache worker successfully locked and started executing tasks")

	// Asynchronous refresh of lock expiration time. If an exception occurs,
	// simply exit and wait for the worker to be rescheduled for execution
	utils.GOWithName("ddbcache_worker_lock_refresher", func() {
		tick := time.NewTicker(time.Second * 10)
		for {
			select {
			case <-tick.C:
				if err := s.cacheManager.RefreshLock(ctx, cacheKey, time.Minute); err != nil {
					if err == service.ErrLockExpired {
						logger.Error(ctx, "ddbcache worker lock has expired, preparing to exit and try again latter", zap.Error(err))
					} else {
						logger.Error(ctx, "ddbcache worker lock refresh failed, preparing to exit and try again latter", zap.Error(err))
					}
					cancel()
					return
				}
			case <-ctx.Done():
				logger.Info(ctx, "context done, dddb cache worker lock refresher exists", zap.Error(ctx.Err()))
				return
			}
		}
	})

	for table, subscriber := range s.ddbStreamsSubscribers {
		subscriber.Subscribe(ctx, s.eventChanMap[table])
	}

	s.runCheckpointStorage()

	var wg sync.WaitGroup
	for table := range s.eventChanMap {
		eventChan := s.eventChanMap[table]
		eventHandler := s.ddbTableHandlerMap[table]
		for i := 0; i < s.workerNumPreTable; i++ {
			wg.Add(1)
			seq := i
			workerName := fmt.Sprintf("%s:%s:%d", s.Name(), table, seq)
			utils.GOWithName(workerName, func() {
				defer wg.Done()
				tick := time.NewTicker(s.heartbeatInterval)
				ctx := contextvar.WithLoggerField(ctx, zap.String("worker_name", workerName))
				for {
					select {
					case event := <-eventChan:
						ctx, span := otel.Tracer(otelName).Start(ctx, "ReceiveTableEvent")
						var seqNumber *string
						if event.Data.Dynamodb != nil {
							seqNumber = event.Data.Dynamodb.SequenceNumber
						}
						ctx = contextvar.WithLoggerField(ctx, zap.String("table", event.Table),
							zap.String("shard_id", event.ShardID), zap.Stringp("sequence_number", seqNumber),
							zap.Stringp("event_id", event.Data.EventID), zap.String("event_name", string(event.Data.EventName)))
						logger.LogWithError(ctx, event.Err, "ddbcache_worker got new event")
						if event.Err != nil {
							span.End()
							continue
						}
						err := eventHandler(ctx, event, s.cacheManager)
						logger.LogWithError(ctx, err, "ddbcache_worker complete handle table event")
						s.checkpointChan <- event
						span.End()
					case <-ctx.Done():
						err := ctx.Err()
						logger.Info(ctx, "context done, ddbcache_worker exists", zap.Error(err))
						return
					case <-tick.C:
						metrics.WorkerHeartbeatGuage.WithLabelValues("ddbcache_worker", workerName).Set(float64(time.Now().Unix()))
					}
				}
			})
		}
	}
	wg.Wait()
}

type Checkpoint struct {
	Table              string
	ShardID            string
	LastStoreTime      time.Time
	LastActiveTime     time.Time
	LastSequenceNumber *big.Int
}

func (s *Syncer) runCheckpointStorage() {
	var checkpointMap sync.Map
	workerName := fmt.Sprintf("%s:checkpointstorage", s.Name())
	utils.GOWithName(workerName, func() {
		tick := time.NewTicker(s.checkpointSaveInterval)
		ctx := contextvar.WithLoggerField(s.ctx, zap.String("worker_name", workerName))
		for {
			select {
			case event := <-s.checkpointChan:
				if event.Data.Dynamodb == nil || event.Data.Dynamodb.SequenceNumber == nil {
					continue
				}
				seqNumber := new(big.Int)
				seqNumber.SetString(*event.Data.Dynamodb.SequenceNumber, 10)
				key := fmt.Sprintf("%s_%s", event.Table, event.ShardID)
				data, ok := checkpointMap.Load(key)
				var checkpoint *Checkpoint
				if ok {
					checkpoint = data.(*Checkpoint)
					if checkpoint.LastSequenceNumber.Cmp(seqNumber) == -1 {
						checkpoint.LastSequenceNumber = seqNumber
					}
				} else {
					checkpoint = new(Checkpoint)
					checkpoint.Table = event.Table
					checkpoint.ShardID = event.ShardID
					checkpoint.LastSequenceNumber = seqNumber
					checkpointMap.Store(key, checkpoint)
				}
				now := time.Now()
				checkpoint.LastActiveTime = now
				if now.Sub(checkpoint.LastStoreTime) > time.Minute {
					if err := s.cacheManager.SetDDBStreamCheckpoint(ctx, event.Table, event.ShardID, checkpoint.LastSequenceNumber.String()); err != nil {
						logger.Error(ctx, "failed to set ddbstream checkpoint", zap.String("table", event.Table), zap.String("shard_id", event.ShardID), zap.String("sequence_number", checkpoint.LastSequenceNumber.String()))
					} else {
						checkpoint.LastStoreTime = now
					}
				}
			case <-ctx.Done():
				err := ctx.Err()
				logger.Info(ctx, "context done, checkpoint storage exists", zap.Error(err))
				return
			case <-tick.C:
				metrics.WorkerHeartbeatGuage.WithLabelValues("ddbcache_worker", workerName).Set(float64(time.Now().Unix()))
				checkpointMap.Range(func(key, value any) bool {
					checkpoint := value.(*Checkpoint)
					table, shardID, seqNumber := checkpoint.Table, checkpoint.ShardID, checkpoint.LastSequenceNumber
					if time.Since(checkpoint.LastActiveTime) > time.Minute*20 {
						// checkpoint has expired, delete it!
						checkpointMap.Delete(key)
					} else if time.Since(checkpoint.LastStoreTime) > time.Minute &&
						checkpoint.LastActiveTime.After(checkpoint.LastStoreTime) {
						// checkpoint is saved more than 1 minute from the last time
						// and there are updates during this period, the checkpoint will be automatically saved
						if err := s.cacheManager.SetDDBStreamCheckpoint(ctx, table, shardID, seqNumber.String()); err != nil {
							logger.Error(ctx, "failed to set ddbstream checkpoint", zap.String("table", table), zap.String("shard_id", shardID), zap.String("sequence_number", seqNumber.String()))
						} else {
							checkpoint.LastStoreTime = time.Now()
						}
					}
					return true
				})
			}
		}
	})
}