package lambda

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/pkg/errors"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
)

// SQSRequestBody represents the body of an SQS message that is actually an http request itself
type SQSRequestBody struct {
	HTTPMethod string            `json:"httpMethod"`
	Path       string            `json:"path"`
	Headers    map[string]string `json:"headers"`
	Body       string            `json:"body"`
}

// ConvertedHTTPRequestAndSQSIdentifiers represents the original SQS message along with it's related HTTP request and
// any errors. This is required since we need to keep track of the original SQS message in case the related HTTP request
// eventually fails
type ConvertedHTTPRequestAndSQSIdentifiers struct {
	Req               *http.Request
	SQSReceiptHandle  string
	SQSMessageId      string
	SQSEventSourceARN string
	Error             error
}

// SafeSQSDeleteMessageBatchRequestEntryMap is a thread safe map that allows us to map SQS ARNs to a list of
// DeleteMessageBatchRequestEntry
type SafeSQSDeleteMessageBatchRequestEntryMap struct {
	mutex                    sync.Mutex
	sqsDeleteMessageBatchMap map[string][]*sqs.DeleteMessageBatchRequestEntry
}

// AddEntry adds an entry to the map. If the key previously existed, the new value will be appended to the existing
// DeleteMessageBatchRequestEntry list. Otherwise, a new list will be created for the provided key, and the  single
// entry will be added to the list
func (m *SafeSQSDeleteMessageBatchRequestEntryMap) AddEntry(arn string, batchEntry *sqs.DeleteMessageBatchRequestEntry) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	// Add the request to the appropriate ARN map
	if existingList, ok := m.sqsDeleteMessageBatchMap[arn]; ok {
		m.sqsDeleteMessageBatchMap[arn] = append(existingList, batchEntry)
	} else {
		m.sqsDeleteMessageBatchMap[arn] = append([]*sqs.DeleteMessageBatchRequestEntry{}, batchEntry)
	}
}

// GetMapSize gets the size of the map
func (m *SafeSQSDeleteMessageBatchRequestEntryMap) GetMapSize() int {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	return len(m.sqsDeleteMessageBatchMap)
}

// GetMapCopy gets a thread unsafe copy of the map (just a standard map)
func (m *SafeSQSDeleteMessageBatchRequestEntryMap) GetMapCopy() map[string][]*sqs.DeleteMessageBatchRequestEntry {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	mapCopy := make(map[string][]*sqs.DeleteMessageBatchRequestEntry, len(m.sqsDeleteMessageBatchMap))
	for k, v := range m.sqsDeleteMessageBatchMap {
		mapCopy[k] = v
	}
	return mapCopy
}

// NewSQSRequest returns a list of new http.Requests from the given list of SQS messages (SQS "Records")
func NewSQSRequest(ctx context.Context, records []events.SQSMessage) []ConvertedHTTPRequestAndSQSIdentifiers {
	// Construct the final lists
	numRequests := len(records)
	convertedSQSList := make([]ConvertedHTTPRequestAndSQSIdentifiers, 0, numRequests)

	for _, sqsRecord := range records {

		// Attempt to parse the body of the SQS event as a request
		var convertedBody SQSRequestBody
		err := json.Unmarshal([]byte(sqsRecord.Body), &convertedBody)
		if err != nil {
			convertedSQSList = append(convertedSQSList, createConvertedHTTPRequestAndSQSIdentifiers(
				nil, sqsRecord, errors.Wrap(err, "failed to unmarshal SQS body")))
			continue
		}

		// Construct the new URL path
		u, err := url.Parse(convertedBody.Path)
		if err != nil {
			convertedSQSList = append(convertedSQSList, createConvertedHTTPRequestAndSQSIdentifiers(
				nil, sqsRecord, errors.Wrap(err, "failed to parse SQS body URL path")))
			continue
		}

		// Construct a new HTTP Request
		req, err := http.NewRequest(convertedBody.HTTPMethod, u.String(), strings.NewReader(convertedBody.Body))
		if err != nil {
			convertedSQSList = append(convertedSQSList, createConvertedHTTPRequestAndSQSIdentifiers(
				nil, sqsRecord, errors.Wrap(err, "failed to create new request from SQS body")))
			continue
		}

		// Add a remote address header
		// Since we are dealing with an SQS queue, use the source ARN instead of an IP here so we note what queue the
		// request came from
		req.RemoteAddr = sqsRecord.EventSourceARN

		// Add headers
		for k, v := range convertedBody.Headers {
			req.Header.Set(k, v)
		}

		// Add the content-length header
		if req.Header.Get("Content-Length") == "" && convertedBody.Body != "" {
			req.Header.Set("Content-Length", strconv.Itoa(len(convertedBody.Body)))
		}

		// Add additional custom fields
		req.Header.Set("X-Request-Id", sqsRecord.MessageId)

		// Add X-Ray support
		if traceID := ctx.Value("x-amzn-trace-id"); traceID != nil {
			req.Header.Set("X-Amzn-Trace-Id", fmt.Sprintf("%v", traceID))
		}

		// Add the host (if possible)
		req.URL.Host = req.Header.Get("Host")
		req.Host = req.URL.Host

		// Add the completed request and note that there are no errors
		convertedSQSList = append(convertedSQSList, createConvertedHTTPRequestAndSQSIdentifiers(req, sqsRecord, nil))
	}

	return convertedSQSList
}

// createConvertedHTTPRequestAndSQSIdentifiers is a helper method to construct ConvertedHTTPRequestAndSQSIdentifiers
// from the original SQS message and the related http.Request
func createConvertedHTTPRequestAndSQSIdentifiers(req *http.Request, msg events.SQSMessage, err error) ConvertedHTTPRequestAndSQSIdentifiers {
	return ConvertedHTTPRequestAndSQSIdentifiers{
		Req:               req,
		SQSReceiptHandle:  msg.ReceiptHandle,
		SQSMessageId:      msg.MessageId,
		SQSEventSourceARN: msg.EventSourceARN,
		Error:             err,
	}
}
