package correlation

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestInjectCorrelationID(t *testing.T) {
	tests := []struct {
		name                      string
		opts                      []InboundHandlerOption
		header                    http.Header
		shouldMatch               string
		shouldNotMatch            string
		shouldSetResponseHeader   bool
		expectedResponseHeader    string
		downstreamResponseHeaders http.Header
	}{
		{
			name: "without_propagation",
		},
		{
			name: "without_propagation_ignore_incoming_header",
			header: map[string][]string{
				propagationHeader: []string{"123"},
			},
			shouldNotMatch: "123",
		},
		{
			name: "with_propagation_no_incoming_header",
			opts: []InboundHandlerOption{WithPropagation()},
		},
		{
			name: "with_propagation_incoming_header",
			opts: []InboundHandlerOption{WithPropagation()},
			header: map[string][]string{
				propagationHeader: []string{"123"},
			},
			shouldMatch: "123",
		},
		{
			name: "with_propagation_double_incoming_header",
			opts: []InboundHandlerOption{WithPropagation()},
			header: map[string][]string{
				propagationHeader: []string{"123", "456"},
			},
			shouldMatch: "123",
		},
		{
			name:                    "with_set_response_header",
			opts:                    []InboundHandlerOption{WithSetResponseHeader()},
			shouldSetResponseHeader: true,
		},
		{
			name:                    "with_set_response_header_and_with_propagation_incoming_header",
			opts:                    []InboundHandlerOption{WithSetResponseHeader(), WithPropagation()},
			shouldSetResponseHeader: true,
		},
		{
			name:                    "with_set_response_header_and_with_propagation_incoming_header_set",
			opts:                    []InboundHandlerOption{WithSetResponseHeader(), WithPropagation()},
			shouldSetResponseHeader: true,
			header: map[string][]string{
				propagationHeader: []string{"123"},
			},
			expectedResponseHeader: "123",
		},
		{
			name: "mismatching_correlation_ids_without_propagation",
			opts: []InboundHandlerOption{},
			downstreamResponseHeaders: map[string][]string{
				propagationHeader: []string{"CLIENT_CORRELATION_ID"},
			},
			shouldSetResponseHeader: true,
			expectedResponseHeader:  "CLIENT_CORRELATION_ID",
		},

		{
			name: "mismatching_correlation_ids_with_propagation",
			opts: []InboundHandlerOption{WithSetResponseHeader()},
			downstreamResponseHeaders: map[string][]string{
				propagationHeader: []string{"CLIENT_CORRELATION_ID"},
			},
			shouldSetResponseHeader: true,
			expectedResponseHeader:  "CLIENT_CORRELATION_ID",
		},
		{
			name:                    "with_set_response_header_and_with_propagation_incoming_header_set",
			opts:                    []InboundHandlerOption{WithSetResponseHeader(), WithPropagation()},
			shouldSetResponseHeader: true,
			header: map[string][]string{
				propagationHeader: []string{"123"},
			},
			downstreamResponseHeaders: map[string][]string{
				propagationHeader: []string{"CLIENT_CORRELATION_ID"},
			},
			expectedResponseHeader: "CLIENT_CORRELATION_ID",
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			invoked := false

			h := InjectCorrelationID(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				invoked = true

				ctx := r.Context()
				correlationID := ExtractFromContext(ctx)
				require.NotNil(t, correlationID, "CorrelationID is missing")
				require.NotEmpty(t, correlationID, "CorrelationID is missing")
				if test.shouldMatch != "" {
					require.Equal(t, test.shouldMatch, correlationID, "CorrelationID should match")
				}

				if test.shouldNotMatch != "" {
					require.NotEqual(t, test.shouldNotMatch, correlationID, "CorrelationID should not match")
				}

				for k, vs := range test.downstreamResponseHeaders {
					for _, v := range vs {
						w.Header().Add(k, v)
					}
				}

			}), test.opts...)

			r := httptest.NewRequest("GET", "http://example.com", nil)
			for k, v := range test.header {
				r.Header[http.CanonicalHeaderKey(k)] = v
			}

			w := httptest.NewRecorder()
			h.ServeHTTP(w, r)

			require.True(t, invoked, "handler not executed")

			v, ok := w.HeaderMap[http.CanonicalHeaderKey(propagationHeader)]
			require.Equal(t, test.shouldSetResponseHeader, ok, "response header existence mismatch")
			if test.shouldSetResponseHeader {
				require.Equal(t, 1, len(v), "expected exactly one correlation header")
				require.NotEqual(t, "", v[0], "expected non-empty correlation header")
			}

			if test.expectedResponseHeader != "" {
				responseHeader := w.Header().Get(propagationHeader)
				require.Equal(t, test.expectedResponseHeader, responseHeader, "response header should match")
			}
		})
	}
}
