package lookup

import (
	"context"
	"encoding/base64"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/store/identity"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

type countingSessions struct {
	calls int
	err   error
}

func (s *countingSessions) AuthenticateSession(context.Context, string) (identity.AuthenticatedUser, error) {
	s.calls++
	if s.err != nil {
		return identity.AuthenticatedUser{}, s.err
	}
	return identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}}, nil
}

func TestAuthenticatedResolverBoundsUnknownBearerToOneSessionLookup(t *testing.T) {
	sessions := &countingSessions{err: context.Canceled}
	workspaces := &countingWorkspaces{}
	resolver := NewAuthenticatedResolver(panicLegacy{}, sessions, workspaces)
	workspaceID, resolved, err := resolver.ResolveWorkspaceRouteIdentity(
		context.Background(), "GET", "/api/v3/repos/acme/repo/branches", "", "",
		"Bearer sssssssssssssssssssssssssssssssssssssssssss", "session_candidate", "wsp_a",
	)
	require.Error(t, err)
	require.False(t, resolved)
	require.Empty(t, workspaceID)
	require.Equal(t, 1, sessions.calls)
	require.Zero(t, workspaces.calls)
}

type countingWorkspaces struct{ calls int }

func (w *countingWorkspaces) GetWorkspaceForUser(context.Context, string, string) (workspace.Summary, error) {
	w.calls++
	return workspace.Summary{WorkspaceID: "wsp_a"}, nil
}

type panicLegacy struct{}

func (panicLegacy) ResolveWorkspaceRouteIdentity(context.Context, string, string, string, string, string, string, string) (string, bool, error) {
	panic("legacy resolver called for AGS session route")
}

func TestAuthenticatedResolverUsesOneBoundedSessionAndMembershipLookup(t *testing.T) {
	sessions := &countingSessions{}
	workspaces := &countingWorkspaces{}
	resolver := NewAuthenticatedResolver(panicLegacy{}, sessions, workspaces)
	workspaceID, resolved, err := resolver.ResolveWorkspaceRouteIdentity(context.Background(), "GET", "/api/v3/repos/acme/repo/branches", "", "", "Bearer sssssssssssssssssssssssssssssssssssssssssss", "session_candidate", "wsp_a")
	require.NoError(t, err)
	require.True(t, resolved)
	require.Equal(t, "wsp_a", workspaceID)
	require.Equal(t, 1, sessions.calls)
	require.Equal(t, 1, workspaces.calls)
}

func TestAuthenticatedResolverUsesOuterWorkspaceCandidate(t *testing.T) {
	sessions := &countingSessions{}
	workspaces := &countingWorkspaces{}
	resolver := NewAuthenticatedResolver(panicLegacy{}, sessions, workspaces)
	path := "/api/v1/workspaces/wsp_a/ags/api/v3/repos/acme/repo"

	workspaceID, resolved, err := resolver.ResolveWorkspaceRouteIdentity(
		context.Background(), "GET", path, "", "", "Bearer sssssssssssssssssssssssssssssssssssssssssss", "session_candidate", "wsp_a",
	)
	require.NoError(t, err)
	require.True(t, resolved)
	require.Equal(t, "wsp_a", workspaceID)
	require.Equal(t, 1, sessions.calls)
	require.Equal(t, 1, workspaces.calls)

	workspaceID, resolved, err = resolver.ResolveWorkspaceRouteIdentity(
		context.Background(), "GET", path, "", "", "Bearer sssssssssssssssssssssssssssssssssssssssssss", "session_candidate", "wsp_b",
	)
	require.NoError(t, err)
	require.False(t, resolved)
	require.Empty(t, workspaceID)
	require.Equal(t, 1, sessions.calls)
	require.Equal(t, 1, workspaces.calls)
}

func TestAuthenticatedResolverRejectsNonSessionOrMismatchedAuthorityBeforeDatabase(t *testing.T) {
	validToken := base64.RawURLEncoding.EncodeToString(make([]byte, 32))
	nonMinimalPaddingBits := validToken[:len(validToken)-1] + "B"
	for _, test := range []struct {
		name          string
		authorization string
		authority     string
	}{
		{name: "repository bearer", authorization: "Bearer ags_rat_secret", authority: "repository_credential"},
		{name: "repository token", authorization: "Token ags_rat_secret", authority: "repository_credential"},
		{name: "repository basic", authorization: "Basic Z2l0OmFnc19yYXRfc2VjcmV0", authority: "repository_credential"},
		{name: "malformed", authorization: "Bearer  session", authority: "malformed"},
		{name: "tampered declaration", authorization: "Bearer session", authority: "repository_credential"},
		{name: "redundant padding bits", authorization: "Bearer " + nonMinimalPaddingBits, authority: "session_candidate"},
	} {
		t.Run(test.name, func(t *testing.T) {
			sessions := &countingSessions{}
			workspaces := &countingWorkspaces{}
			resolver := NewAuthenticatedResolver(panicLegacy{}, sessions, workspaces)
			workspaceID, resolved, err := resolver.ResolveWorkspaceRouteIdentity(
				context.Background(), "GET", "/api/v3/repos/acme/repo/branches", "", "",
				test.authorization, test.authority, "wsp_a",
			)
			require.Error(t, err)
			require.False(t, resolved)
			require.Empty(t, workspaceID)
			require.Zero(t, sessions.calls)
			require.Zero(t, workspaces.calls)
		})
	}
}