package lambda

import (
	"context"
	"net/http"

	"encoding/json"
	"fmt"
	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/pkg/errors"
	"os"
	"sync"
	"sync/atomic"
)

const (
	sqsEventSource      = "aws:sqs"
	recordsKey          = "Records"
	maxSQSClientRetries = 3
)

// ListenAndServe is a drop-in replacement for
// http.ListenAndServe for use within AWS Lambda.
//
// ListenAndServe always returns a non-nil error.
func ListenAndServe(addr string, h http.Handler) error {
	if h == nil {
		h = http.DefaultServeMux
	}

	// Make an SQS client (required for some SQS failure cases)
	sqsClient := sqs.New(session.Must(session.NewSession(&aws.Config{
		Region:     aws.String(os.Getenv("AWS_REGION")),
		MaxRetries: aws.Int(maxSQSClientRetries),
	})))

	lambda.Start(func(ctx context.Context, e map[string]interface{}) (events.APIGatewayProxyResponse, error) {
		// Right now, the only event with Records that we accept is SQS messages, assume "Records" implies SQS
		// In the future, we will want to dig a bit deeper into this map to verify additional components before t
		if val, ok := e[recordsKey]; ok {
			return processSQSEvent(sqsClient, h, ctx, val)
		} else {
			// If there wasn't a recordsKey, assume the event should be processed as an API Gateway request
			return processAPIGatewayEvent(h, ctx, e)
		}
	})

	return nil
}

// processSQSEvent is a helper method that takes the input interface and processes it as an SQS Event
func processSQSEvent(sqsClient *sqs.SQS, h http.Handler, ctx context.Context, e interface{}) (events.APIGatewayProxyResponse, error) {

	// Marshal the interface to JSON and then unmarshal the JSON to a list of SQSMessage
	var records []events.SQSMessage
	recordsJsonString, err := json.Marshal(e)
	if err != nil {
		return events.APIGatewayProxyResponse{}, errors.Wrap(err, "failed to marshal map")
	}
	err = json.Unmarshal(recordsJsonString, &records)
	if err != nil || len(records) == 0 || records[0].EventSource != sqsEventSource {
		return events.APIGatewayProxyResponse{}, errors.Wrap(err, "failed to unmarshal to an SQSMessage list")
	}

	// Convert the SQS messages into separate Lambda HTTP requests with their related SQS messages
	convertedHTTPRequestAndSQSIdentifiers := NewSQSRequest(ctx, records)

	// Construct a thread safe container for the map of the SQS ARN to delete requests (for messages that are safe to
	// delete). This is required since we might have received messages from different queues, and we also need to ensure
	// we only delete messages that were successful
	deleteSQS := SafeSQSDeleteMessageBatchRequestEntryMap{
		mutex:                    sync.Mutex{},
		sqsDeleteMessageBatchMap: make(map[string][]*sqs.DeleteMessageBatchRequestEntry),
	}

	// Create a WaitGroup to ensure every SQS request is processed before the lambda returns a responses
	var wg sync.WaitGroup
	wg.Add(len(convertedHTTPRequestAndSQSIdentifiers))

	// Keep track if any requests failed
	var atomicFailedRequestNum int32

	for _, convertedRequest := range convertedHTTPRequestAndSQSIdentifiers {
		// Concurrently call the lambda for every SQS message
		go func(convertedRequest ConvertedHTTPRequestAndSQSIdentifiers) {
			defer wg.Done()

			// If the request had any errors during the conversion process, don't attempt to send it to the server
			if convertedRequest.Error != nil {
				atomic.AddInt32(&atomicFailedRequestNum, 1)
				return
			}

			// Otherwise, it should be safe to send
			w := NewGatewayResponse()
			h.ServeHTTP(w, convertedRequest.Req.WithContext(ctx))
			resp := w.End()

			// If we didn't get a 2xx response, then we should assume the request failed
			if resp.StatusCode >= 300 {
				atomic.AddInt32(&atomicFailedRequestNum, 1)
			} else {
				// Otherwise, the request succeeded. This means we should remove the SQS message from the queue so
				// we don't retry it. Therefore, construct a DeleteMessageBatchRequestEntry for the message
				successfulSQS := &sqs.DeleteMessageBatchRequestEntry{}
				successfulSQS = successfulSQS.SetId(convertedRequest.SQSMessageId)
				successfulSQS = successfulSQS.SetReceiptHandle(convertedRequest.SQSReceiptHandle)

				// Add the request to the appropriate ARN map
				deleteSQS.AddEntry(convertedRequest.SQSEventSourceARN, successfulSQS)
			}
		}(convertedRequest)
	}
	wg.Wait()

	// If no request failed, return a happy response. If SQS doesn't get back an error, it assumes every message worked
	// and will delete every SQS message sent in the batch for us
	failedRequestNum := atomic.LoadInt32(&atomicFailedRequestNum)
	if failedRequestNum == 0 {
		return events.APIGatewayProxyResponse{}, nil
	} else {
		// Otherwise, at least one message failed. This means we have to manually call SQS to delete any successful
		// messages and then error so that any remaining (failed) messages are retried according to the SQS queue's
		// retry policy
		if deleteSQS.GetMapSize() > 0 {
			// Use the sqs client logger
			logger := sqsClient.Config.Logger

			// Time for more concurrent requests to delete batches
			wg.Add(deleteSQS.GetMapSize())

			// For every queue ARN we have, try deleting its successful messages
			for queueARNString, batchEntries := range deleteSQS.GetMapCopy() {
				go func(queueARNString string, batchEntries []*sqs.DeleteMessageBatchRequestEntry) {
					defer wg.Done()

					// Get the queue name and account id from the ARN and call GetQueueUrlWithContext to get the Queue URL
					// SQS messages are modeled as arn:aws:sqs:region:account-id:queuename
					// https://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html#genref-arns
					queueARN, err := arn.Parse(queueARNString)
					if err != nil {
						// If we don't have an ARN, there isn't much we can do
						logger.Log(fmt.Sprintf("Cannot delete successful SQS messages. Invalid queue ARN given for %v. err: %v", batchEntries, err))
						return
					}
					getQueueInput := &sqs.GetQueueUrlInput{}
					getQueueInput.SetQueueName(queueARN.Resource)
					getQueueInput.SetQueueOwnerAWSAccountId(queueARN.AccountID)
					getQueueResp, err := sqsClient.GetQueueUrlWithContext(ctx, getQueueInput)

					if err != nil || getQueueResp == nil {
						// If the error is not nil or the response was empty, we can't delete the batch
						// This is technically fine, since SQS -> Lambda should be idempotent
						logger.Log(fmt.Sprintf("Cannot delete successful SQS messages. Failed GetQueueUrl call for %v. err: %v", batchEntries, err))
						return
					}

					// Now, delete the batch
					delBatchInput := &sqs.DeleteMessageBatchInput{}
					delBatchInput.SetEntries(batchEntries)
					delBatchInput.SetQueueUrl(*getQueueResp.QueueUrl)
					// We technically don't care about the response, since a failed batch will just retry, and these
					// should be idempotent. However, we'll just log failures so they're easier to debug later on
					delMessageResp, err := sqsClient.DeleteMessageBatchWithContext(ctx, delBatchInput)
					if len(delMessageResp.Failed) > 0 || err != nil {
						logger.Log(fmt.Sprintf("Cannot delete successful SQS messages. Failed DeleteMessageBatch call for %v. err: %v", delMessageResp.Failed, err))
					}
				}(queueARNString, batchEntries)
			}
		}
		wg.Wait()
		return events.APIGatewayProxyResponse{}, errors.Errorf("at least one SQS message failed")
	}
}

// processAPIGatewayEvent is a helper method that takes the input map and processes it as APIGatewayProxyRequest event
func processAPIGatewayEvent(h http.Handler, ctx context.Context, e map[string]interface{}) (events.APIGatewayProxyResponse, error) {

	// Marshal the map to JSON and then unmarshal the JSON to an APIGatewayProxyRequest
	var apiGatewayRequest events.APIGatewayProxyRequest
	jsonString, err := json.Marshal(e)
	if err != nil {
		return events.APIGatewayProxyResponse{}, errors.Wrap(err, "failed to marshal map")
	}
	err = json.Unmarshal(jsonString, &apiGatewayRequest)
	if err != nil {
		return events.APIGatewayProxyResponse{}, errors.Wrap(err, "failed to unmarshal to APIGatewayProxyRequest")
	}

	r, err := NewGatewayRequest(ctx, apiGatewayRequest)
	if err != nil {
		return events.APIGatewayProxyResponse{}, err
	}

	w := NewGatewayResponse()
	h.ServeHTTP(w, r.WithContext(ctx))
	return w.End(), nil
}
