package handler_utils

import (
	"bytes"
	"context"
	"github.com/aws/aws-sdk-go/aws/credentials"
	v4 "github.com/aws/aws-sdk-go/aws/signer/v4"
	"github.com/go-resty/resty/v2"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-lambda-go/lambdacontext"
)

func RequestIDFromLambdaContext(ctx context.Context) *string {
	// Get request ID from context
	if ctx != nil {
		lambdaContext, _ := lambdacontext.FromContext(ctx)
		if lambdaContext != nil {
			return &lambdaContext.AwsRequestID
		}
	}
	return nil
}

func RequestIDFromHeaders(headers map[string]string, requestIDHeaderKey string) *string {
	// Get request ID from parent micro service
	if requestID, ok := headers[requestIDHeaderKey]; ok {
		return &requestID
	}
	return nil
}

func AddRequestIDToHeaders(requestID *string, headers map[string]string, requestIDHeaderKey string) {
	if requestID != nil && headers != nil {
		headers[requestIDHeaderKey] = *requestID
	}
}

func SignAWSRestyClient(client *resty.Client, accessKeyID, secretAccessKey string, bodyBytes []byte) error {
	var bodySeeker io.ReadSeeker = bytes.NewReader(bodyBytes)

	// Add a hook to sign the request before sending
	client.SetPreRequestHook(func(_ *resty.Client, r *http.Request) error {
		err := SignAWSHttpRequest(r, accessKeyID, secretAccessKey, bodySeeker)
		if err != nil {
			return err
		}
		return nil
	})

	return nil
}

// SignAWSRequest wraps and executes http.NewRequest and adds a sig version 4 signature for AWS API Gateway
func SignAWSHttpRequest(request *http.Request, accessKeyID, secretAccessKey string, bodySeeker io.ReadSeeker) error {
	// Use AWS SDK to sign the request for API gateway,  i.e. execute-api, and the current region
	_, err := v4.NewSigner(credentials.NewStaticCredentials(accessKeyID, secretAccessKey, "")).
		Sign(request, bodySeeker, "execute-api", os.Getenv("AWS_REGION"), time.Now())
	if err != nil {
		return err
	}

	return nil
}

func FindHeaderValue(headers map[string]string, key string) string {
	for k, v := range headers {
		if strings.ToLower(k) == strings.ToLower(key) {
			return v
		}
	}
	return ""
}