package worker

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"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/repo"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/util"
	"github.com/tidbcloud/aws-shared-provider/pkg/storage"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/pkgv2/import/cloud_storage"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/pkgv2/logid"
	"go.uber.org/zap"
)

const (
	taskConcurrency       = 3
	taskMismatchThreshold = 10
	downloadConcurrency   = 10

	downloadRetryTimes    = 3
	downloadRetryInterval = 1 * time.Second
)

type Worker struct {
	config               config.AuditLogWorkerConfig
	s3Client             *s3.Client
	o11yS3Client         *s3.Client
	auditLogTaskRepo     repo.AuditLogTaskRepo
	auditLogScheduleRepo repo.AuditLogScheduleRepo
}

func NewWorker(config config.AuditLogWorkerConfig) (*Worker, error) {
	if err := config.Check(); err != nil {
		return nil, err
	}
	s3Client, err := util.GetS3Client(config.CloudProviderType, config.Region)
	if err != nil {
		return nil, fmt.Errorf("get s3 client err: %w", err)
	}
	o11yS3Clients, err := util.GetO11yS3Client(config.CloudProviderType, config.Region)
	if err != nil {
		return nil, fmt.Errorf("get o11y s3 clients err: %w", err)
	}

	rdsSecretRegion := config.Region
	if config.RdsSecretRegion != "" {
		rdsSecretRegion = config.RdsSecretRegion
	}
	rdsConfig, err := util.GetRdsConfig(config.CloudProviderType, config.RdsSecretArn, rdsSecretRegion)
	if err != nil {
		return nil, fmt.Errorf("get rds config err: %w", err)
	}
	db, err := storage.NewGormDB(rdsConfig)
	if err != nil {
		return nil, fmt.Errorf("new gorm db err: %w", err)
	}
	dbProvider, _ := storage.NewDBProvier(db)

	return &Worker{
		s3Client:             s3Client,
		o11yS3Client:         o11yS3Clients,
		config:               config,
		auditLogTaskRepo:     repo.NewAuditLogTaskRepo(dbProvider),
		auditLogScheduleRepo: repo.NewAuditLogScheduleRepo(dbProvider),
	}, nil
}

func (w *Worker) Run() error {
	ctx := logid.WithLogid(context.Background(), logid.GenLogid())
	logger := log.WithContextL(ctx)

	logger.Info("start to run audit log worker")

	tasks, err := w.auditLogTaskRepo.List(context.Background(), &entity.ListAuditLogTaskOptions{
		TaskIds: w.config.TaskIds,
	})
	if err != nil {
		logger.Error("failed to list audit log tasks", zap.Error(err))
		return err
	}

	var wg sync.WaitGroup
	var hasErr bool
	tasksCh := make(chan *entity.AuditLogTask, len(tasks))
	for range taskConcurrency {
		wg.Add(1)
		go w.consumer(ctx, tasksCh, &wg, &hasErr)
	}
	for _, task := range tasks {
		tasksCh <- task
	}
	close(tasksCh)
	wg.Wait()

	if hasErr {
		return errors.New("one or more tasks failed")
	}

	logger.Info("run audit log worker success")
	return nil
}

func (w *Worker) consumer(ctx context.Context, tasksCh <-chan *entity.AuditLogTask, wg *sync.WaitGroup, hasErr *bool) {
	defer wg.Done()
	logger := log.WithContextL(ctx)
	for task := range tasksCh {
		// error occurs when task state is not updated
		// will not return error when task state is updated to failed or corrupted
		if err := w.executeOneTask(ctx, task); err != nil {
			logger.Error("failed to execute one task", zap.String("task_id", task.TaskId),
				zap.String("cluster_id", task.ClusterId), zap.Error(err))
			*hasErr = true
		}
	}
}

func (w *Worker) executeOneTask(ctx context.Context, task *entity.AuditLogTask) (err error) {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))

	var schedule *entity.AuditLogSchedule

	defer func() {
		task.ExecuteCount++
		task.State = entity.AuditLogTaskStateSucceeded
		if err != nil {
			task.State = entity.AuditLogTaskStateFailed
			task.ErrorMsg = err.Error()
			if strings.HasPrefix(err.Error(), "corrupted:") {
				task.State = entity.AuditLogTaskStateCorrupted
			}
		}
		if schedule != nil {
			task.CloudStorageType = schedule.CloudStorageType
			task.CloudStorage = schedule.CloudStorage
		}
		err = w.auditLogTaskRepo.UpdateRunningState(ctx, task)
	}()

	// get schedule
	schedule, err = w.auditLogScheduleRepo.GetByClusterId(ctx, task.ClusterId)
	if err != nil {
		logger.Error("failed to get audit log schedule", zap.Error(err))
		return err
	}

	// check task
	if task.State == entity.AuditLogTaskStatePending || task.State == entity.AuditLogTaskStateSucceeded ||
		task.State == entity.AuditLogTaskStateCorrupted {
		logger.Warn("skip the task", zap.String("state", string(task.State)))
		return nil
	}

	// list files
	logger.Info("start to list objects")
	objects, err := w.ListObjects(ctx, task)
	if err != nil {
		return err
	}

	// merge files
	logger.Info("start to merge files", zap.Int("object_count", len(objects)))
	file, err := os.Create(task.FileName)
	if err != nil {
		logger.Error("unable to create file", zap.Error(err))
		return err
	}
	defer file.Close() //nolint:errcheck // no need check error
	if err := w.mergeFiles(ctx, objects, task, file); err != nil {
		return err
	}
	// the pointer of file is at the end of file, need to seek to the start of file
	if _, err := file.Seek(0, io.SeekStart); err != nil {
		logger.Error("failed to seek file", zap.Error(err))
		return err
	}

	// upload file
	logger.Info("start to upload file")
	switch schedule.CloudStorageType {
	case entity.AuditLogCloudStorageTypeS3:
		return w.uploadFileToS3(ctx, task, schedule, file)
	case entity.AuditLogCloudStorageTypeOSS:
		return w.uploadFileToOSS(ctx, task, schedule, file)
	case entity.AuditLogCloudStorageTypeAzureBlob:
		return w.uploadFileToAzureBlob(ctx, task, schedule, file)
	case entity.AuditLogCloudStorageTypeGCS:
		return w.uploadFileToGCS(ctx, task, schedule, file)
	case "", entity.AuditLogCloudStorageTypeTiDBCloud:
		return w.uploadFilesToTiDBCloud(ctx, task, file)
	default:
		return fmt.Errorf("unsupported cloud storage type: %s", schedule.CloudStorageType)
	}
}

func (w *Worker) uploadFilesToTiDBCloud(ctx context.Context, task *entity.AuditLogTask, file *os.File) error {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))
	key := fmt.Sprintf("%s/%s/%s/%s", task.OrgId, task.ProjectId, task.ClusterId, task.FileName)
	_, err := w.s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: &w.config.BucketName,
		Key:    &key,
		Body:   file,
	})
	if err != nil {
		logger.Error("failed to upload file", zap.Error(err))
		return err
	}
	return nil
}

func (w *Worker) uploadFileToS3(ctx context.Context, task *entity.AuditLogTask, schedule *entity.AuditLogSchedule, file *os.File) error {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))

	// build s3 cloud storage
	cloudStorageS3 := schedule.CloudStorage.S3
	options := make([]cloud_storage.S3Option, 0)
	if cloudStorageS3.RoleArn != "" && cloudStorageS3.ExternalId != "" {
		options = append(options, cloud_storage.WithAssumeRole(cloudStorageS3.RoleArn, cloudStorageS3.ExternalId))
	} else if cloudStorageS3.AccessKeyId != "" && cloudStorageS3.SecretAccessKey != "" {
		options = append(options, cloud_storage.WithAccessKey(cloudStorageS3.AccessKeyId, cloudStorageS3.SecretAccessKey))
	} else {
		return errors.New("invalid S3Config: must provide either (accessKeyID & secretAccessKey) or (roleArn & externalId)")
	}
	if region := cloudStorageS3.Region; region != "" {
		options = append(options, cloud_storage.WithS3Region(region))
	}
	cs, bizErr := cloud_storage.NewS3(ctx, cloudStorageS3.Uri, options...)
	if bizErr != nil {
		logger.Error("failed to create S3 client", zap.Error(bizErr))
		return fmt.Errorf("failed to create S3 client")
	}
	csS3, ok := cs.(*cloud_storage.S3)
	if !ok || csS3 == nil {
		logger.Error("failed to transfer to S3 client")
		return errors.New("failed to transfer to S3 client")
	}

	// upload file
	key := task.FileName
	if csS3.Prefix != nil && *csS3.Prefix != "" {
		if !strings.HasSuffix(*csS3.Prefix, "/") {
			key = fmt.Sprintf("%s/%s", *csS3.Prefix, task.FileName)
		} else {
			key = fmt.Sprintf("%s%s", *csS3.Prefix, task.FileName)
		}
	}
	_, err := csS3.Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(csS3.Bucket),
		Key:    &key,
		Body:   file,
	})
	if err != nil {
		logger.Error("failed to upload file", zap.Error(err))
		return err
	}
	return nil
}

func (w *Worker) uploadFileToOSS(ctx context.Context, task *entity.AuditLogTask, schedule *entity.AuditLogSchedule, file *os.File) error {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))

	// build oss cloud storage
	cloudStorageOSS := schedule.CloudStorage.OSS
	options := make([]cloud_storage.OSSOption, 0)
	if cloudStorageOSS.RoleArn != "" && cloudStorageOSS.ExternalId != "" {
		options = append(options, cloud_storage.WithOSSAssumeRole(cloudStorageOSS.RoleArn, cloudStorageOSS.ExternalId))
	} else if cloudStorageOSS.AccessKeyId != "" && cloudStorageOSS.AccessKeySecret != "" {
		options = append(options, cloud_storage.WithOSSAccessKey(cloudStorageOSS.AccessKeyId, cloudStorageOSS.AccessKeySecret))
	} else {
		logger.Error("invalid OSSConfig: must provide either (accessKeyID & secretAccessKey) or (roleArn & externalId)")
		return errors.New("invalid OSSConfig: must provide either (accessKeyID & secretAccessKey) or (roleArn & externalId)")
	}

	region := cloudStorageOSS.Region
	if region == "" {
		bucketRegion, bizErr := cloud_storage.GetOssBucketRegion(ctx, cloudStorageOSS.Uri, options...)
		if bizErr != nil {
			logger.Error("failed to get OSS bucket region", zap.Error(bizErr))
			return errors.New("failed to get OSS bucket region")
		}
		region = bucketRegion
	}
	options = append(options, cloud_storage.WithOSSRegion(region))
	if w.config.CloudProviderType == types.ProviderTypeAliCloudFreeTier && w.config.Region == region {
		logger.Info("use internal endpoint for OSS", zap.String("region", region))
		options = append(options, cloud_storage.UseInternalEndpoint())
	}
	cs, bizErr := cloud_storage.NewOSS(ctx, cloudStorageOSS.Uri, options...)
	if bizErr != nil {
		logger.Error("failed to create S3 client", zap.Error(bizErr))
		return fmt.Errorf("failed to create S3 client")
	}
	csS3, ok := cs.(*cloud_storage.S3)
	if !ok || csS3 == nil {
		logger.Error("failed to create OSS client")
		return errors.New("failed to create OSS client")
	}

	// upload using s3 sdk
	key := task.FileName
	if csS3.Prefix != nil && *csS3.Prefix != "" {
		if !strings.HasSuffix(*csS3.Prefix, "/") {
			key = fmt.Sprintf("%s/%s", *csS3.Prefix, task.FileName)
		} else {
			key = fmt.Sprintf("%s%s", *csS3.Prefix, task.FileName)
		}
	}
	_, err := csS3.Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(csS3.Bucket),
		Key:    &key,
		Body:   file,
	})
	if err != nil {
		logger.Error("failed to upload file", zap.Error(err))
		return err
	}
	return nil
}

func (w *Worker) uploadFileToGCS(ctx context.Context, task *entity.AuditLogTask, schedule *entity.AuditLogSchedule, file *os.File) error {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))

	cloudStorageGCS := schedule.CloudStorage.GCS
	options := make([]cloud_storage.GCSOption, 0)
	if cloudStorageGCS.ServiceAccountKey != "" {
		options = append(options, cloud_storage.WithGCSAccountKey(cloudStorageGCS.ServiceAccountKey))
	} else {
		return errors.New("invalid GCSConfig: must provide serviceAccountKey")
	}
	cs, bizErr := cloud_storage.NewGCS(ctx, cloudStorageGCS.Uri, options...)
	if bizErr != nil {
		logger.Error("failed to create GCS client", zap.Error(bizErr))
		return fmt.Errorf("failed to create GCS client")
	}
	csGCS, ok := cs.(*cloud_storage.GCS)
	if !ok || csGCS == nil {
		logger.Error("failed to transfer to GCS client")
		return errors.New("failed to transfer to GCS client")
	}
	defer csGCS.Client.Close()

	// upload file
	key := task.FileName
	if csGCS.Prefix != nil && *csGCS.Prefix != "" {
		if !strings.HasSuffix(*csGCS.Prefix, "/") {
			key = fmt.Sprintf("%s/%s", *csGCS.Prefix, task.FileName)
		} else {
			key = fmt.Sprintf("%s%s", *csGCS.Prefix, task.FileName)
		}
	}

	bucketWriter := csGCS.Client.Bucket(csGCS.Bucket).Object(key).NewWriter(ctx)
	if _, err := io.Copy(bucketWriter, file); err != nil {
		logger.Error("failed to upload file", zap.Error(err))
		return err
	}
	// must call close the writer to upload the file
	if err := bucketWriter.Close(); err != nil {
		logger.Error("failed to close writer", zap.Error(err))
		return err
	}
	return nil
}

func (w *Worker) uploadFileToAzureBlob(ctx context.Context, task *entity.AuditLogTask, schedule *entity.AuditLogSchedule, file *os.File) error {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))

	cloudStorageAzBlob := schedule.CloudStorage.AzureBlob
	options := make([]cloud_storage.AzBlobOption, 0)
	if cloudStorageAzBlob.SASToken != "" {
		options = append(options, cloud_storage.WithAzBlobSasToken(cloudStorageAzBlob.SASToken))
	} else {
		return errors.New("invalid AzureBlobConfig: must provide sasToken")
	}
	cs, bizErr := cloud_storage.NewAzBlob(ctx, cloudStorageAzBlob.Uri, options...)
	if bizErr != nil {
		logger.Error("failed to create az blob client", zap.Error(bizErr))
		return fmt.Errorf("failed to create az blob client")
	}
	csAzBlob, ok := cs.(*cloud_storage.AzBlob)
	if !ok || csAzBlob == nil {
		logger.Error("failed to transfer to az blob client")
		return errors.New("failed to transfer to az blob client")
	}

	// upload file
	key := task.FileName
	if csAzBlob.Prefix != nil && *csAzBlob.Prefix != "" {
		if !strings.HasSuffix(*csAzBlob.Prefix, "/") {
			key = fmt.Sprintf("%s/%s", *csAzBlob.Prefix, task.FileName)
		} else {
			key = fmt.Sprintf("%s%s", *csAzBlob.Prefix, task.FileName)
		}
	}
	if _, err := csAzBlob.ContainerClient.NewBlockBlobClient(key).UploadFile(ctx, file, nil); err != nil {
		logger.Error("failed to upload file", zap.Error(err))
		return err
	}
	return nil
}

func (w *Worker) ListObjects(ctx context.Context, task *entity.AuditLogTask) ([]string, error) {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))

	// handle the first object
	objects := make([]string, 0, task.ObjectCount)
	objects = append(objects, task.StartObjectKey)
	if task.ObjectCount == 1 {
		return objects, nil
	}
	// handle the rest objects
	mergeCount := 1
	finish := false
	prefix := config.GetO11yAuditLogPrefix(
		w.config.Env,
		task.OrgId,
		task.ProjectId,
		task.ClusterId,
		w.config.O11yBucketName,
		w.config.CloudProviderType)
	input := &s3.ListObjectsV2Input{
		Bucket:     &w.config.O11yBucketName,
		Prefix:     &prefix,
		StartAfter: &task.StartObjectKey,
	}
	for {
		result, err := w.o11yS3Client.ListObjectsV2(ctx, input)
		if err != nil {
			logger.Error("failed to list objects", zap.Error(err))
			return nil, err
		}

		for _, obj := range result.Contents {
			// skip the directory
			if strings.HasSuffix(*obj.Key, "/") {
				continue
			}
			objects = append(objects, *obj.Key)
			mergeCount++
			if *obj.Key == task.EndObjectKey || mergeCount-task.ObjectCount > taskMismatchThreshold {
				finish = true
				break
			}
		}
		if finish || result.NextContinuationToken == nil || *result.NextContinuationToken == "" {
			break
		}
		input.ContinuationToken = result.NextContinuationToken
	}

	if mergeCount < task.ObjectCount {
		logger.Error("corrupted: data lost")
		return nil, fmt.Errorf("corrupted: data lost")
	}

	if mergeCount-task.ObjectCount > taskMismatchThreshold {
		logger.Error("corrupted: data mismatch too much")
		return nil, fmt.Errorf("corrupted: data mismatch too much")
	}
	if mergeCount != task.ObjectCount {
		logger.Warn("data mismatch, please check it")
	}
	return objects, nil
}

type downloadResult struct {
	name string
	err  error
	body io.ReadCloser
}

// mergeFiles download files concurrently and merge them into one file
// once one file download failed or merge failed, stop the merge process and return the error
func (w *Worker) mergeFiles(ctx context.Context, objects []string,
	task *entity.AuditLogTask, file *os.File) error {
	logger := log.WithContextL(ctx).With(zap.String("task_id", task.TaskId)).
		With(zap.String("cluster_id", task.ClusterId))

	objectsCh := make(chan string, len(objects))
	filesCh := make(chan *downloadResult, len(objects))
	stop := false
	for range downloadConcurrency {
		go w.downloadFile(ctx, objectsCh, filesCh, &stop)
	}
	for _, object := range objects {
		objectsCh <- object
	}
	close(objectsCh)

	// Create a map to store results by object name for ordering
	resultsMap := make(map[string]*downloadResult)
	resultCount := 0

	for result := range filesCh {
		if result.err != nil {
			logger.Error("failed to download file", zap.String("object", result.name), zap.Error(result.err))
			return result.err
		}
		resultsMap[result.name] = result
		resultCount++

		if resultCount == len(objects) {
			close(filesCh)
			break
		}
	}

	// Write files to the output in the correct order
	for _, object := range objects {
		result, ok := resultsMap[object]
		if !ok {
			err := fmt.Errorf("missing file for object: %s", object)
			logger.Error("failed to find downloaded file", zap.String("object", object), zap.Error(err))
			return err
		}
		_, err := io.Copy(file, result.body)
		result.body.Close() //nolint:errcheck // no need check error
		if err != nil {
			logger.Error("failed to copy file", zap.String("object", result.name), zap.Error(err))
			return err
		}
	}

	return nil
}

func (w *Worker) downloadFile(ctx context.Context, objectsCh <-chan string,
	filesCh chan<- *downloadResult, stop *bool) {
	for object := range objectsCh {
		if *stop {
			return
		}

		input := &s3.GetObjectInput{
			Bucket: aws.String(w.config.O11yBucketName),
			Key:    &object,
		}
		var err error
		var result *s3.GetObjectOutput

		for range downloadRetryTimes {
			if result, err = w.o11yS3Client.GetObject(ctx, input); err == nil {
				break
			}
			time.Sleep(downloadRetryInterval)
		}

		filesCh <- &downloadResult{
			name: object,
			err:  err,
			body: func() io.ReadCloser {
				if err == nil {
					return result.Body
				}
				return nil
			}(),
		}

		if err != nil {
			*stop = true
		}
	}
}