package pkulaw

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

	"github.com/stretchr/testify/require"

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

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

	testCases := []wireCase{
		{
			action: ActionSearchLawArticles,
			input:  `{"query":"  劳动合同解除  "}`,
			path:   "/mcp-law-agg/1.0.0/mcp",
			tool:   "search_article",
			arguments: `{
				"query":"劳动合同解除"
			}`,
		},
		{
			action: ActionGetLawArticle,
			input:  `{"title":" 中华人民共和国民法典 ","number":" 第一千零七十七条 "}`,
			path:   "/mcp-law-agg/1.0.0/mcp",
			tool:   "get_article",
			arguments: `{
				"title":"中华人民共和国民法典",
				"number":"第一千零七十七条"
			}`,
		},
		{
			action: ActionSearchLaws,
			input:  `{"title":" 民法典 ","fulltext":" 居住权 "}`,
			path:   "/mcp-law-agg/1.0.0/mcp",
			tool:   "get_law_list",
			arguments: `{
				"title":"民法典",
				"fulltext":"居住权"
			}`,
		},
		{
			action: ActionSearchCases,
			input:  `{"text":"  网络侵权责任  "}`,
			path:   "/mcp-law-agg/1.0.0/mcp",
			tool:   "search_case",
			arguments: `{
				"text":"网络侵权责任"
			}`,
		},
		{
			action: ActionSearchCaseDocuments,
			input:  `{"fulltext":"  正当防卫  "}`,
			path:   "/mcp-law-agg/1.0.0/mcp",
			tool:   "get_case_list",
			arguments: `{
				"fulltext":"正当防卫"
			}`,
		},
	}
	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.action, func(t *testing.T) {
			t.Parallel()
			assertPKULawWireCase(t, TopologyAggregate, testCase)
		})
	}
}

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

	testCases := []wireCase{
		{
			action: ActionSearchLawArticles,
			input:  `{"query":"劳动合同解除"}`,
			path:   "/mcp-law-search-service/mcp",
			tool:   "search_article",
			arguments: `{
				"text":"劳动合同解除"
			}`,
		},
		{
			action: ActionGetLawArticle,
			input:  `{"title":"中华人民共和国民法典","number":"第一千零七十七条"}`,
			path:   "/mcp-law-search-service/mcp",
			tool:   "get_article",
			arguments: `{
				"title":"中华人民共和国民法典",
				"number":"第一千零七十七条"
			}`,
		},
		{
			action: ActionSearchLaws,
			input:  `{"title":"民法典"}`,
			path:   "/mcp-law/mcp",
			tool:   "get_law_list",
			arguments: `{
				"lawInput":{"Title":"民法典"}
			}`,
		},
		{
			action: ActionSearchCases,
			input:  `{"text":"网络侵权责任"}`,
			path:   "/mcp-case-search-service/mcp",
			tool:   "search_case",
			arguments: `{
				"text":"网络侵权责任"
			}`,
		},
		{
			action: ActionSearchCaseDocuments,
			input:  `{"title":"指导案例","fulltext":"正当防卫"}`,
			path:   "/mcp-case/mcp",
			tool:   "get_case_list",
			arguments: `{
				"caseInput":{
					"Title":"指导案例",
					"Fulltext":"正当防卫"
				}
			}`,
		},
	}
	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.action, func(t *testing.T) {
			t.Parallel()
			assertPKULawWireCase(t, TopologyFallback, testCase)
		})
	}
}

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

	server := &pkulawMCPWireDouble{
		t:          t,
		credential: "pkulaw-test-secret",
		path:       "/mcp-law-agg/1.0.0/mcp",
		result: func(pkulawRPCRequest) any {
			return map[string]any{
				"content": []map[string]any{{
					"type": "text",
					"text": `{"records":[{"title":"民法典"}]}`,
				}},
			}
		},
	}
	executor := newWireExecutor(t, server, TopologyAggregate)
	result, err := executor.Invoke(
		context.Background(),
		registry.BuiltinInvocation{
			WorkspaceID:           "workspace-1",
			ServiceInstallationID: "installation-1",
			Action:                ActionSearchLawArticles,
			Input:                 json.RawMessage(`{"query":"民法典"}`),
		},
	)
	require.NoError(t, err)
	require.JSONEq(
		t,
		`{"result":{"records":[{"title":"民法典"}]}}`,
		string(result.Data),
	)
}

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

	server := &pkulawMCPWireDouble{
		t:          t,
		credential: "pkulaw-test-secret",
		path:       "/mcp-law-agg/1.0.0/mcp",
		result: func(pkulawRPCRequest) any {
			return map[string]any{
				"isError": true,
				"content": []map[string]any{{
					"type": "text",
					"text": `{"error":"failed"}`,
				}},
			}
		},
	}
	executor := newWireExecutor(t, server, TopologyAggregate)
	result, err := executor.Invoke(
		context.Background(),
		registry.BuiltinInvocation{
			WorkspaceID:           "workspace-1",
			ServiceInstallationID: "installation-1",
			Action:                ActionSearchLawArticles,
			Input:                 json.RawMessage(`{"query":"民法典"}`),
		},
	)
	require.Nil(t, result.Data)
	code, ok := extensionflow.InvocationErrorCodeOf(err)
	require.True(t, ok)
	require.Equal(t, extensionflow.InvocationUpstreamFailure, code)
}

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

	testCases := []struct {
		name   string
		config runner.ConnectionConfig
	}{
		{
			name: "unsupported topology",
			config: runner.ConnectionConfig{
				"topology": "future-topology",
			},
		},
		{
			name: "valid topology with extra field",
			config: runner.ConnectionConfig{
				"topology": TopologyAggregate,
				"endpoint": "aggregate",
			},
		},
		{
			name:   "missing topology",
			config: runner.ConnectionConfig{},
		},
		{
			name:   "nil config",
			config: nil,
		},
	}
	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.name, func(t *testing.T) {
			t.Parallel()

			server := &pkulawMCPWireDouble{
				t:          t,
				credential: "never-used",
				path:       "/never-used",
			}
			clients, err := newOperationClients(server)
			require.NoError(t, err)
			source := &staticConnectionSource{
				resolved: runner.ResolvedConnection{
					Credentials: runner.Credentials{
						"apiKey": "never-used",
					},
					Config: testCase.config,
				},
			}
			executor, err := newExecutor(source, clients)
			require.NoError(t, err)
			_, err = executor.Invoke(
				context.Background(),
				registry.BuiltinInvocation{
					WorkspaceID:           "workspace-1",
					ServiceInstallationID: "installation-1",
					Action:                ActionSearchCases,
					Input:                 json.RawMessage(`{"text":"案件"}`),
				},
			)
			code, ok := extensionflow.InvocationErrorCodeOf(err)
			require.True(t, ok)
			require.Equal(
				t,
				extensionflow.InvocationProviderAuthUnavailable,
				code,
			)
			require.Empty(t, server.recordedMethods())

			calls, contract := source.snapshot()
			require.Equal(t, 1, calls)
			require.Equal(t, ConnectionConfigContract(), contract)
		})
	}
}

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

	server := &pkulawMCPWireDouble{
		t:          t,
		credential: "never-used",
		path:       "/never-used",
	}
	clients, err := newOperationClients(server)
	require.NoError(t, err)
	source := &staticConnectionSource{
		err: errors.New("database contained pkulaw-secret"),
	}
	executor, err := newExecutor(source, clients)
	require.NoError(t, err)
	_, err = executor.Invoke(
		context.Background(),
		registry.BuiltinInvocation{
			WorkspaceID:           "workspace-1",
			ServiceInstallationID: "installation-1",
			Action:                "call-tool",
			Input:                 json.RawMessage(`{}`),
		},
	)
	code, ok := extensionflow.InvocationErrorCodeOf(err)
	require.True(t, ok)
	require.Equal(t, extensionflow.InvocationNotFound, code)
	require.Empty(t, server.recordedMethods())
	require.NotContains(t, err.Error(), "pkulaw-secret")
	calls, _ := source.snapshot()
	require.Zero(t, calls)
}

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

	server := &pkulawMCPWireDouble{
		t:          t,
		credential: "never-used",
		path:       "/never-used",
	}
	clients, err := newOperationClients(server)
	require.NoError(t, err)
	source := &staticConnectionSource{
		err: errors.New("database contained pkulaw-secret"),
	}
	executor, err := newExecutor(source, clients)
	require.NoError(t, err)
	_, err = executor.Invoke(
		context.Background(),
		registry.BuiltinInvocation{
			WorkspaceID:           "workspace-1",
			ServiceInstallationID: "installation-1",
			Action:                ActionSearchCases,
			Input:                 json.RawMessage(`{"text":"案件"}`),
		},
	)
	code, ok := extensionflow.InvocationErrorCodeOf(err)
	require.True(t, ok)
	require.Equal(t, extensionflow.InvocationProviderAuthUnavailable, code)
	require.Empty(t, server.recordedMethods())
	require.NotContains(t, err.Error(), "pkulaw-secret")
}

type wireCase struct {
	action    string
	input     string
	path      string
	tool      string
	arguments string
}

func assertPKULawWireCase(
	t *testing.T,
	topology string,
	testCase wireCase,
) {
	t.Helper()
	server := &pkulawMCPWireDouble{
		t:          t,
		credential: "pkulaw-test-secret",
		path:       testCase.path,
	}
	server.result = func(request pkulawRPCRequest) any {
		require.Equal(t, "tools/call", request.Method)
		var params struct {
			Name      string          `json:"name"`
			Arguments json.RawMessage `json:"arguments"`
		}
		require.NoError(t, json.Unmarshal(request.Params, &params))
		require.Equal(t, testCase.tool, params.Name)
		require.JSONEq(t, testCase.arguments, string(params.Arguments))
		return map[string]any{
			"content": []map[string]any{{
				"type": "text",
				"text": "must not be exposed",
			}},
			"structuredContent": map[string]any{
				"source":  "pkulaw",
				"records": []any{},
			},
		}
	}
	executor := newWireExecutor(t, server, topology)
	result, err := executor.Invoke(
		context.Background(),
		registry.BuiltinInvocation{
			WorkspaceID:           "workspace-1",
			ServiceInstallationID: "installation-1",
			Action:                testCase.action,
			Input:                 json.RawMessage(testCase.input),
		},
	)
	require.NoError(t, err)
	require.JSONEq(
		t,
		`{"result":{"source":"pkulaw","records":[]}}`,
		string(result.Data),
	)
	require.Equal(t, []string{
		"server/discover",
		"initialize",
		"notifications/initialized",
		"tools/call",
		"DELETE",
	}, server.recordedMethods())
}

func newWireExecutor(
	t *testing.T,
	server *pkulawMCPWireDouble,
	topology string,
) *Executor {
	t.Helper()
	clients, err := newOperationClients(server)
	require.NoError(t, err)
	executor, err := newExecutor(
		&staticConnectionSource{
			resolved: runner.ResolvedConnection{
				Credentials: runner.Credentials{
					"apiKey": "pkulaw-test-secret",
				},
				Config: runner.ConnectionConfig{
					"topology": topology,
				},
			},
		},
		clients,
	)
	require.NoError(t, err)
	return executor
}