package xray

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/aws/request"
)

var xRayBuildHandler = request.NamedHandler{
	Name: "XRayBuildHandler",
	Fn: func(r *request.Request) {
		ctx := r.HTTPRequest.Context()

		ctx, subseg := newSegment(ctx, r.ClientInfo.ServiceName)
		r.HTTPRequest = r.HTTPRequest.WithContext(ctx)

		r.HTTPRequest.Header.Set("x-amzn-trace-id", fmt.Sprintf("Root=%s; Parent=%s; Sampled=%d", subseg.root().TraceID, subseg.ID, btoi(subseg.root().sampled)))
	},
}

var xRayCompleteHandler = request.NamedHandler{
	Name: "XRayCompleteHandler",
	Fn: func(r *request.Request) {
		subseg := getSegment(r.HTTPRequest.Context())

		subseg.Lock()
		subseg.Namespace = "AWS"

		subseg.aws().Region = r.ClientInfo.SigningRegion
		subseg.aws().Operation = r.Operation.Name
		subseg.aws().RequestID = r.RequestID
		subseg.aws().Retries = r.RetryCount

		subseg.http().response().Status = r.HTTPResponse.StatusCode
		subseg.http().response().ContentLength = int(r.HTTPResponse.ContentLength)
		subseg.Unlock()

		subseg.close(r.Error)
	},
}

// AWS adds XRay tracing to an AWS client
func AWS(c *client.Client) {
	c.Handlers.Build.PushBackNamed(xRayBuildHandler)
	c.Handlers.Complete.PushBackNamed(xRayCompleteHandler)
}
