package correlation

import (
	"net"
	"net/http"

	"github.com/sebest/xff"
)

// InjectCorrelationID is an HTTP middleware to generate an Correlation-ID for the incoming request,
// or extract the existing Correlation-ID from the incoming request. By default, any upstream Correlation-ID,
// passed in via the `X-Request-ID` header will be ignored. To enable this behaviour, the `WithPropagation`
// option should be passed into the options.
// Whether the Correlation-ID is generated or propagated, once inside this handler the request context
// will have a Correlation-ID associated with it.
func InjectCorrelationID(h http.Handler, opts ...InboundHandlerOption) http.Handler {
	config := applyInboundHandlerOptions(opts)

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		parent := r.Context()

		correlationID := ""
		clientName := ""
		ipAddress := getRemoteIP(r, config)
		if config.shouldPropagate(ipAddress) {
			correlationID = r.Header.Get(propagationHeader)
			clientName = r.Header.Get(clientNameHeader)
		}

		if correlationID == "" {
			correlationID = SafeRandomID()
		}

		ctx := ContextWithCorrelation(parent, correlationID)

		if clientName != "" {
			ctx = ContextWithClientName(ctx, clientName)
		}
		if config.sendResponseHeader {
			// Set the response header.
			w.Header().Set(propagationHeader, correlationID)
		}
		h.ServeHTTP(w, r.WithContext(ctx))
	})
}

func getRemoteIP(r *http.Request, c inboundHandlerConfig) string {
	remoteAddr := r.RemoteAddr

	// Unix domain sockets have a remote addr of @. This will make the
	// xff package lookup the X-Forwarded-For address if available.
	if remoteAddr == "@" {
		r.RemoteAddr = "127.0.0.1:0"
	}

	if c.xffAllowed != nil {
		remoteAddr = xff.GetRemoteAddrIfAllowed(r, c.xffAllowed)
	}

	host, _, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		return r.RemoteAddr
	}

	return host
}
