package server

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/extension"
	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
)

func TestExtensionProviderIndexProjectsOnlyOptedInConnectionCapability(t *testing.T) {
	index := &extensionProviderIndex{byAppID: map[string]extension.Extension{
		"app_generic": {
			Definition: registry.Definition{Domain: "provider-domain"},
			Distribution: extension.Distribution{
				AppID:                     "app_generic",
				GenericProviderConnection: true,
			},
		},
		"app_bespoke": {
			Definition: registry.Definition{Domain: "bespoke-domain"},
			Distribution: extension.Distribution{
				AppID: "app_bespoke",
			},
		},
		"app_malformed": {
			Distribution: extension.Distribution{
				AppID:                     "app_malformed",
				GenericProviderConnection: true,
			},
		},
	}}

	capability, ok := index.ProviderConnectionCapabilityForAppID(" app_generic ")
	require.True(t, ok)
	require.Equal(t, "provider-domain", capability.Domain,
		"domain must come from the trusted registered Definition")

	for _, appID := range []string{"app_bespoke", "app_malformed", "app_unknown", ""} {
		_, ok := index.ProviderConnectionCapabilityForAppID(appID)
		require.Falsef(t, ok, "%s must fail closed", appID)
	}
}

func TestExtensionProviderIndexResolvesOnlyValidGenericConnectionOptIn(
	t *testing.T,
) {
	t.Parallel()

	validator := func(context.Context, runner.Credentials) error {
		return nil
	}
	auth := runner.AuthSpec{
		Kind: runner.AuthBearer,
		Fields: []runner.CredentialField{
			runner.APIKeyField("API key", "Provider API key."),
		},
	}
	index := &extensionProviderIndex{byAppID: map[string]extension.Extension{
		"app_generic": {
			Definition: registry.Definition{
				Domain: "generic",
				Name:   "Generic",
			},
			Distribution: extension.Distribution{
				AppID:                     "app_generic",
				AppSlug:                   "generic",
				GenericProviderConnection: true,
			},
			Implementation: runner.Implementation{
				Auth:                auth,
				ValidateCredentials: validator,
			},
		},
		"app_bespoke": {
			Definition: registry.Definition{
				Domain: "bespoke",
				Name:   "Bespoke",
			},
			Distribution: extension.Distribution{
				AppID:   "app_bespoke",
				AppSlug: "bespoke",
			},
			Implementation: runner.Implementation{
				Auth:                auth,
				ValidateCredentials: validator,
			},
		},
		"app_ambiguous": {
			Definition: registry.Definition{
				Domain: "ambiguous",
				Name:   "Ambiguous",
			},
			Distribution: extension.Distribution{
				AppID:                     "app_ambiguous",
				AppSlug:                   "ambiguous",
				GenericProviderConnection: true,
			},
			Implementation: runner.Implementation{
				Auth:                auth,
				ValidateCredentials: validator,
				PrepareConnection: func(
					context.Context,
					runner.Credentials,
				) (runner.ConnectionConfig, error) {
					return nil, nil
				},
			},
		},
	}}

	resolved, err := index.resolveRegisteredProvider("app_generic")
	require.NoError(t, err)
	require.Equal(t, "app_generic", resolved.Target.AppID)
	require.NotNil(t, resolved.PrepareConnection)

	for _, appID := range []string{
		"app_bespoke",
		"app_ambiguous",
		"app_unknown",
	} {
		_, err := index.resolveRegisteredProvider(appID)
		require.ErrorIs(t, err, api.ErrNotFound)
	}
}