package searchengine_test

import (
	"context"
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/extension/providers/brave"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/searchengine"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/searchengine/runtime"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/searxng"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/serper"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/tavily"
	"github.com/tidbcloud/AgentTeam/internal/extension/registry"
	"github.com/tidbcloud/AgentTeam/internal/service/extensionflow"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

func TestProviderDefinitionsExposeConcreteSearchAction(t *testing.T) {
	t.Parallel()

	definitions := []registry.Definition{
		searxng.Definition(),
		brave.Definition(),
		tavily.Definition(),
		serper.Definition(),
	}
	wantDomains := []string{"searxng", "brave", "tavily", "serper"}
	for index, definition := range definitions {
		require.Equal(t, wantDomains[index], definition.Domain)
		require.Len(t, definition.Actions, 1)
		action := definition.Actions[0]
		require.Equal(t, searchengine.ActionSearch, action.Name)
		require.Equal(t, "read", action.Safety.Class)
		require.True(t, action.Safety.Idempotent)
		require.False(t, action.Safety.Destructive)
		properties := action.OutputSchema["properties"].(map[string]any)
		require.Contains(t, properties, "results")
		require.NotContains(t, properties, "provider")
		inputProperties := action.InputSchema["properties"].(map[string]any)
		require.NotContains(t, inputProperties, "provider")
	}
}

func TestProviderRegistrationsAreIndependentBuiltins(t *testing.T) {
	t.Parallel()

	reg := registry.New()
	invoker := &recordingInvoker{}
	require.NoError(t, searxng.Register(reg, invoker))
	require.NoError(t, brave.Register(reg, invoker))
	require.NoError(t, tavily.Register(reg, invoker))
	require.NoError(t, serper.Register(reg, invoker))

	registered := reg.List()
	require.Len(t, registered, 4)
	wantDigests := map[string]string{
		runtime.ProviderSearXNG: model.SearXNGBuiltinExtensionDefinitionDigest,
		runtime.ProviderBrave:   model.BraveBuiltinExtensionDefinitionDigest,
		runtime.ProviderTavily:  model.TavilyBuiltinExtensionDefinitionDigest,
		runtime.ProviderSerper:  model.SerperBuiltinExtensionDefinitionDigest,
	}
	for _, extension := range registered {
		require.Equal(t, registry.Builtin, extension.Kind)
		require.Equal(t, wantDigests[extension.Definition.Domain], extension.Definition.Digest)
	}
}

func TestProviderRegistrationsRejectMissingSearchService(t *testing.T) {
	t.Parallel()

	for name, register := range map[string]func(*registry.Registry, searchengine.Invoker) error{
		"searxng": searxng.Register,
		"brave":   brave.Register,
		"tavily":  tavily.Register,
		"serper":  serper.Register,
	} {
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			reg := registry.New()
			require.Error(t, register(reg, nil))
			_, ok := reg.LookupByIdentity(name, 1)
			require.False(t, ok)
		})
	}
}

func TestExecutorBindsProviderIdentity(t *testing.T) {
	t.Parallel()

	invoker := &recordingInvoker{result: runtime.SearchResult{
		Results: []runtime.Result{{
			URL:     "https://example.com",
			Title:   "Example",
			Snippet: "snippet",
		}},
	}}
	executor := searchengine.NewExecutor(runtime.ProviderSearXNG, invoker)
	result, err := executor.Invoke(context.Background(), registry.BuiltinInvocation{
		WorkspaceID:           "wsp_1",
		AgentID:               "agt_1",
		AgentMemberID:         "mem_1",
		ServiceInstallationID: "svcinst_1",
		ExtensionName:         "searxng",
		Action:                searchengine.ActionSearch,
		Input:                 json.RawMessage(`{"query":"AgentTeam","limit":2}`),
	})
	require.NoError(t, err)
	require.Equal(t, runtime.ProviderSearXNG, invoker.request.Provider)
	require.Equal(t, "AgentTeam", invoker.request.Query)
	require.Equal(t, 2, invoker.request.Limit)
	require.JSONEq(t, `{"results":[{"url":"https://example.com","title":"Example","snippet":"snippet"}]}`, string(result.Data))
}

func TestExecutorMapsSearchErrors(t *testing.T) {
	t.Parallel()

	executor := searchengine.NewExecutor(runtime.ProviderBrave, &recordingInvoker{
		err: &runtime.Error{
			Code:    runtime.ErrorRateLimited,
			Message: "provider rate limited",
		},
	})
	_, err := executor.Invoke(context.Background(), registry.BuiltinInvocation{
		WorkspaceID:           "wsp_1",
		AgentID:               "agt_1",
		AgentMemberID:         "mem_1",
		ServiceInstallationID: "svcinst_1",
		ExtensionName:         "brave",
		Action:                searchengine.ActionSearch,
		Input:                 json.RawMessage(`{"query":"AgentTeam"}`),
	})
	require.Error(t, err)
	code, ok := extensionflow.InvocationErrorCodeOf(err)
	require.True(t, ok)
	require.Equal(t, extensionflow.InvocationRateLimited, code)
}

type recordingInvoker struct {
	request runtime.SearchRequest
	result  runtime.SearchResult
	err     error
}

func (invoker *recordingInvoker) Search(
	_ context.Context,
	request runtime.SearchRequest,
) (runtime.SearchResult, error) {
	invoker.request = request
	return invoker.result, invoker.err
}