package scheduler

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	s3Types "github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/config"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/entity"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/metric"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/pkgv2/logid"

	"go.uber.org/zap"
)

const (
	auditLogScheduleConcurrency = 5
	checkPointDelay             = 5 * time.Minute
)

type MergeContext struct {
	startObjectKey         string
	endObjectKey           string
	endObjectKeyLastModify time.Time
	totalFileSize          int64
	objectCount            int

	currentTime time.Time
	checkpoint  time.Time
	taskNumber  int
}

func (m *MergeContext) reset(obj s3Types.Object) {
	m.startObjectKey = *obj.Key
	m.endObjectKey = *obj.Key
	m.endObjectKeyLastModify = *obj.LastModified
	m.totalFileSize = *obj.Size
	m.objectCount = 1
}

func (m *MergeContext) update(obj s3Types.Object) {
	m.totalFileSize += *obj.Size
	m.objectCount++
	m.endObjectKey = *obj.Key
	m.endObjectKeyLastModify = *obj.LastModified
}

// Merge check and generate audit log tasks
func (s *Scheduler) Merge() {
	ctx := logid.WithLogid(context.Background(), logid.GenLogid())
	logger := log.WithContextL(ctx)

	logger.Info("start to run merge task")

	schedules, err := s.auditLogScheduleRepo.List(ctx, &entity.ListAuditLogScheduleOptions{
		State: entity.AuditLogScheduleStateEnabled,
	})
	if err != nil {
		logger.Error("failed to list audit log schedules", zap.Error(err))
		metric.AuditLogScheduleMergeFailedGauge.WithLabelValues(metric.MergeGlobal, "list audit log schedules error").Inc()
		return
	}

	var wg sync.WaitGroup
	jobs := make(chan *entity.AuditLogSchedule, len(schedules))

	for i := 0; i < auditLogScheduleConcurrency; i++ {
		wg.Add(1)
		go s.consumer(ctx, jobs, &wg)
	}
	// producer
	for _, schedule := range schedules {
		jobs <- schedule
	}
	close(jobs)

	wg.Wait()

	logger.Info("run merge task finish")
}

func (s *Scheduler) consumer(ctx context.Context, jobs <-chan *entity.AuditLogSchedule, wg *sync.WaitGroup) {
	defer wg.Done()
	logger := log.WithContextL(ctx)
	for schedule := range jobs {
		if err := s.mergeOneCluster(ctx, schedule); err != nil {
			logger.Error("error occurs when merge one cluster", zap.String("cluster_id", schedule.ClusterId),
				zap.String("start_after", schedule.S3StartAfter), zap.Error(err))
			metric.AuditLogScheduleMergeFailedGauge.WithLabelValues(schedule.ClusterId, err.Error()).Inc()
		}
	}
}

func (s *Scheduler) mergeOneCluster(ctx context.Context, schedule *entity.AuditLogSchedule) (err error) {
	logger := log.WithContextL(ctx).With(zap.String("cluster_id", schedule.ClusterId))
	now := time.Now()
	// Assume that all the audit logs before the checkpoint have been generated and uploaded to S3.
	checkpoint := now.Add(-checkPointDelay)
	mergeCtx := &MergeContext{
		checkpoint:  checkpoint,
		currentTime: now,
		taskNumber:  0,
	}
	logger.Info("start to merge one cluster", zap.Time("checkpoint", checkpoint))

	defer func() {
		if err == nil {
			err = s.EnsureScheduleState(ctx, checkpoint, schedule)
		}
	}()

	exceedCheckpoint := false
	input := s.prepareS3ListInput(schedule)
	for {
		result, err := s.o11yS3Client.ListObjectsV2(ctx, input)
		if err != nil {
			logger.Error("failed to list objects", zap.Error(err))
			return err
		}
		for _, obj := range result.Contents {
			if obj.LastModified.After(checkpoint) {
				exceedCheckpoint = true
				break
			}
			// skip the directory
			if strings.HasSuffix(*obj.Key, "/") {
				continue
			}
			// init the merge context
			if mergeCtx.objectCount == 0 {
				mergeCtx.reset(obj)
				continue
			}
			// merge the previous files when the current size exceeds the limit
			if schedule.ExceedMaxFileSize(mergeCtx.totalFileSize + *obj.Size) {
				logger.Info("total file size exceeds the limit, generate task")
				err := s.GenerateTask(ctx, schedule, mergeCtx)
				if err != nil {
					logger.Error("failed to generate task", zap.Error(err))
					return err
				}
				mergeCtx.reset(obj)
			} else {
				mergeCtx.update(obj)
			}
		}
		// quit when exceed the checkpoint or there is no more object
		if exceedCheckpoint || result.NextContinuationToken == nil || *result.NextContinuationToken == "" {
			break
		}
		input.ContinuationToken = result.NextContinuationToken
	}

	// handle the rest file
	if err := s.handleRestFiles(ctx, schedule, mergeCtx); err != nil {
		logger.Error("failed to handle rest files", zap.Error(err))
		return err
	}
	logger.Info("finish to merge one cluster", zap.Int("generate_task_number", mergeCtx.taskNumber))
	return nil
}

func (s *Scheduler) prepareS3ListInput(schedule *entity.AuditLogSchedule) *s3.ListObjectsV2Input {
	input := &s3.ListObjectsV2Input{
		Bucket: &s.config.O11yBucketName,
		Prefix: aws.String(config.GetO11yAuditLogPrefix(
			s.config.Env,
			schedule.OrgId,
			schedule.ProjectId,
			schedule.ClusterId,
			s.config.O11yBucketName,
			s.config.CloudProviderType)),
	}
	// s3 start after is empty when it is the first time to merge
	if schedule.S3StartAfter != "" {
		input.StartAfter = aws.String(schedule.S3StartAfter)
	}
	return input
}

func (s *Scheduler) handleRestFiles(
	ctx context.Context,
	schedule *entity.AuditLogSchedule,
	mergeCtx *MergeContext) error {
	logger := log.WithContextL(ctx).With(zap.String("cluster_id", schedule.ClusterId))
	if mergeCtx.objectCount == 0 {
		return nil
	}
	// it will occur when the last single file is larger than the max file size
	if schedule.ExceedMaxFileSize(mergeCtx.totalFileSize) {
		logger.Info("total file size exceeds the limit, generate task")
		return s.GenerateTask(ctx, schedule, mergeCtx)
	}
	if !mergeCtx.currentTime.Before(schedule.RotationTime) {
		logger.Info("rotation time meets, generate task")
		return s.GenerateTask(ctx, schedule, mergeCtx)
	}
	if schedule.DeleteTime != nil && mergeCtx.checkpoint.After(*schedule.DeleteTime) {
		logger.Info("all files before delete time has been handled, generate task")
		return s.GenerateTask(ctx, schedule, mergeCtx)
	}
	if schedule.DisableTime != nil && mergeCtx.checkpoint.After(*schedule.DisableTime) {
		logger.Info("all files before disable time has been handled, generate task")
		return s.GenerateTask(ctx, schedule, mergeCtx)
	}
	return nil
}

// GenerateTask generate audit log task and update schedule in a transaction
func (s *Scheduler) GenerateTask(ctx context.Context, schedule *entity.AuditLogSchedule, mergeCtx *MergeContext) error {
	fileName := generateFileName(schedule.LastFileName, mergeCtx.endObjectKeyLastModify)
	schedule.S3StartAfter = mergeCtx.endObjectKey
	schedule.ResetRotationTime()
	schedule.LastFileName = fileName

	task := &entity.AuditLogTask{
		TaskId:         GenerateIDWithPrefix("alt-"),
		OrgId:          schedule.OrgId,
		ProjectId:      schedule.ProjectId,
		ClusterId:      schedule.ClusterId,
		State:          entity.AuditLogTaskStatePending,
		FileName:       fileName,
		StartObjectKey: mergeCtx.startObjectKey,
		EndObjectKey:   mergeCtx.endObjectKey,
		ObjectCount:    mergeCtx.objectCount,
	}

	return s.txProvider.Transaction(ctx, func(ctx context.Context) error {
		if err := s.auditLogTaskRepo.Create(ctx, task); err != nil {
			return err
		}
		if err := s.auditLogScheduleRepo.Save(ctx, schedule); err != nil {
			return err
		}
		mergeCtx.taskNumber++
		return nil
	})
}

// GenerateFileName use end-object's modify time as the file name
// It will add index when conflict with the last file name, so the file name format is
// 2025-10-01-12-00.log or 2025-10-01-12-00.1.log
func generateFileName(lastFileName string, endObjectKeyLastModify time.Time) string {
	format := "2006-01-02-15-04"
	fileNameTimePart := endObjectKeyLastModify.Format(format)
	// return directly when file name is not conflict with the last file name
	if lastFileName == "" || !strings.HasPrefix(lastFileName, fileNameTimePart) {
		return fmt.Sprintf("%s.log", fileNameTimePart)
	}

	index := 1
	parts := strings.Split(lastFileName, ".")
	if len(parts) == 3 { //nolint:mnd // no need
		if parsedIndex, err := strconv.Atoi(parts[1]); err == nil {
			index = parsedIndex + 1
		}
	}
	return fmt.Sprintf("%s.%d.log", fileNameTimePart, index)
}

func (s *Scheduler) EnsureScheduleState(ctx context.Context, checkpoint time.Time,
	schedule *entity.AuditLogSchedule) error {
	logger := log.WithContextL(ctx).With(zap.String("cluster_id", schedule.ClusterId))
	// judge delete time first
	if schedule.DeleteTime != nil && checkpoint.After(*schedule.DeleteTime) {
		if err := s.auditLogScheduleRepo.SoftDelete(ctx, schedule); err != nil {
			logger.Error("failed to soft delete schedule, will retry in the next loop", zap.Error(err))
			return err
		}
		logger.Info("delete schedule success")
	}
	if schedule.DisableTime != nil && checkpoint.After(*schedule.DisableTime) {
		if err := s.auditLogScheduleRepo.Disable(ctx, schedule); err != nil {
			logger.Error("failed to disable schedule, will retry in the next loop", zap.Error(err))
			return err
		}
		logger.Info("disable schedule success")
	}
	return nil
}