package xray

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
)

// Client creates a shallow copy of the provided http client,
// defaulting to http.DefaultClient, with roundtripper wrapped
// with xray.RoundTripper.
func Client(c *http.Client) *http.Client {
	if c == nil {
		c = http.DefaultClient
	}
	transport := c.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}
	return &http.Client{
		Transport:     RoundTripper(transport),
		CheckRedirect: c.CheckRedirect,
		Jar:           c.Jar,
		Timeout:       c.Timeout,
	}
}

// RoundTripper wraps the provided http roundtripper with xray.Capture,
// sets HTTP-specific xray fields, and adds the trace header to the outbound request.
func RoundTripper(rt http.RoundTripper) http.RoundTripper {
	return &roundtripper{rt}
}

type roundtripper struct {
	Base http.RoundTripper
}

func (rt *roundtripper) RoundTrip(r *http.Request) (*http.Response, error) {
	var resp *http.Response
	err := Capture(r.Context(), r.Host, func(ctx context.Context) error {
		var err error
		r = r.WithContext(ctx)
		seg := getSegment(ctx)

		seg.Lock()
		seg.Namespace = "remote"
		seg.http().request().Method = r.Method
		seg.http().request().URL = r.URL.String()

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

		resp, err = rt.Base.RoundTrip(r)

		if resp != nil {
			seg.Lock()
			seg.http().response().Status = resp.StatusCode
			seg.http().response().ContentLength, _ = strconv.Atoi(resp.Header.Get("Content-Length"))

			if resp.StatusCode >= 400 && resp.StatusCode < 500 {
				seg.Error = true
			}
			if resp.StatusCode == 429 {
				seg.Throttle = true
			}
			if resp.StatusCode >= 500 && resp.StatusCode < 600 {
				seg.Fault = true
			}
			seg.Unlock()
		}

		return err
	})
	return resp, err
}

func btoi(b bool) int {
	if b {
		return 1
	}
	return 0
}
