package server

import (
	"testing"

	"buf.build/go/protovalidate"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/featureflag"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/jwttool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/matcher"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_otel"
	"go.uber.org/mock/gomock"
)

const anyUserID = 84

func TestWebSocketToken_SuccessfulGenerateVerify(t *testing.T) {
	// GIVEN
	wt := setup(t, "test")
	testAgentKey := api.AgentKey{ID: 42, Type: api.AgentTypeKubernetes}
	testImpConfig := &kubernetes_api.ImpersonationConfig{
		Username: "any-username",
		Groups:   []string{"any-group"},
		Uid:      "any-uid",
	}

	// WHEN
	token, err := wt.generate("any-url", testAgentKey, anyUserID, testImpConfig, featureflag.Set{})
	require.NoError(t, err)
	agentKey, impConfig, ffs, err := wt.verify(token, "any-url")
	require.NoError(t, err)

	// THEN
	assert.Equal(t, testAgentKey, agentKey)
	assert.EqualExportedValues(t, testImpConfig, impConfig)
	assert.Equal(t, featureflag.Set{}, ffs)
}

func TestWebSocketToken_WithFeatureFlags(t *testing.T) {
	// GIVEN
	testAgentKey := api.AgentKey{ID: 42, Type: api.AgentTypeKubernetes}
	testImpConfig := &kubernetes_api.ImpersonationConfig{
		Username: "any-username",
		Groups:   []string{"any-group"},
		Uid:      "any-uid",
	}

	ctrl := gomock.NewController(t)
	ffCheckCounterMock := mock_otel.NewMockInt64Counter(ctrl)
	ffFoo := featureflag.NewFeatureFlag("foo", false)
	ffBar := featureflag.NewFeatureFlag("bar", true)

	ffCheckCounterMock.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"name": ffFoo.Name, "enabled": !ffFoo.Default}))
	ffCheckCounterMock.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"name": ffBar.Name, "enabled": !ffBar.Default}))

	v, err := protovalidate.New()
	require.NoError(t, err)
	wt := webSocketToken{
		validator:        v,
		jwtSigningMethod: jwttool.SigningMethodHS3_512,
		jwtSecret:        []byte("test"),
		ffCheckCounter:   ffCheckCounterMock,
	}

	// WHEN
	token, err := wt.generate("any-url", testAgentKey, anyUserID, testImpConfig, featureflag.NewSet(map[string]bool{ffFoo.Name: !ffFoo.Default, ffBar.Name: !ffBar.Default}, nil))
	require.NoError(t, err)
	agentKey, impConfig, ffs, err := wt.verify(token, "any-url")
	require.NoError(t, err)

	// THEN
	assert.Equal(t, testAgentKey, agentKey)
	assert.EqualExportedValues(t, testImpConfig, impConfig)
	assert.True(t, ffs.IsEnabled(ffFoo))
	assert.True(t, ffs.IsDisabled(ffBar))
}

func TestWebSocketToken_InvalidEndpoint(t *testing.T) {
	// GIVEN
	wt := setup(t, "test")

	testImpConfig := &kubernetes_api.ImpersonationConfig{
		Username: "any-username",
		Groups:   []string{"any-group"},
		Uid:      "any-uid",
	}

	testAgentKey := api.AgentKey{ID: 42, Type: api.AgentTypeKubernetes}
	token, err := wt.generate("some-url", testAgentKey, anyUserID, testImpConfig, featureflag.Set{})
	require.NoError(t, err)

	// WHEN
	_, _, _, err = wt.verify(token, "another-url")

	// THEN
	assert.EqualError(t, err, "token has invalid claims: token has invalid endpoint")
}

func TestWebSocketToken_InvalidSignature(t *testing.T) {
	// GIVEN
	wtGenerate := setup(t, "some-secret")
	wtVerify := setup(t, "another-secret")

	// WHEN
	testAgentKey := api.AgentKey{ID: 42, Type: api.AgentTypeKubernetes}
	token, err := wtGenerate.generate("any-url", testAgentKey, anyUserID, nil, featureflag.Set{})
	require.NoError(t, err)
	_, _, _, err = wtVerify.verify(token, "any-url")

	// THEN
	assert.EqualError(t, err, "token signature is invalid: signature is invalid")
}

func TestWebSocketToken_InvalidEmptyToken(t *testing.T) {
	// GIVEN
	wt := setup(t, "any-secret")

	// WHEN
	_, _, _, err := wt.verify("", "any-url")

	// THEN
	assert.EqualError(t, err, "token is malformed: token contains an invalid number of segments")
}

func setup(t *testing.T, secret string) webSocketToken {
	v, err := protovalidate.New()
	require.NoError(t, err)
	wt := webSocketToken{
		validator:        v,
		jwtSigningMethod: jwttool.SigningMethodHS3_512,
		jwtSecret:        []byte(secret),
	}
	return wt
}
