package input

import (
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"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/aws/aws-sdk-go/service/sqs/sqsiface"
	"github.com/jmespath/go-jmespath"
	log "github.com/sirupsen/logrus"

	"github.com/AdRoll/baker"
	"github.com/AdRoll/baker/input/inpututils"
	"github.com/AdRoll/baker/pkg/awsutils"
)

var SQSDesc = baker.InputDesc{
	Name:   "SQS",
	New:    NewSQS,
	Config: &SQSConfig{},
	Help: `This input listens on multiple SQS queues for paths to S3 files.
It can be used with SQS queues subscribed to SNS topics (with raw_message_delivery subscription) and supports arbitrary payload (plain or json).
It never exits.

Supported formats (MessageFormat):
 - "plain": the message payload is a plain S3 path.
 - "sns": the payload is an SNS notification
 - "s3::ObjectCreated": the payload is the message generated by an S3::ObjectCreated notification
 - "json": the message is an arbitrary JSON string, the S3 path is extracted from it via a [jmespath](https://jmespath.org/) expression specified in MessageExpression.
 `,
}

type SQSConfig struct {
	AwsRegion         string   `help:"AWS region to connect to" default:"us-west-2"`
	Bucket            string   `help:"S3 Bucket to use if paths do not have one" default:""`
	QueuePrefixes     []string `help:"Prefixes of the names of the SQS queues to monitor"`
	QueueNames        []string `help:"Names of the SQS queues to monitor"`
	MessageFormat     string   `help:"SQS message format. See help string for supported formats" default:"sns"`
	MessageExpression string   `help:"The expression to extract an S3 path from arbitrary message formats"`
	FilePathFilter    string   `help:"If provided, will only use S3 files with the given path."`

	format sqsFormatType
}

func (cfg *SQSConfig) fillDefaults() error {
	if cfg.AwsRegion == "" {
		cfg.AwsRegion = "us-west-2"
	}

	switch sqsFormatType(strings.ToLower(cfg.MessageFormat)) {
	case sqsFormatSNS, "":
		cfg.format = sqsFormatJSON
		cfg.MessageExpression = "Message"
	case sqsFormatPlain:
		cfg.format = sqsFormatPlain
	case sqsFormatS3ObjectCreated:
		cfg.format = sqsFormatJSON
		cfg.MessageExpression = `Records[*].join('/',['s3:/', s3.bucket.name, s3.object.key]) | [0]`
	case sqsFormatJSON:
		cfg.format = sqsFormatJSON
		if cfg.MessageExpression == "" {
			return fmt.Errorf("MessageExpression is required with json MessageFormat")
		}
	default:
		return fmt.Errorf("unknown MessageFormat %q", cfg.MessageFormat)
	}

	return nil
}

type sqsFormatType string

const (
	sqsFormatPlain           sqsFormatType = "plain"
	sqsFormatSNS             sqsFormatType = "sns"
	sqsFormatS3ObjectCreated sqsFormatType = "s3::objectcreated"
	sqsFormatJSON            sqsFormatType = "json"
)

type SQS struct {
	s3Input *inpututils.S3Input

	Cfg  *SQSConfig
	svc  sqsiface.SQSAPI
	done chan bool

	filepathRx *regexp.Regexp
	parse      parseSQSMessageFunc
}

func NewSQS(cfg baker.InputParams) (baker.Input, error) {
	dcfg := cfg.DecodedConfig.(*SQSConfig)
	if err := dcfg.fillDefaults(); err != nil {
		return nil, fmt.Errorf("SQS: configuration error: %v", err)
	}

	var filepathRx *regexp.Regexp
	if dcfg.FilePathFilter != "" {
		var err error
		if filepathRx, err = regexp.Compile(dcfg.FilePathFilter); err != nil {
			return nil, fmt.Errorf("SQS: can't compile FilePathFilter: %v", err)
		}
	}

	sess, err := session.NewSession(&aws.Config{Region: aws.String(dcfg.AwsRegion)})
	if err != nil {
		return nil, fmt.Errorf("SQS: can't create aws session: %v", err)
	}

	parseFunc, err := sqsParseFunction(dcfg)
	if err != nil {
		return nil, fmt.Errorf("SQS: can't configure message parsing")
	}

	s3Input, err := inpututils.NewS3Input(dcfg.AwsRegion, dcfg.Bucket)
	if err != nil {
		return nil, fmt.Errorf("SQS: %v", err)
	}

	if len(dcfg.QueuePrefixes) == 0 && len(dcfg.QueueNames) == 0 {
		return nil, fmt.Errorf("SQS: QueuePrefixes or QueueNames must be set")
	}

	sqs := &SQS{
		s3Input:    s3Input,
		Cfg:        dcfg,
		svc:        sqs.New(sess),
		filepathRx: filepathRx,
		done:       make(chan bool),
		parse:      parseFunc,
	}
	return sqs, nil
}

type parseSQSMessageFunc func(string) (string, error)

func sqsParseFunction(cfg *SQSConfig) (parseSQSMessageFunc, error) {
	var f parseSQSMessageFunc

	switch cfg.format {
	case sqsFormatPlain:
		f = func(s string) (string, error) { return s, nil }

	case sqsFormatJSON:
		jsexpr, err := jmespath.Compile(cfg.MessageExpression)
		if err != nil {
			return nil, fmt.Errorf("can't compile jmespath expression: %v", err)
		}
		f = func(s string) (string, error) {
			dec := json.NewDecoder(strings.NewReader(s))
			var js interface{}
			if err := dec.Decode(&js); err != nil {
				return "", fmt.Errorf("can't decode json from SQS message")
			}
			iface, err := jsexpr.Search(js)
			if err != nil {
				return "", fmt.Errorf("can't extract path from SQS message")
			}
			if iface == nil {
				return "", fmt.Errorf("can't find S3 path field in SQS message")
			}
			s3Path, ok := iface.(string)
			if !ok {
				return "", fmt.Errorf("extracted field is not a string")
			}

			u, err := url.Parse(s3Path)
			if err != nil {
				return "", fmt.Errorf("can't parse received URL in SQS message: %v", err)
			}
			// If bucket was not provided, use it from the S3 path.
			if cfg.Bucket == "" {
				return s3Path, nil
			}
			return u.Path[1:], nil
		}
	default:
		panic("unexpected message format: " + cfg.format)
	}
	return f, nil
}

// pollQueue polls the given queue as long as the given context is alive.
func (s *SQS) pollQueue(ctx context.Context, sqsurl string) {
	ctxLog := log.WithFields(log.Fields{"f": "SQS.pollQueue", "url": sqsurl})
	backoff := awsutils.DefaultBackoff
	for {
		resp, err := s.svc.ReceiveMessageWithContext(ctx, &sqs.ReceiveMessageInput{
			QueueUrl:        aws.String(sqsurl),
			WaitTimeSeconds: aws.Int64(20),
			// We ask only for 1 message at a time, because the
			// parseFile() call below could block, and we want to
			// receive messages and not process them immediately,
			// or they could get rescheduled to other readers.
			MaxNumberOfMessages: aws.Int64(1),
		})
		if ctx.Err() == context.Canceled || ctx.Err() == context.DeadlineExceeded {
			return
		}

		if err != nil {
			ctxLog.WithError(err).Error("error from ReceiveMessage")
			time.Sleep(backoff.Duration())
			continue
		}
		backoff.Reset()

		for _, msg := range resp.Messages {
			s3FilePath, err := s.parse(*msg.Body)
			if err != nil {
				ctxLog.WithError(err).Error("error parsing message")
				continue
			}

			// Try to unescape URL. Voluntarily ignore an error here since the
			// result might be a valid URL anyway.
			s3FilePath, _ = url.QueryUnescape(s3FilePath)

			// Skip the file if it doesn't match the filter provided.
			if s.filepathRx == nil || s.filepathRx.MatchString(s3FilePath) {
				// FIXME: we should check if the bucket matches what was configured
				// or even better, change s3Input to not be limited to a single bucket
				s.s3Input.ParseFile(s3FilePath)
			}

			_, err = s.svc.DeleteMessageWithContext(ctx, &sqs.DeleteMessageInput{
				QueueUrl:      aws.String(sqsurl),
				ReceiptHandle: msg.ReceiptHandle,
			})
			if ctx.Err() == context.Canceled || ctx.Err() == context.DeadlineExceeded {
				return
			}
			if err != nil {
				ctxLog.WithError(err).Error("error from DeleteMessage")
			}
		}
	}
}

func (s *SQS) Run(inch chan<- *baker.Data) error {
	s.s3Input.SetOutputChannel(inch)

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

	queueUrls := []string{}

	for _, queueName := range s.Cfg.QueueNames {

		resp, err := s.svc.GetQueueUrlWithContext(ctx, &sqs.GetQueueUrlInput{
			QueueName: aws.String(queueName),
		})
		if err != nil {
			return err
		}
		queueUrls = append(queueUrls, *resp.QueueUrl)
	}

	for _, prefix := range s.Cfg.QueuePrefixes {

		resp, err := s.svc.ListQueuesWithContext(ctx, &sqs.ListQueuesInput{
			QueueNamePrefix: aws.String(prefix),
		})
		if err != nil {
			return err
		}

		for _, url := range resp.QueueUrls {
			queueUrls = append(queueUrls, *url)
		}
	}

	var wg sync.WaitGroup
	for _, url := range queueUrls {
		wg.Add(1)
		go func(url string) {
			defer wg.Done()

			s.pollQueue(ctx, url)
		}(url)
	}

	// The correct order of operation to cleanly stop the whole pipeline is the
	// following:
	//  - first we close the 'done' channel, this in turns cancel polling via
	//    context cancellation
	//  - we then wait for all the polling goroutines to end.
	//  - After this point we are guaranteed to not receive any more files so
	//    we notify the embedded S3input.
	//  - now we ask S3Input to stop as soon as it has finished processing files
	//  - finally Run exits after being signaled from S3Input that we can
	<-s.done
	cancel()
	wg.Wait()
	s.s3Input.NoMoreFiles()
	s.s3Input.Stop()
	<-s.s3Input.Done
	return nil
}

func (s *SQS) Stop() {
	// See the comment at the end of the Run method for details about
	// how this brings the whole topology down cleanly.
	close(s.done)
}

func (s *SQS) Stats() baker.InputStats {
	return s.s3Input.Stats()
}

func (s *SQS) FreeMem(data *baker.Data) {
	s.s3Input.FreeMem(data)
}
