package runtime

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync/atomic"
	"testing"

	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"

	"github.com/stretchr/testify/require"
)

func TestMaaSClientSendsExactNativeExplorationContract(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		require.Equal(t, "/maas/v1/chat/completions", request.URL.Path)
		require.Equal(t, "Bearer capability", request.Header.Get("Authorization"))
		require.Equal(t, "kqa_1", request.Header.Get("Idempotency-Key"))
		require.Equal(t, chatCompletionsProfileOpenAIRequiredToolsV1, request.Header.Get(chatCompletionsProfileHeader))
		require.Equal(t, "application/json", request.Header.Get("Accept"))
		var payload map[string]any
		require.NoError(t, json.NewDecoder(request.Body).Decode(&payload))
		require.Equal(t, "model-a", payload["model"])
		require.Equal(t, false, payload["stream"])
		require.Equal(t, false, payload["store"])
		require.Equal(t, float64(1), payload["n"])
		require.Equal(t, "high", payload["reasoning_effort"])
		require.Equal(t, false, payload["parallel_tool_calls"])
		require.Equal(t, "required", payload["tool_choice"])
		require.NotContains(t, payload, "response_format")

		messages := payload["messages"].([]any)
		require.Len(t, messages, 2)
		require.Equal(t, "system", messages[0].(map[string]any)["role"])
		require.Equal(t, "user", messages[1].(map[string]any)["role"])
		tools := payload["tools"].([]any)
		require.Len(t, tools, 6)
		names := make([]string, 0, len(tools))
		for _, rawTool := range tools {
			tool := rawTool.(map[string]any)
			require.Equal(t, "function", tool["type"])
			function := tool["function"].(map[string]any)
			names = append(names, function["name"].(string))
			parameters := function["parameters"].(map[string]any)
			require.Equal(t, "object", parameters["type"])
			require.Equal(t, false, parameters["additionalProperties"])
			properties := parameters["properties"].(map[string]any)
			if function["name"] == knowledgeNoAnswerFunction {
				require.Empty(t, properties)
			} else {
				require.NotEmpty(t, properties)
			}
			switch function["name"] {
			case knowledgeSearchFunction:
				require.ElementsMatch(t, []any{"query"}, parameters["required"])
				require.Equal(t, float64(maxSearchPatternBytes), properties["query"].(map[string]any)["maxLength"])
				require.Equal(t, float64(1), properties["rootAliases"].(map[string]any)["minItems"])
				require.Equal(t, float64(maxSearchPathGlobs), properties["pathGlobs"].(map[string]any)["maxItems"])
			case knowledgeListFunction:
				require.Len(t, parameters["anyOf"], 2)
				require.NotContains(t, parameters, "oneOf")
				require.Equal(t, `^r[1-9][0-9]*$`, properties["rootAlias"].(map[string]any)["pattern"])
				require.Equal(t, float64(maxReferenceBytes), properties["rootAlias"].(map[string]any)["maxLength"])
				require.Equal(t, `^d[1-9][0-9]*$`, properties["directoryRef"].(map[string]any)["pattern"])
			case knowledgeReadFunction:
				require.Equal(t, []any{"fileRef", "lineStart", "lineEnd"}, parameters["required"])
				require.Equal(t, float64(maxReferenceBytes), properties["fileRef"].(map[string]any)["maxLength"])
				require.Equal(t, "integer", properties["lineStart"].(map[string]any)["type"])
				require.Equal(t, float64(1), properties["lineEnd"].(map[string]any)["minimum"])
			case knowledgeRefFunction:
				require.Equal(t, []any{"rootAlias", "ref", "refType"}, parameters["required"])
				require.Equal(t, float64(maxRefNameBytes), properties["ref"].(map[string]any)["maxLength"])
				require.Equal(t, []any{RefKindBranch, RefKindTag}, properties["refType"].(map[string]any)["enum"])
			case knowledgeAnswerFunction:
				description := function["description"].(string)
				require.Contains(t, description, "We recommend citing supporting reads")
				require.Contains(t, description, "a claim cannot contain only evidence markers")
				require.NotContains(t, description, "containing exact [E1] markers")
				require.Equal(t, []any{"claims"}, parameters["required"])
				require.NotContains(t, parameters, "oneOf")
				claimsSchema := properties["claims"].(map[string]any)
				require.Equal(t, "array", claimsSchema["type"])
				require.NotContains(t, claimsSchema, "anyOf")
				require.NotContains(t, claimsSchema, "minItems")
				require.Equal(t, "string", claimsSchema["items"].(map[string]any)["type"])
			case knowledgeNoAnswerFunction:
				require.NotContains(t, parameters, "required")
			}
		}
		require.Equal(t, []string{
			knowledgeSearchFunction, knowledgeListFunction, knowledgeReadFunction,
			knowledgeRefFunction, knowledgeAnswerFunction, knowledgeNoAnswerFunction,
		}, names)
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("X-AgentTeam-MaaS-Usage-Event-ID", "usage_header")
		_, _ = fmt.Fprint(w, validChatCompletion("response_1", "call_1", knowledgeSearchFunction, `{"query":"server"}`))
	}))
	defer server.Close()

	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: "kqa_1",
		Mode: InferenceModeExploration, System: "system",
		Messages: []InferenceMessage{{Role: "user", Content: `EXPLORE THIS TURN: {"authorizedRoots":[]}`}},
		Profile:  mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
	})
	require.NoError(t, err)
	require.Equal(t, "assistant", result.Message.Role)
	require.True(t, result.Message.ContentNull)
	require.Len(t, result.Message.ToolCalls, 1)
	require.Equal(t, knowledgeSearchFunction, result.Message.ToolCalls[0].Function.Name)
	require.JSONEq(t, `{"query":"server"}`, result.Message.ToolCalls[0].Function.Arguments)
	require.Equal(t, "usage_header", result.SettlementRef)
	require.Equal(t, float64(17), result.Usage["input_tokens"])
}

func TestMaaSClientResolvesExactProfileFromMaaSModels(t *testing.T) {
	profile := mustSupportedProfile(t, chatCompletionsProfileDeepSeekV4ThinkingV1)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		require.Equal(t, http.MethodGet, request.Method)
		require.Equal(t, "/maas/v1/models", request.URL.Path)
		require.Equal(t, "Bearer capability", request.Header.Get("Authorization"))
		require.NoError(t, json.NewEncoder(w).Encode(map[string]any{
			"object": "list",
			"data": []any{map[string]any{
				"id":                                  "agentteam/deepseek-v4-pro",
				"x_agentteam_executable_wire_api_families": []string{"chat_completions"},
				"x_agentteam_chat_completions_features": []string{
					"chat_completions.tools",
					"chat_completions.reasoning_effort",
				},
				"x_agentteam_chat_completions_compatibility_profile": profile,
			}},
		}))
	}))
	defer server.Close()

	resolved, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).
		ResolveChatCompletionsProfile(context.Background(), "agentteam/deepseek-v4-pro", "capability")

	require.NoError(t, err)
	require.Equal(t, chatCompletionsProfileDeepSeekV4ThinkingV1, resolved.ID)
	require.Equal(t, "max_tokens", resolved.OutputTokenField)
	require.Equal(t, "high", resolved.ReasoningEffortMap["low"])
	require.Equal(t, "high", resolved.ReasoningEffortMap["medium"])
}

func TestMaaSClientFailsClosedUnknownMissingOrConflictingProfile(t *testing.T) {
	exactProfile := func() map[string]any {
		profile := mustSupportedProfile(t, chatCompletionsProfileDeepSeekV4ThinkingV1)
		raw, err := json.Marshal(profile)
		require.NoError(t, err)
		var result map[string]any
		require.NoError(t, json.Unmarshal(raw, &result))
		return result
	}
	tests := []struct {
		name    string
		profile any
	}{
		{name: "missing"},
		{name: "unknown", profile: map[string]any{"id": "chat_completions_future_v99"}},
		{name: "extra semantic field", profile: func() map[string]any {
			profile := exactProfile()
			profile["future_semantic"] = true
			return profile
		}()},
		{name: "missing false field", profile: func() map[string]any {
			profile := exactProfile()
			delete(profile, "supports_n_one")
			return profile
		}()},
		{name: "conflicting", profile: func() ChatCompletionsCompatibilityProfile {
			profile := mustSupportedProfile(t, chatCompletionsProfileDeepSeekV4ThinkingV1)
			profile.OutputTokenField = "max_completion_tokens"
			return profile
		}()},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				require.NoError(t, json.NewEncoder(w).Encode(map[string]any{
					"data": []any{map[string]any{
						"id":                                  "model-a",
						"x_agentteam_executable_wire_api_families": []string{"chat_completions"},
						"x_agentteam_chat_completions_features": []string{
							"chat_completions.tools",
							"chat_completions.reasoning_effort",
						},
						"x_agentteam_chat_completions_compatibility_profile": test.profile,
					}},
				}))
			}))
			defer server.Close()

			_, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).
				ResolveChatCompletionsProfile(context.Background(), "model-a", "capability")

			var runtimeErr *Error
			require.ErrorAs(t, err, &runtimeErr)
			require.Equal(t, "knowledge_qa.compatibility_profile_unavailable", runtimeErr.Code)
			require.NotContains(t, err.Error(), "future_v99")
		})
	}
}

func TestMaaSClientFailsClosedWhenProfileDrivenRequestFeatureIsMissing(t *testing.T) {
	profile := mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1)
	for _, missing := range []string{
		"chat_completions.tool_choice",
		"chat_completions.parallel_tool_calls",
	} {
		t.Run(missing, func(t *testing.T) {
			features := []string{
				"chat_completions.tools",
				"chat_completions.tool_choice",
				"chat_completions.parallel_tool_calls",
				"chat_completions.reasoning_effort",
			}
			features = slices.DeleteFunc(features, func(feature string) bool { return feature == missing })
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				require.NoError(t, json.NewEncoder(w).Encode(map[string]any{
					"data": []any{map[string]any{
						"id":                                  "model-a",
						"x_agentteam_executable_wire_api_families": []string{"chat_completions"},
						"x_agentteam_chat_completions_features":    features,
						"x_agentteam_chat_completions_compatibility_profile": profile,
					}},
				}))
			}))
			defer server.Close()

			_, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).
				ResolveChatCompletionsProfile(context.Background(), "model-a", "capability")

			var runtimeErr *Error
			require.ErrorAs(t, err, &runtimeErr)
			require.Equal(t, "knowledge_qa.compatibility_profile_unavailable", runtimeErr.Code)
		})
	}
}

func TestMaaSClientSendsExactDeepSeekShapeAndPreservesReasoningReplay(t *testing.T) {
	profile := mustSupportedProfile(t, chatCompletionsProfileDeepSeekV4ThinkingV1)
	var requests atomic.Int32
	const reasoning = "exact reasoning bytes π preserved"
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		round := requests.Add(1)
		require.Equal(t, chatCompletionsProfileDeepSeekV4ThinkingV1, request.Header.Get(chatCompletionsProfileHeader))
		var payload map[string]any
		require.NoError(t, json.NewDecoder(request.Body).Decode(&payload))
		require.Equal(t, "high", payload["reasoning_effort"])
		require.NotContains(t, payload, "store")
		require.NotContains(t, payload, "n")
		require.NotContains(t, payload, "tool_choice")
		require.NotContains(t, payload, "parallel_tool_calls")
		require.NotContains(t, payload, "max_completion_tokens")
		require.Equal(t, float64(512), payload["max_tokens"])
		messages := payload["messages"].([]any)
		if round == 1 {
			require.Len(t, messages, 2)
			_, _ = fmt.Fprintf(w,
				`{"id":"response_1","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":"","reasoning_content":%q,"tool_calls":[{"id":"call_search","type":"function","function":{"name":"knowledge_search","arguments":"{\"query\":\"server\"}"}},{"id":"call_list","type":"function","function":{"name":"knowledge_list","arguments":"{\"rootAlias\":\"r1\"}"}}]}}],"usage":{"input_tokens":17,"output_tokens":3}}`,
				reasoning,
			)
			return
		}
		require.Len(t, messages, 5)
		assistant := messages[2].(map[string]any)
		require.Equal(t, "", assistant["content"])
		require.Equal(t, reasoning, assistant["reasoning_content"])
		require.Len(t, assistant["tool_calls"].([]any), 2)
		require.Equal(t, "call_search", messages[3].(map[string]any)["tool_call_id"])
		require.Equal(t, "call_list", messages[4].(map[string]any)["tool_call_id"])
		_, _ = fmt.Fprintf(w,
			`{"id":"response_2","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":"","reasoning_content":"finish reasoning","tool_calls":[{"id":"call_answer","type":"function","function":{"name":"knowledge_answer","arguments":"{\"claims\":[\"Grounded. [E1]\"]}"}}]}}],"usage":{"input_tokens":23,"output_tokens":4}}`,
		)
	}))
	defer server.Close()
	client := MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}
	first, err := client.Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: "round_1",
		Mode: InferenceModeExploration, System: "system", MaxOutputTokens: 512, Profile: profile,
		Messages: []InferenceMessage{{Role: "user", Content: "explore"}},
	})
	require.NoError(t, err)
	require.NotNil(t, first.Message.ReasoningContent)
	require.Equal(t, reasoning, *first.Message.ReasoningContent)
	require.False(t, first.Message.ContentNull)
	require.Len(t, first.Message.ToolCalls, 2)

	second, err := client.Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: "round_2",
		Mode: InferenceModeExploration, System: "system", MaxOutputTokens: 512, Profile: profile,
		Messages: []InferenceMessage{
			{Role: "user", Content: "explore"},
			first.Message,
			{Role: "tool", ToolCallID: "call_search", Content: `{"ok":true}`},
			{Role: "tool", ToolCallID: "call_list", Content: `{"ok":true}`},
		},
	})
	require.NoError(t, err)
	require.Equal(t, knowledgeAnswerFunction, second.Message.ToolCalls[0].Function.Name)
	action, err := decodeExplorationToolCall(second.Message.ToolCalls[0])
	require.NoError(t, err)
	require.Equal(t, []string{"Grounded. [E1]"}, action.Answer.Claims)
	require.Equal(t, int32(2), requests.Load())
}

func TestMaaSClientFailsClosedWhenResolvedProfileRemapsExplorationReasoning(t *testing.T) {
	profile := mustSupportedProfile(t, chatCompletionsProfileDeepSeekV4ThinkingV1)
	var requests atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		requests.Add(1)
	}))
	defer server.Close()

	for _, reasonLevel := range []string{"low", "medium"} {
		_, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(
			context.Background(),
			InferenceRequest{
				Model: "model-a", ReasonLevel: reasonLevel, Capability: "capability",
				IdempotencyKey: "profile_changed_" + reasonLevel,
				Mode:           InferenceModeExploration, System: "system", Profile: profile,
				Messages: []InferenceMessage{{Role: "user", Content: "explore"}},
			},
		)

		var runtimeErr *Error
		require.ErrorAs(t, err, &runtimeErr)
		require.Equal(t, CodeInvalidRequest, runtimeErr.Code)
	}
	require.Equal(t, int32(0), requests.Load())
}

func TestMaaSClientDeepSeekRejectsMissingReasoningOrNullContentWithoutRepair(t *testing.T) {
	profile := mustSupportedProfile(t, chatCompletionsProfileDeepSeekV4ThinkingV1)
	const privateReasoning = "provider-private-reasoning-sentinel"
	tests := []struct {
		name string
		body string
	}{
		{name: "plain text", body: `{"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"answer"}}]}`},
		{name: "missing reasoning", body: `{"choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"knowledge_search","arguments":"{}"}}]}}]}`},
		{name: "null content", body: `{"choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"reasoning_content":"provider-private-reasoning-sentinel","tool_calls":[{"id":"call_1","type":"function","function":{"name":"knowledge_search","arguments":"{}"}}]}}]}`},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var calls atomic.Int32
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				calls.Add(1)
				_, _ = fmt.Fprint(w, test.body)
			}))
			defer server.Close()

			_, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
				Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: "invalid",
				Mode: InferenceModeExploration, System: "system", Profile: profile,
				Messages: []InferenceMessage{{Role: "user", Content: "explore"}},
			})

			var runtimeErr *Error
			require.ErrorAs(t, err, &runtimeErr)
			require.Equal(t, CodeResponseInvalid, runtimeErr.Code)
			require.NotContains(t, err.Error(), privateReasoning)
			require.Equal(t, int32(1), calls.Load())
		})
	}
}

func TestMaaSClientDeepSeekSummaryUsesMaxTokens(t *testing.T) {
	profile := mustSupportedProfile(t, chatCompletionsProfileDeepSeekV4ThinkingV1)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		var payload map[string]any
		require.NoError(t, json.NewDecoder(request.Body).Decode(&payload))
		require.Equal(t, float64(contextSummaryMaxTokens), payload["max_tokens"])
		require.NotContains(t, payload, "max_completion_tokens")
		require.Equal(t, "high", payload["reasoning_effort"])
		_, _ = fmt.Fprint(w,
			`{"id":"summary","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":"","reasoning_content":"summary reasoning","tool_calls":[{"id":"call_summary","type":"function","function":{"name":"knowledge_context_summary","arguments":"{\"summary\":\"faithful\"}"}}]}}]}`,
		)
	}))
	defer server.Close()

	_, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "low", Capability: "capability", IdempotencyKey: "summary",
		Mode: InferenceModeContextSummary, System: contextSummarySystemPrompt,
		Messages:        []InferenceMessage{{Role: "user", Content: "compact"}},
		MaxOutputTokens: contextSummaryMaxTokens, Profile: profile,
	})
	require.NoError(t, err)
}

func TestMaaSClientRoundTripsFullAssistantAndMatchingToolMessages(t *testing.T) {
	var requests atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		requests.Add(1)
		var payload map[string]any
		require.NoError(t, json.NewDecoder(request.Body).Decode(&payload))
		messages := payload["messages"].([]any)
		require.Len(t, messages, 4)
		assistant := messages[2].(map[string]any)
		require.Equal(t, "assistant", assistant["role"])
		require.Contains(t, assistant, "content")
		require.Nil(t, assistant["content"])
		calls := assistant["tool_calls"].([]any)
		require.Len(t, calls, 1)
		require.Equal(t, "call_search", calls[0].(map[string]any)["id"])
		tool := messages[3].(map[string]any)
		require.Equal(t, "tool", tool["role"])
		require.Equal(t, "call_search", tool["tool_call_id"])
		require.JSONEq(t, `{"ok":true,"matches":[]}`, tool["content"].(string))
		_, _ = fmt.Fprint(w, validChatCompletion("response_2", "call_no_answer", knowledgeNoAnswerFunction, `{}`))
	}))
	defer server.Close()

	previous := InferenceToolCall{
		ID: "call_search", Type: "function",
		Function: InferenceFunctionCall{Name: knowledgeSearchFunction, Arguments: `{"query":"server"}`},
	}
	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "medium", Capability: "capability", IdempotencyKey: "kqa_2",
		Mode: InferenceModeExploration, System: "system",
		Messages: []InferenceMessage{
			{Role: "user", Content: "explore"},
			{Role: "assistant", ContentNull: true, ToolCalls: []InferenceToolCall{previous}},
			{Role: "tool", ToolCallID: "call_search", Content: `{"ok":true,"matches":[]}`},
		},
		Profile: mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
	})
	require.NoError(t, err)
	require.Equal(t, int32(1), requests.Load())
	require.Equal(t, knowledgeNoAnswerFunction, result.Message.ToolCalls[0].Function.Name)
}

func TestMaaSClientUsesSeparateSummaryFunctionSchema(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		var payload map[string]any
		require.NoError(t, json.NewDecoder(request.Body).Decode(&payload))
		require.Equal(t, "low", payload["reasoning_effort"])
		require.Equal(t, float64(contextSummaryMaxTokens), payload["max_completion_tokens"])
		tools := payload["tools"].([]any)
		require.Len(t, tools, 1)
		function := tools[0].(map[string]any)["function"].(map[string]any)
		require.Equal(t, knowledgeContextSummaryFunction, function["name"])
		parameters := function["parameters"].(map[string]any)
		summary := parameters["properties"].(map[string]any)["summary"].(map[string]any)
		require.Equal(t, float64(maxContextSummaryCharacters), summary["maxLength"])
		encoded, err := json.Marshal(tools)
		require.NoError(t, err)
		require.NotContains(t, string(encoded), knowledgeSearchFunction)
		require.NotContains(t, string(encoded), knowledgeAnswerFunction)
		require.NotContains(t, string(encoded), knowledgeNoAnswerFunction)
		_, _ = fmt.Fprint(w, validChatCompletion(
			"response_summary", "call_summary", knowledgeContextSummaryFunction, `{"summary":"faithful context"}`,
		))
	}))
	defer server.Close()

	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "low", Capability: "capability", IdempotencyKey: "kqa_summary",
		Mode: InferenceModeContextSummary, System: contextSummarySystemPrompt,
		Messages: []InferenceMessage{{Role: "user", Content: "compact"}}, MaxOutputTokens: contextSummaryMaxTokens,
		Profile: mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
	})
	require.NoError(t, err)
	require.Equal(t, knowledgeContextSummaryFunction, result.Message.ToolCalls[0].Function.Name)
}

func TestSchemaValidContextSummaryFitsAssistantEnvelope(t *testing.T) {
	summaryArguments, err := json.Marshal(map[string]string{
		"summary": strings.Repeat("\x00", maxContextSummaryCharacters),
	})
	require.NoError(t, err)
	summary := nativeToolResult(knowledgeContextSummaryFunction, string(summaryArguments))
	require.LessOrEqual(t, inferenceMessageSize(summary.Message), maxExplorationMessage)
}

func TestMaaSClientRejectsInvalidNativeToolCallResponsesWithoutTextFallback(t *testing.T) {
	tests := []struct {
		name string
		body string
	}{
		{name: "plain text", body: `{"id":"r1","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"{\"type\":\"answer\"}"}}],"usage":{"input_tokens":1}}`},
		{name: "missing choices", body: `{"id":"r0","usage":{"input_tokens":1}}`},
		{name: "missing choice index", body: `{"id":"r0","choices":[{"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"tool_calls":[{"id":"a","type":"function","function":{"name":"knowledge_search","arguments":"{}"}}]}}]}`},
		{name: "wrong choice index", body: `{"id":"r0","choices":[{"index":1,"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"tool_calls":[{"id":"a","type":"function","function":{"name":"knowledge_search","arguments":"{}"}}]}}]}`},
		{name: "missing call", body: `{"id":"r0","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":null}}]}`},
		{name: "multiple choices", body: `{"id":"r2","choices":[{},{}],"usage":{"input_tokens":2}}`},
		{name: "duplicate call ids", body: `{"id":"r3","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"tool_calls":[{"id":"a","type":"function","function":{"name":"knowledge_search","arguments":"{}"}},{"id":"a","type":"function","function":{"name":"knowledge_list","arguments":"{}"}}]}}],"usage":{"input_tokens":3}}`},
		{name: "unknown call", body: validChatCompletion("r4", "call_1", "legacy_answer", `{}`)},
		{name: "missing id", body: validChatCompletion("r5", "", knowledgeSearchFunction, `{"query":"server"}`)},
		{name: "untrimmed id", body: validChatCompletion("r5", " call_1 ", knowledgeSearchFunction, `{"query":"server"}`)},
		{name: "untrimmed name", body: validChatCompletion("r5", "call_1", " knowledge_search ", `{"query":"server"}`)},
		{name: "wrong type", body: strings.Replace(validChatCompletion("r5", "call_1", knowledgeSearchFunction, `{"query":"server"}`), `"type":"function"`, `"type":"custom"`, 1)},
		{name: "arguments not object", body: validChatCompletion("r6", "call_1", knowledgeSearchFunction, `[]`)},
		{name: "malformed json", body: `{"id":"r7","choices":[`},
		{name: "oversized assistant message", body: strings.Replace(
			validChatCompletion("r8", "call_1", knowledgeSearchFunction, `{"query":"server"}`),
			`"content":null`, `"content":`+fmt.Sprintf("%q", strings.Repeat("x", maxExplorationMessage)), 1,
		)},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			result, _, err := readMaaSResponse(
				strings.NewReader(test.body),
				InferenceModeExploration,
				mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
				false,
				false,
			)
			var runtimeErr *Error
			require.ErrorAs(t, err, &runtimeErr)
			require.Equal(t, CodeResponseInvalid, runtimeErr.Code)
			require.Empty(t, result.Message.ToolCalls)
		})
	}
}

func TestMaaSClientRejectsMultipleContextSummaryCalls(t *testing.T) {
	body := `{"id":"summary","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"tool_calls":[{"id":"a","type":"function","function":{"name":"knowledge_context_summary","arguments":"{\"summary\":\"first\"}"}},{"id":"b","type":"function","function":{"name":"knowledge_context_summary","arguments":"{\"summary\":\"second\"}"}}]}}]}`
	result, _, err := readMaaSResponse(
		strings.NewReader(body),
		InferenceModeContextSummary,
		mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
		false,
		false,
	)
	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeResponseInvalid, runtimeErr.Code)
	require.Empty(t, result.Message.ToolCalls)
}

func TestMaaSClientRejectsOverLimitExplorationCallBatch(t *testing.T) {
	calls := make([]map[string]any, maxExplorationCallsPerResponse+1)
	for index := range calls {
		calls[index] = map[string]any{
			"id": fmt.Sprintf("call_%d", index), "type": "function",
			"function": map[string]any{"name": knowledgeSearchFunction, "arguments": `{}`},
		}
	}
	body, err := json.Marshal(map[string]any{
		"choices": []any{map[string]any{
			"index": 0, "finish_reason": "tool_calls",
			"message": map[string]any{"role": "assistant", "content": nil, "tool_calls": calls},
		}},
	})
	require.NoError(t, err)
	result, _, err := readMaaSResponse(
		bytes.NewReader(body),
		InferenceModeExploration,
		mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
		false,
		false,
	)
	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeResponseInvalid, runtimeErr.Code)
	require.Empty(t, result.Message.ToolCalls)
}

func TestMaaSClientRejectsMismatchedToolResultIDBeforeDispatch(t *testing.T) {
	call := InferenceToolCall{
		ID: "call_expected", Type: "function",
		Function: InferenceFunctionCall{Name: knowledgeSearchFunction, Arguments: `{"query":"server"}`},
	}
	err := validateInferenceMessageHistory(
		InferenceModeExploration,
		[]InferenceMessage{
			{Role: "user", Content: "explore"},
			{Role: "assistant", ContentNull: true, ToolCalls: []InferenceToolCall{call}},
			{Role: "tool", ToolCallID: "call_forged", Content: `{}`},
		},
		mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
		false,
	)
	require.Error(t, err)
}

func TestMaaSClientRetainsCompleteDebugOutputAndAccountingOnInvalidResponse(t *testing.T) {
	rawResponse := fmt.Sprintf(
		`{"id":"response_invalid","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":%q}}],"usage":{"input_tokens":17}}`,
		strings.Repeat("full-debug-output-", 5000),
	)
	require.Greater(t, len(rawResponse), 64<<10)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = fmt.Fprint(w, rawResponse)
	}))
	defer server.Close()
	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: "kqa_invalid",
		Mode: InferenceModeExploration, System: "system", Messages: []InferenceMessage{{Role: "user", Content: "explore"}},
		Profile:          mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
		DebugModelOutput: true,
	})
	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeResponseInvalid, runtimeErr.Code)
	require.Equal(t, float64(17), result.Usage["input_tokens"])
	require.Equal(t, "response_invalid", result.SettlementRef)
	require.Equal(t, rawResponse, result.RawResponse)
	require.NotContains(t, result.RawResponse, "debug output truncated")
}

func TestMaaSClientPreservesStructuredProviderFailureCodeAndHeaderSettlement(t *testing.T) {
	for _, test := range []struct {
		name   string
		status int
		code   string
	}{
		{name: "context window", status: http.StatusConflict, code: "upstream.context_window.exceeded"},
		{name: "timeout", status: http.StatusGatewayTimeout, code: upstreamTimeoutCode},
		{name: "unavailable", status: http.StatusServiceUnavailable, code: upstreamUnavailableCode},
		{name: "rate limit", status: http.StatusTooManyRequests, code: upstreamRateLimitCode},
	} {
		t.Run(test.name, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				w.Header().Set("X-AgentTeam-MaaS-Usage-Event-ID", "usage_rejected")
				w.WriteHeader(test.status)
				_, _ = fmt.Fprintf(w, `{"error":{"code":%q,"retryable":true,"requires_new_idempotency_key":true}}`, test.code)
			}))
			defer server.Close()
			result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
				Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: "kqa_rejected",
				Mode: InferenceModeExploration, System: "system", Messages: []InferenceMessage{{Role: "user", Content: "explore"}},
				Profile: mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
			})
			var runtimeErr *Error
			require.ErrorAs(t, err, &runtimeErr)
			require.Equal(t, test.code, runtimeErr.Code)
			require.Equal(t, "usage_rejected", result.SettlementRef)
		})
	}
}

func TestMaaSClientTreatsHostedRequestStatusAsAccountingRecovery(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = fmt.Fprint(w, `{"object":"hosted.request.status","request_status":"reconciliation_pending","usage_event_id":"hue_recovery_1","usage":{"input_tokens":17}}`)
	}))
	defer server.Close()
	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(context.Background(), InferenceRequest{
		Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: "kqa_recovery",
		Mode: InferenceModeExploration, System: "system", Messages: []InferenceMessage{{Role: "user", Content: "explore"}},
		Profile: mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
	})
	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, "knowledge_qa.maas_accounting_recovery", runtimeErr.Code)
	require.Equal(t, float64(17), result.Usage["input_tokens"])
	require.Equal(t, "hue_recovery_1", result.SettlementRef)
	require.Empty(t, result.Message.ToolCalls)
}

func TestMaaSClientBoundsSuccessfulAccountingAtProducerBoundary(t *testing.T) {
	rawRef := strings.Repeat("private-settlement-reference", 16)
	require.Greater(t, len(rawRef), protocol.MaxAccountingSettlementReferenceBytes)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("X-AgentTeam-MaaS-Usage-Event-ID", rawRef)
		_, _ = fmt.Fprint(w, validChatCompletion(
			"response_1", "call_1", knowledgeSearchFunction, `{"query":"server"}`,
		))
	}))
	defer server.Close()

	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(
		context.Background(), maasGenerateTestRequest(t, "kqa_oversized_success_ref"),
	)

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeMaaSAccountingRecovery, runtimeErr.Code)
	require.Equal(t, accountingRecoveryReference("kqa_oversized_success_ref"), result.SettlementRef)
	require.NotContains(t, result.SettlementRef, rawRef)
	require.Equal(t, float64(17), result.Usage["input_tokens"])
	require.Len(t, result.Message.ToolCalls, 1)
}

func TestMaaSClientBoundsRejectedAccountingAtProducerBoundary(t *testing.T) {
	rawRef := strings.Repeat("private-settlement-reference", 16)
	require.Greater(t, len(rawRef), protocol.MaxAccountingSettlementReferenceBytes)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("X-AgentTeam-MaaS-Usage-Event-ID", rawRef)
		w.WriteHeader(http.StatusServiceUnavailable)
		_, _ = fmt.Fprint(w, `{"error":{"code":"maas.unavailable","retryable":true}}`)
	}))
	defer server.Close()

	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(
		context.Background(), maasGenerateTestRequest(t, "kqa_oversized_rejected_ref"),
	)

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeMaaSAccountingRecovery, runtimeErr.Code)
	require.Equal(t, accountingRecoveryReference("kqa_oversized_rejected_ref"), result.SettlementRef)
	require.NotContains(t, result.SettlementRef, rawRef)
}

func TestMaaSClientProjectsOversizedUsageAtProducerBoundary(t *testing.T) {
	rawUsage := map[string]any{
		"input_tokens":             17,
		"output_tokens":            3,
		"provider_private_payload": strings.Repeat("x", protocol.MaxAccountingUsageJSONBytes),
	}
	rawUsageJSON, err := json.Marshal(rawUsage)
	require.NoError(t, err)
	require.Greater(t, len(rawUsageJSON), protocol.MaxAccountingUsageJSONBytes)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		require.NoError(t, json.NewEncoder(w).Encode(map[string]any{
			"id": "response_1",
			"choices": []any{map[string]any{
				"index": 0, "finish_reason": "tool_calls",
				"message": map[string]any{
					"role": "assistant", "content": nil,
					"tool_calls": []any{map[string]any{
						"id": "call_1", "type": "function",
						"function": map[string]any{
							"name": knowledgeSearchFunction, "arguments": `{"query":"server"}`,
						},
					}},
				},
			}},
			"usage": rawUsage,
		}))
	}))
	defer server.Close()

	result, err := (MaaSClient{BaseURL: server.URL, HTTPClient: server.Client()}).Generate(
		context.Background(), maasGenerateTestRequest(t, "kqa_oversized_usage"),
	)

	require.NoError(t, err)
	require.Equal(t, map[string]any{
		"input_tokens": float64(17), "output_tokens": float64(3),
	}, result.Usage)
	encoded, err := json.Marshal(result.Usage)
	require.NoError(t, err)
	require.LessOrEqual(t, len(encoded), protocol.MaxAccountingUsageJSONBytes)
}
func maasGenerateTestRequest(t *testing.T, idempotencyKey string) InferenceRequest {
	t.Helper()
	return InferenceRequest{
		Model: "model-a", ReasonLevel: "high", Capability: "capability", IdempotencyKey: idempotencyKey,
		Mode: InferenceModeExploration, System: "system", Messages: []InferenceMessage{{Role: "user", Content: "explore"}},
		Profile: mustSupportedProfile(t, chatCompletionsProfileOpenAIRequiredToolsV1),
	}
}

func validChatCompletion(responseID, callID, name, arguments string) string {
	return fmt.Sprintf(
		`{"id":%q,"choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":null,"tool_calls":[{"id":%q,"type":"function","function":{"name":%q,"arguments":%q}}]}}],"usage":{"input_tokens":17,"output_tokens":3}}`,
		responseID, callID, name, arguments,
	)
}
func mustSupportedProfile(t *testing.T, id string) ChatCompletionsCompatibilityProfile {
	t.Helper()
	profile, ok := lookupSupportedChatCompletionsProfile(id)
	require.True(t, ok)
	return profile
}