package jobqueue

// jobqueue is a package based on SQS

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/account/commons/config"
	"github.com/tidbcloud/tiinf/monitor/tilog"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"go.uber.org/zap"
)

const (
	jobNameAttr = "JOB_NAME"
	SUCCEED     = "SUCCEED"
	FAILED      = "FAILED"
)

var (
	sqsInitOnce     sync.Once
	defaultJobQueue *JobQueue

	jobQueueHistogram = promauto.NewHistogramVec(prometheus.HistogramOpts{
		Name: "jobqueue_histogram",
		Help: "jobqueue histogram",
	}, []string{"job_name", "status"})

	ErrJobAlreadyRegistered = errors.New("Job already registered")
	ErrBadJob               = errors.New("Bad job")
	ErrJobHandlerNotFound   = errors.New("Job handler not found")
)

type Handler func(ctx context.Context, payload []byte) error

type SQSIface interface {
	SendMessage(input *sqs.SendMessageInput) (*sqs.SendMessageOutput, error)
	ReceiveMessage(input *sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error)
	DeleteMessage(input *sqs.DeleteMessageInput) (*sqs.DeleteMessageOutput, error)
}

// GetQueueArn gets the ARN of a queue based on its URL
func GetQueueArn(queueURL *string) *string {
	// snippet-start:[sqs.go.get_queue_url.arn]
	parts := strings.Split(*queueURL, "/")
	subParts := strings.Split(parts[2], ".")

	arn := "arn:aws:" + subParts[0] + ":" + subParts[1] + ":" + parts[3] + ":" + parts[4]
	// snippet-end:[sqs.go.get_queue_url.arn]

	return &arn
}

// ConfigureDeadLetterQueue configures an Amazon SQS queue for messages that could not be delivered to another queue
// Inputs:
//
//	sess is the current session, which provides configuration for the SDK's service clients
//	deadLetterQueueARN is the ARN of the dead-letter queue
//	queueURL is the URL of the queue that did not get messages
//
// Output:
//
//	If success, the URL of the queue and nil
//	Otherwise, an empty string and an error from the call to json.Marshal or SetQueueAttributes
//
// from: https://github.com/awsdocs/aws-doc-sdk-examples/blob/6f27ad9bab61e32c2a7417751fafd43c8ae8d454/go/sqs/DeadLetterQueue/DeadLetterQueue.go#L70
func ConfigureDeadLetterQueue(sess *session.Session, dlQueueARN *string, queueURL *string) error {
	// Create an SQS service client
	svc := sqs.New(sess)

	// Our redrive policy for our queue
	// snippet-start:[sqs.go.dead_letter_queue.policy]
	policy := map[string]string{
		"deadLetterTargetArn": *dlQueueARN,
		"maxReceiveCount":     "10", // NOTE: so the maxium job retry count is 10 times, after that, the job will be put into dead letter queue
	}
	// snippet-end:[sqs.go.dead_letter_queue.policy]

	// Marshal policy for SetQueueAttributes
	// snippet-start:[sqs.go.dead_letter_queue.marshall]
	b, err := json.Marshal(policy)
	// snippet-end:[sqs.go.dead_letter_queue.marshall]
	if err != nil {
		return err
	}

	// snippet-start:[sqs.go.dead_letter_queue.set_attributes]
	_, err = svc.SetQueueAttributes(&sqs.SetQueueAttributesInput{
		QueueUrl: queueURL,
		Attributes: map[string]*string{
			sqs.QueueAttributeNameRedrivePolicy: aws.String(string(b)),
		},
	})
	// snippet-end:[sqs.go.dead_letter_queue.set_attributes]
	if err != nil {
		return err
	}

	return nil
}

type JobQueue struct {
	sqsSvc      SQSIface
	sqsQueueURL string
	timeout     int64
	async       bool

	mapper map[string]Handler
}

func NewJobQueue(ctx context.Context, queueURL, dlQueueURL, region string, timeout int64) (*JobQueue, error) {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		Config: aws.Config{
			Region: &region,
		},
		SharedConfigState: session.SharedConfigStateFromEnv,
	}))

	sqsSvc := sqs.New(sess)
	if !config.IsLocalEnv() {
		arn := GetQueueArn(&dlQueueURL)
		err := ConfigureDeadLetterQueue(sess, arn, &queueURL)
		if err != nil {
			tilog.Error(ctx, "configure dead letter queue", zap.Error(err))
			return nil, err
		}
	}

	return &JobQueue{
		sqsSvc:      sqsSvc,
		sqsQueueURL: queueURL,
		timeout:     timeout,
		mapper:      make(map[string]Handler),
	}, nil
}

// Enqueue pushes a job to the queue
func (q *JobQueue) Enqueue(ctx context.Context, jobName string, v interface{}) error {
	var (
		payload []byte
		err     error
	)
	if payload, err = json.Marshal(v); err != nil {
		return err
	}

	_, err = q.sqsSvc.SendMessage(&sqs.SendMessageInput{
		//MessageGroupId: aws.String("account"),
		MessageAttributes: map[string]*sqs.MessageAttributeValue{
			jobNameAttr: {DataType: aws.String("String"), StringValue: aws.String(jobName)},
		},
		MessageBody: aws.String(string(payload)),
		QueueUrl:    aws.String(q.sqsQueueURL),
	})
	tilog.Info(ctx, "enqueue message", zap.String("job_name", jobName), zap.Any("payload", v), zap.Error(err))

	return err
}

// Register register a handler for a job
func (q *JobQueue) Register(jobName string, handler Handler) error {
	if _, ok := q.mapper[jobName]; ok {
		return fmt.Errorf("register job %s with handler %v failed: %w", jobName, handler, ErrJobAlreadyRegistered)
	}

	q.mapper[jobName] = handler

	return nil
}

func (q *JobQueue) handleMsg(ctx context.Context, msg *sqs.Message) error {
	if msg.Body == nil {
		tilog.Error(ctx, "message body is nil", zap.Any("message", msg))
		return fmt.Errorf("message body is nil: %w", ErrBadJob)
	}

	if msg.MessageAttributes == nil || msg.MessageAttributes[jobNameAttr] == nil || msg.MessageAttributes[jobNameAttr].StringValue == nil {
		tilog.Error(ctx, "message attributes is nil", zap.Any("message", msg))
		return fmt.Errorf("message attributes is nil: %w", ErrBadJob)
	}

	jobName := *msg.MessageAttributes[jobNameAttr].StringValue
	handler, ok := q.mapper[jobName]
	if !ok {
		tilog.Error(ctx, "handler not found", zap.String("jobName", jobName))
		return fmt.Errorf("handler of %s not found: %w", jobName, ErrJobHandlerNotFound)
	}

	tilog.Info(ctx, "handler start to executing", zap.String("jobName", jobName), zap.Any("msg", msg))
	start := time.Now()
	err := handler(ctx, []byte(*msg.Body))
	duration := time.Since(start)
	status := SUCCEED
	if err != nil {
		status = FAILED
	}
	jobQueueHistogram.WithLabelValues(jobName, status).Observe(duration.Seconds())

	if err != nil {
		tilog.Warn(ctx, "handler executed failed", zap.Error(err), zap.String("jobName", jobName), zap.Any("msg", msg))
		return err
	}
	tilog.Info(ctx, "handler executed succeeded", zap.String("jobName", jobName), zap.Any("msg", msg))

	_, err = q.sqsSvc.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      aws.String(q.sqsQueueURL),
		ReceiptHandle: msg.ReceiptHandle,
	})
	tilog.Info(ctx, "delete message", zap.String("receiptHandle", *msg.ReceiptHandle), zap.Error(err), zap.String("jobName", jobName))
	return err
}

func (q *JobQueue) doPullAndExec(ctx context.Context) error {
	defer func() {
		// recover from panic if one occurred. Set err to nil otherwise.
		if err := recover(); err != nil {
			tilog.Error(ctx, "panic occurred", zap.Any("panic", err))
		}
	}()

	msgResult, err := q.sqsSvc.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            aws.String(q.sqsQueueURL),
		MaxNumberOfMessages: aws.Int64(10), // https://docs.aws.amazon.com/lambda/latest/operatorguide/sqs-retries.html
		VisibilityTimeout:   aws.Int64(q.timeout),
		WaitTimeSeconds:     aws.Int64(10),
	})
	if err != nil {
		tilog.Error(ctx, "error occurred", zap.Error(err))
		return err
	}

	tilog.Info(ctx, "received message", zap.Int("count", len(msgResult.Messages)))

	for _, msg := range msgResult.Messages {
		if q.async {
			go q.handleMsg(ctx, msg) //nolint
		} else {
			q.handleMsg(ctx, msg) //nolint
		}
	}

	return nil
}

// StartPolling pulls messages from the queue and calls the handler
func (q *JobQueue) StartPolling(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			// do the real work in following code
		}

		err := q.doPullAndExec(ctx)
		if err != nil {
			time.Sleep(time.Minute)
		}
	}
}

func (q *JobQueue) EnableAsync() {
	q.async = true
}

/////////////// global methods

func InitDefaultJobQueue(ctx context.Context, queueURL, dlQueueURL, region string, timeout int64) error {
	var err error
	sqsInitOnce.Do(func() {
		defaultJobQueue, err = NewJobQueue(ctx, queueURL, dlQueueURL, region, timeout)
	})

	return err
}

func Enqueue(ctx context.Context, jobName string, v interface{}) error {
	if defaultJobQueue == nil {
		return errors.New("jobqueue not initialized")
	}

	return defaultJobQueue.Enqueue(ctx, jobName, v)
}

func Register(jobName string, handler Handler) error {
	if defaultJobQueue == nil {
		return errors.New("jobqueue not initialized")
	}

	return defaultJobQueue.Register(jobName, handler)
}

func MustRegister(jobName string, handler Handler) {
	if err := Register(jobName, handler); err != nil {
		panic(err)
	}
}

func StartPolling(ctx context.Context) error {
	if defaultJobQueue == nil {
		return errors.New("jobqueue not initialized")
	}

	return defaultJobQueue.StartPolling(ctx)
}

func EnableAsync() {
	if defaultJobQueue == nil {
		return
	}

	defaultJobQueue.EnableAsync()
}

// CatchError catch panic and error result, print the error and ignore it, so the job framework
// will delete the SQS message and never retry the job.
func CatchError(ctx context.Context, fn func() error) error {
	defer func() {
		// recover from panic if one occurred. Set err to nil otherwise.
		if err := recover(); err != nil {
			tilog.Error(ctx, "panic occurred", zap.Any("panic", err))
		}
	}()

	err := fn()
	if err != nil {
		tilog.Error(ctx, "error occurred", zap.Error(err))
	}

	return nil
}