package xray

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

// Handler wraps the provided http handler with xray.Capture
// using the request's context, parsing the incoming headers,
// adding response headers if needed, and sets HTTP sepecific trace fields.
func Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// As requests can contain spoofed host headers,
		// we allow users to overwrite the segment name in the config
		name := config.Name()
		if name == "" {
			name = r.Host
		}

		// We also attempt to detect IP addresses and use a default name instead
		host, _, err := net.SplitHostPort(name)
		if err != nil {
			host = name
		}
		if net.ParseIP(host) != nil {
			name = config.DefaultName()
		}

		Capture(r.Context(), name, func(ctx context.Context) error {
			r = r.WithContext(ctx)
			seg := getSegment(ctx)

			seg.Lock()
			seg.http().request().Method = r.Method
			seg.http().request().URL = r.URL.String()
			seg.http().request().ClientIP, seg.http().request().XForwardedFor = clientIP(r)
			seg.http().request().UserAgent = r.UserAgent()

			trace := parseHeaders(r.Header)
			if trace["Root"] != "" {
				seg.TraceID = trace["Root"]
				seg.requestWasTraced = true
			}
			if trace["Parent"] != "" {
				seg.ParentID = trace["Parent"]
			}
			switch trace["Sampled"] {
			case "0":
				seg.sampled = false
			case "1":
				seg.sampled = true
			case "?":
				w.Header().Set("x-amzn-trace-id", fmt.Sprintf("Root=%s; Sampled=%d", seg.TraceID, btoi(seg.sampled)))
			}
			seg.Unlock()

			resp := &responseCapturer{w, 200, 0}
			h.ServeHTTP(resp, r)

			seg.Lock()
			seg.http().response().Status = resp.status
			seg.http().response().ContentLength = resp.length

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

			return nil
		})
	})
}

func parseHeaders(h http.Header) map[string]string {
	m := map[string]string{}
	s := h.Get("x-amzn-trace-id")
	for _, c := range strings.Split(s, ";") {
		p := strings.SplitN(c, "=", 2)
		k := strings.TrimSpace(p[0])
		v := ""
		if len(p) > 1 {
			v = strings.TrimSpace(p[1])
		}
		m[k] = v
	}
	return m
}

func clientIP(r *http.Request) (string, bool) {
	forwardedFor := r.Header.Get("X-Forwarded-For")
	if forwardedFor != "" {
		return strings.TrimSpace(strings.Split(forwardedFor, ",")[0]), true
	}

	return r.RemoteAddr, false
}

type responseCapturer struct {
	http.ResponseWriter
	status int
	length int
}

func (w *responseCapturer) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *responseCapturer) Write(data []byte) (int, error) {
	w.length += len(data)
	return w.ResponseWriter.Write(data)
}
