package pkulaw

import (
	"encoding/json"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
	"github.com/tidbcloud/AgentTeam/internal/extension/runner/mcp"
)

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

	got, err := normalizeToolResult(mcp.ToolResult{
		StructuredContent: json.RawMessage(
			`{"records":[{"title":"中华人民共和国民法典"}]}`,
		),
		Content: []json.RawMessage{
			json.RawMessage(`{"type":"image","data":"ignored"}`),
			json.RawMessage(`{"type":"text","text":"not json"}`),
		},
	})
	require.NoError(t, err)
	require.JSONEq(
		t,
		`{"result":{"records":[{"title":"中华人民共和国民法典"}]}}`,
		string(got),
	)
}

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

	for _, testCase := range []struct {
		content string
		want    string
	}{
		{
			content: `{"type":"text","text":" [1, true, \"第一条\"] "}`,
			want:    `{"result":[1,true,"第一条"]}`,
		},
		{
			content: `{"type":"text","text":"{\"ok\":true}","annotations":{"audience":["user"],"priority":0.8},"_meta":{"trace":"opaque"}}`,
			want:    `{"result":{"ok":true}}`,
		},
	} {
		got, err := normalizeToolResult(mcp.ToolResult{
			Content: []json.RawMessage{json.RawMessage(testCase.content)},
		})
		require.NoError(t, err, testCase.content)
		require.JSONEq(t, testCase.want, string(got))
	}
}

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

	for _, raw := range []string{
		`true`,
		`false`,
		`42`,
		`"第一条"`,
		`[{"title":"民法典"}]`,
		`{"data":[],"error":null}`,
		`{"mcp_tool_error_code":0,"mcp_tool_error_msg":"ok","data":[]}`,
		`{"mcp_tool_error_code":"0","data":[]}`,
		`{"articles":[{"content":["第一条","第二条"]}]}`,
		`{"data":{"code":0,"message":"ok","content":[{"section":"第一条"}]}}`,
	} {
		got, err := normalizeToolResult(mcp.ToolResult{
			StructuredContent: json.RawMessage(raw),
		})
		require.NoError(t, err, raw)
		require.JSONEq(t, `{"result":`+raw+`}`, string(got))
	}
}

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

	oversized := `"` +
		strings.Repeat("x", runner.MaxResponseBytes) +
		`"`
	testCases := []struct {
		name   string
		result mcp.ToolResult
	}{
		{name: "missing result"},
		{
			name: "structured whitespace is present but invalid",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(" \n\t "),
				Content: []json.RawMessage{
					json.RawMessage(`{"type":"text","text":"{\"ok\":true}"}`),
				},
			},
		},
		{
			name: "structured null does not fall back",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(`null`),
				Content: []json.RawMessage{
					json.RawMessage(`{"type":"text","text":"{\"ok\":true}"}`),
				},
			},
		},
		{
			name: "wire structured null does not fall back",
			result: mcp.ToolResult{
				StructuredContentPresent: true,
				Content: []json.RawMessage{
					json.RawMessage(`{"type":"text","text":"{\"ok\":true}"}`),
				},
			},
		},
		{
			name: "empty fallback text",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(`{"type":"text","text":"  "}`),
			}},
		},
		{
			name: "multiple fallback content",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(`{"type":"text","text":"{\"ok\":true}"}`),
				json.RawMessage(`{"type":"text","text":"{\"ok\":true}"}`),
			}},
		},
		{
			name: "heterogeneous fallback content",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(`{"type":"image","data":"x"}`),
			}},
		},
		{
			name: "fallback text content has unknown property",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(
					`{"type":"text","text":"{\"ok\":true}","unknown":{}}`,
				),
			}},
		},
		{
			name: "fallback annotations is not an object",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(
					`{"type":"text","text":"{\"ok\":true}","annotations":[]}`,
				),
			}},
		},
		{
			name: "fallback meta is not an object",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(
					`{"type":"text","text":"{\"ok\":true}","_meta":"opaque"}`,
				),
			}},
		},
		{
			name: "fallback content duplicate type",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(
					`{"type":"image","type":"text","text":"{\"ok\":true}"}`,
				),
			}},
		},
		{
			name: "fallback non json",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(`{"type":"text","text":"not json"}`),
			}},
		},
		{
			name: "fallback multiple json documents",
			result: mcp.ToolResult{Content: []json.RawMessage{
				json.RawMessage(
					`{"type":"text","text":"{\"ok\":true} {\"ok\":true}"}`,
				),
			}},
		},
		{
			name: "structured malformed json",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(`{"value":NaN}`),
			},
		},
		{
			name: "structured invalid utf8",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage{
					'{', '"', 'x', '"', ':', '"', 0xff, '"', '}',
				},
			},
		},
		{
			name: "structured non finite number",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(`{"value":1e999}`),
			},
		},
		{
			name: "recursive duplicate keys",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"data":{"title":"bad","title":"hidden"}}`,
				),
			},
		},
		{
			name: "excessive JSON nesting",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					strings.Repeat("[", 130) +
						"0" +
						strings.Repeat("]", 130),
				),
			},
		},
		{
			name: "embedded isError false MCP content envelope",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"isError":false,"content":[{"type":"text","text":"[]"}]}`,
				),
			},
		},
		{
			name: "embedded capital IsError",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"IsError":true,"records":[]}`,
				),
			},
		},
		{
			name: "embedded structured content MCP envelope",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"content":[],"structuredContent":{"records":[1]}}`,
				),
			},
		},
		{
			name: "nested high confidence MCP content envelope",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"data":{"content":[{"type":"text","text":"[]"}]}}`,
				),
			},
		},
		{
			name: "embedded JSON RPC envelope",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"jsonrpc":"2.0","id":1,"result":{"records":[1]}}`,
				),
			},
		},
		{
			name: "embedded interactive MCP result",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"resultType":"input_required","inputRequests":{"query":{}}}`,
				),
			},
		},
		{
			name: "embedded MCP tool error code",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"mcp_tool_error_code":1001,"mcp_tool_error_msg":"failed"}`,
				),
			},
		},
		{
			name: "embedded error object",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"data":null,"error":{"code":"QUERY_FAILED","message":"failed"}}`,
				),
			},
		},
		{
			name: "nested nonzero code and message",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"data":{"code":500,"message":"query failed","records":[]}}`,
				),
			},
		},
		{
			name: "nested string nonzero code and message",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"data":{"code":"E_QUERY","message":"query failed"}}`,
				),
			},
		},
		{
			name: "nested HTTP status phrase and message",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"data":{"status":"500 Internal Server Error","message":"query failed"}}`,
				),
			},
		},
		{
			name: "nested embedded error string",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"data":{"error":"query failed"}}`,
				),
			},
		},
		{
			name: "success false",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(
					`{"success":false,"message":"failed"}`,
				),
			},
		},
		{
			name: "null result",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(`null`),
			},
		},
		{
			name: "empty object result",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(`{}`),
			},
		},
		{
			name: "empty array result",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(`[]`),
			},
		},
		{
			name: "empty string result",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(`"  "`),
			},
		},
		{
			name: "oversized result",
			result: mcp.ToolResult{
				StructuredContent: json.RawMessage(oversized),
			},
		},
	}

	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.name, func(t *testing.T) {
			t.Parallel()
			_, err := normalizeToolResult(testCase.result)
			require.Error(t, err)
			require.Equal(
				t,
				runner.ErrorUpstreamFailure,
				runner.CodeOf(err),
			)
		})
	}
}