package tracing

import (
	"context"
	"os"
	"reflect"
	"testing"

	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/labkit/correlation"
)

func TestNewEnvInjector(t *testing.T) {
	tests := []struct {
		name          string
		tracingEnv    string
		correlationID string
		env           []string
		wantEnv       []string
	}{
		{
			name: "trivial",
		},
		{
			name:       "gitlab_tracing_config",
			tracingEnv: "opentracing://null",
			wantEnv: []string{
				"GITLAB_TRACING=opentracing://null",
			},
		},
		{
			name:       "merge",
			tracingEnv: "opentracing://null",
			env: []string{
				"A=1",
			},
			wantEnv: []string{
				"A=1",
				"GITLAB_TRACING=opentracing://null",
			},
		},
		{
			name:          "correlation",
			correlationID: "abc123",
			env:           []string{},
			wantEnv: []string{
				"CORRELATION_ID=abc123",
			},
		},
		{
			name:          "correlation_merge",
			tracingEnv:    "opentracing://null",
			correlationID: "abc123",
			env:           []string{"A=1"},
			wantEnv: []string{
				"A=1",
				"CORRELATION_ID=abc123",
				"GITLAB_TRACING=opentracing://null",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.tracingEnv != "" {
				resetEnvironment := modifyEnvironment(tracingEnvKey, tt.tracingEnv)
				defer resetEnvironment()
			} else {
				resetEnvironment := clearEnvironment(tracingEnvKey)
				defer resetEnvironment()
			}
			envInjector := NewEnvInjector()

			ctx := context.Background()
			if tt.correlationID != "" {
				ctx = correlation.ContextWithCorrelation(ctx, tt.correlationID)
			}

			got := envInjector(ctx, tt.env)
			require.Equal(t, tt.wantEnv, got)
		})
	}
}

// modifyEnvironment will change an environment variable and return a func suitable
// for `defer` to change the value back.
func modifyEnvironment(key string, value string) func() {
	oldValue, hasOldValue := os.LookupEnv(key)
	os.Setenv(key, value)
	return func() {
		if hasOldValue {
			os.Setenv(key, oldValue)
		} else {
			os.Unsetenv(key)
		}
	}
}

// clearEnvironment will clear an environment variable and return a func suitable
// for `defer` to change the value back.
func clearEnvironment(key string) func() {
	oldValue, hasOldValue := os.LookupEnv(key)

	if !hasOldValue {
		return func() {
			// Nothing to do
		}
	}

	os.Unsetenv(key)
	return func() {
		if hasOldValue {
			os.Setenv(key, oldValue)
		}
	}
}

func Test_appendMapToEnv(t *testing.T) {
	tests := []struct {
		name   string
		env    []string
		envMap map[string]string
		want   []string
	}{
		{
			name:   "empty_case",
			env:    nil,
			envMap: nil,
			want:   nil,
		},
		{
			name:   "no_new_values",
			env:    []string{"A=1"},
			envMap: nil,
			want:   []string{"A=1"},
		},
		{
			name:   "no_original_values",
			env:    nil,
			envMap: map[string]string{"A": "1"},
			want:   []string{"A=1"},
		},
		{
			name:   "merge_values",
			env:    []string{"A=1"},
			envMap: map[string]string{"B": "2"},
			want:   []string{"A=1", "B=2"},
		},
		{
			name:   "merge_conflicts",
			env:    []string{"A=1", "C=3"},
			envMap: map[string]string{"B": "2", "C": "4"},
			want:   []string{"A=1", "C=3", "B=2", "C=4"},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := appendMapToEnv(tt.env, tt.envMap); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("appendMapToEnv() = %v, want %v", got, tt.want)
			}
		})
	}
}
