package lambda

import (
	"net/http"

	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/lambda/lambdaiface"
	"github.com/twitchtv/twirp"
)

const (
	contentTypeProto = "application/protobuf"
	contentTypeJSON  = "application/json"
)

// Ref https://docs.aws.amazon.com/lambda/latest/dg/API_Invoke.html#API_Invoke_Errors
var awserrCodeToHttpStatusCode = map[string]int{
	lambda.ErrCodeEC2AccessDeniedException:             502,
	lambda.ErrCodeEC2ThrottledException:                502,
	lambda.ErrCodeEC2UnexpectedException:               502,
	lambda.ErrCodeENILimitReachedException:             502,
	lambda.ErrCodeInvalidParameterValueException:       400,
	lambda.ErrCodeInvalidRequestContentException:       400,
	lambda.ErrCodeInvalidRuntimeException:              502,
	lambda.ErrCodeInvalidSecurityGroupIDException:      502,
	lambda.ErrCodeInvalidSubnetIDException:             502,
	lambda.ErrCodeInvalidZipFileException:              502,
	lambda.ErrCodeKMSAccessDeniedException:             502,
	lambda.ErrCodeKMSDisabledException:                 502,
	lambda.ErrCodeKMSInvalidStateException:             502,
	lambda.ErrCodeKMSNotFoundException:                 502,
	lambda.ErrCodeRequestTooLargeException:             400,
	lambda.ErrCodeResourceNotFoundException:            404,
	lambda.ErrCodeServiceException:                     503,
	lambda.ErrCodeSubnetIPAddressLimitReachedException: 502,
	lambda.ErrCodeTooManyRequestsException:             429,
	lambda.ErrCodeUnsupportedMediaTypeException:        400,
}

// Client is the lambda implementation of the Twirp HTTPClient interface, it's job is just to know how to invoke
// twirp lambdas
type Client struct {
	lambdaAPI  lambdaiface.LambdaAPI
	lambdaName string
}

// lambdaErrorJSON represents an error response from AWS Lambda
type lambdaErrorJSON struct {
	Message    string   `json:"errorMessage"`
	Type       string   `json:"errorType"`
	StackTrace []string `json:"stackTrace"`
}

// twerrJSON represents a Twirp error's JSON format
type twerrJSON struct {
	Code string            `json:"code"`
	Msg  string            `json:"msg"`
	Meta map[string]string `json:"meta,omitempty"`
}

// NewClient creates a new lambda twirp client that implements twirp.HTTPClient for use with Twirp Clients
func NewClient(awsLambda lambdaiface.LambdaAPI, lambdaName string) *Client {
	return &Client{
		lambdaAPI:  awsLambda,
		lambdaName: lambdaName,
	}
}

// LambdaName returns the underlying lambda we are calling
// The AWS Go SDK allows for this to be either the function name (ex:
// haberdasher) or the lambda arn (ex:
// arn:aws:lambda:us-west-2:123456789012:function:haberdasher).
//
// If you're relying on this to identify your lambda, this could be either.
func (l *Client) LambdaName() string {
	return l.lambdaName
}

func (l *Client) invokeInputFromHTTPRequest(req *http.Request) (*lambda.InvokeInput, error) {
	header := req.Header.Get("Content-Type")
	i := strings.Index(header, ";")
	if i == -1 {
		i = len(header)
	}
	contentType := strings.TrimSpace(strings.ToLower(header[:i]))

	var payload string
	if contentType == contentTypeJSON {
		buf := new(bytes.Buffer)
		_, err := buf.ReadFrom(req.Body)
		if err != nil {
			return nil, err
		}
		payload = buf.String()
	} else if contentType == contentTypeProto { //
		buf := new(bytes.Buffer)
		_, err := buf.ReadFrom(req.Body)
		if err != nil {
			return nil, err
		}
		payload = base64.StdEncoding.EncodeToString(buf.Bytes())
	} else {
		return nil, fmt.Errorf("labmbda_transport: unknown content type %q for Client", contentType)
	}

	request := events.APIGatewayProxyRequest{
		HTTPMethod:      req.Method,
		Path:            req.URL.Path,
		Headers:         headersToMap(req.Header),
		Body:            payload,
		IsBase64Encoded: contentType == contentTypeProto,
	}

	encoded, err := json.Marshal(request)
	if err != nil {
		return nil, err
	}

	return &lambda.InvokeInput{
		FunctionName: aws.String(l.lambdaName),
		Payload:      encoded,
	}, nil
}

func headersToMap(headers http.Header) map[string]string {
	hMap := make(map[string]string)
	for k, vs := range headers {
		for _, v := range vs {
			if existingValue, exists := hMap[k]; exists {
				hMap[k] = existingValue + "," + v
			} else {
				hMap[k] = v
			}
		}
	}
	return hMap
}

func bodyFromResponse(gatewayResponse events.APIGatewayProxyResponse) (io.ReadCloser, error) {
	header := gatewayResponse.Headers["Content-Type"]
	i := strings.Index(header, ";")
	if i == -1 {
		i = len(header)
	}
	contentType := strings.TrimSpace(strings.ToLower(header[:i]))
	if contentType == contentTypeJSON {
		return ioutil.NopCloser(strings.NewReader(gatewayResponse.Body)), nil
	}
	if contentType == contentTypeProto {
		b, err := base64.StdEncoding.DecodeString(gatewayResponse.Body)
		if err != nil {
			return nil, err
		}
		return ioutil.NopCloser(bytes.NewBuffer(b)), nil
	}
	return nil, fmt.Errorf("labmbda_transport: unknown content-type: %q", contentType)
}

// Do function to implement Twirp HTTP
func (l *Client) Do(req *http.Request) (*http.Response, error) {
	invokeInput, err := l.invokeInputFromHTTPRequest(req)

	if err != nil {
		return nil, err
	}

	ctx := req.Context()
	var result *lambda.InvokeOutput
	if ctx == nil {
		result, err = l.lambdaAPI.Invoke(invokeInput)
	} else {
		result, err = l.lambdaAPI.InvokeWithContext(ctx, invokeInput)
	}

	// TODO: Need to return a response on this if there was one, like in a 403 case
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			return convertAwsErrorCode(aerr)
		}
		return nil, err
	}

	// Check if there was an error executing the funtion
	if result.FunctionError != nil {
		return convertLambdaFunctionError(result)
	}

	var gatewayResponse events.APIGatewayProxyResponse
	err = json.Unmarshal(result.Payload, &gatewayResponse)
	if err != nil {
		return nil, err
	}

	body, err := bodyFromResponse(gatewayResponse)
	if err != nil {
		return nil, err
	}

	return &http.Response{
		StatusCode: gatewayResponse.StatusCode,
		Body:       body,
	}, nil
}

func convertAwsErrorCode(aerr awserr.Error) (*http.Response, error) {
	if statusCode, ok := awserrCodeToHttpStatusCode[aerr.Code()]; ok {
		return &http.Response{
			StatusCode: statusCode,
			Body:       ioutil.NopCloser(bytes.NewBuffer([]byte(aerr.Message()))),
		}, nil
	} else {
		// Bubble up unhandled errors. They will become internal
		// client errors.
		return nil, aerr
	}
}

// convertLambdaFunctionError converts a lambda function error response to an
// http.Response following the Twirp error spec
func convertLambdaFunctionError(output *lambda.InvokeOutput) (*http.Response, error) {
	if output.FunctionError == nil {
		return nil, errors.New("Result was not a Function error")
	}

	var errorResponse lambdaErrorJSON
	err := json.Unmarshal(output.Payload, &errorResponse)
	if err != nil {
		return nil, err
	}

	twirpErr := twirp.InternalError("There was an error executing the Lambda function.")
	// AWS Lambda Functions errors can have non-string error messages
	twirpErr = twirpErr.WithMeta("errorMessage", errorResponse.Message)
	twirpErr = twirpErr.WithMeta("errorType", errorResponse.Type)
	twirpErr = twirpErr.WithMeta("stackTrace", strings.Join(errorResponse.StackTrace, "\n"))
	errorBody := marshalErrorToJSON(twirpErr)

	return &http.Response{
		StatusCode: 500,
		Body:       ioutil.NopCloser(bytes.NewBuffer(errorBody)),
	}, nil
}

// marshalErrorToJSON returns JSON from a twirp.Error, that can be used as HTTP error response body.
// If serialization fails, it will use a descriptive Internal error instead.
func marshalErrorToJSON(twerr twirp.Error) []byte {
	// make sure that msg is not too large
	msg := twerr.Msg()
	if len(msg) > 1e6 {
		msg = msg[:1e6]
	}

	tj := twerrJSON{
		Code: string(twerr.Code()),
		Msg:  msg,
		Meta: twerr.MetaMap(),
	}

	buf, err := json.Marshal(&tj)
	if err != nil {
		buf = []byte("{\"type\": \"" + twirp.Internal + "\", \"msg\": \"There was an error but it could not be serialized into JSON\"}") // fallback
	}

	return buf
}
