package openai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"testing"
	"time"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

type mockTunnel struct {
	frames chan *iop.ProviderTunnelFrame
	closed bool
}

func (t *mockTunnel) Dispatch() edgeservice.RunDispatch { return edgeservice.RunDispatch{} }
func (t *mockTunnel) Stream() edgeservice.ProviderTunnelStream {
	return edgeservice.ProviderTunnelStream{Frames: t.frames}
}
func (t *mockTunnel) Close()                       { t.closed = true }
func (t *mockTunnel) WaitTimeout() time.Duration   { return 0 }
func (t *mockTunnel) SetHeaders(map[string]string) {}

type mockService struct {
	submit func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error)
}

func (m *mockService) SubmitProviderPool(ctx context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
	return m.submit(ctx, req)
}

func validDispatch() edgeservice.SingleRequestStageDispatchBinding {
	return edgeservice.SingleRequestStageDispatchBinding{
		Managed:              true,
		ModelGroupKey:        "plan-model",
		RouteID:              "route-1",
		ProfileID:            "profile-1",
		CredentialSlotRef:    "slot-1",
		CredentialRevision:   1,
		RouteRevision:        1,
		PrincipalRef:         "principal-1",
		ProjectionGeneration: 1,
		ProviderID:           "gemini",
		UpstreamModel:        "gemini-3.6-flash",
		TimeoutSec:           60,
		MaxQueue:             10,
		QueueTimeoutMS:       5000,
		CandidatePredicate:   func(c edgeservice.ProviderPoolCandidate) bool { return c.ProviderID == "gemini" },
	}
}

func validStageBinding() edgeservice.SingleRequestStageBinding {
	d := validDispatch()
	return edgeservice.SingleRequestStageBinding{
		Model:    "plan-model",
		Options:  map[string]any{"reasoning_effort": "high", "temperature": 0.2, "tools": []any{"should_be_ignored"}, "model": "override_ignored"},
		Dispatch: &d,
	}
}

func validLimits() edgeservice.SingleRequestLimits {
	return edgeservice.SingleRequestLimits{
		WallClockMS:       1000,
		StageTimeoutMS:    500,
		MaxToolIterations: 1,
		MaxOutputBytes:    4096,
	}
}

func successBody(content string) []byte {
	b, _ := json.Marshal(map[string]any{
		"id":      "id",
		"object":  "chat.completion",
		"created": 1,
		"model":   "gemini-3.6-flash",
		"choices": []any{
			map[string]any{
				"index":         0,
				"finish_reason": "stop",
				"message": map[string]any{
					"role":              "assistant",
					"content":           content,
					"reasoning_content": "provider-private-reasoning",
				},
			},
		},
		"usage": map[string]any{
			"prompt_tokens":     11,
			"completion_tokens": 2,
			"total_tokens":      13,
			"prompt_tokens_details": map[string]any{
				"cached_tokens": 3,
			},
		},
	})
	return b
}

func successBodyWithThoughtSignature(content string) []byte {
	var body map[string]any
	_ = json.Unmarshal(successBody(content), &body)
	choices := body["choices"].([]any)
	message := choices[0].(map[string]any)["message"].(map[string]any)
	message["extra_content"] = map[string]any{"google": map[string]any{"thought_signature": "provider-private-thought-signature"}}
	b, _ := json.Marshal(body)
	return b
}

func matchingDispatch() edgeservice.RunDispatch {
	return edgeservice.RunDispatch{
		ModelGroupKey:      "plan-model",
		ProviderID:         "gemini",
		Target:             "gemini-3.6-flash",
		ProfileID:          "profile-1",
		ProfileDriver:      string(config.ProtocolDriverOpenAIChat),
		CredentialSlotRef:  "slot-1",
		CredentialRevision: 1,
		ExecutionPath:      string(edgeservice.ProviderPoolPathTunnel),
	}
}

func prepareSingleRequestTestTunnel(req edgeservice.ProviderPoolDispatchRequest, candidate edgeservice.ProviderPoolCandidate) (edgeservice.SubmitProviderTunnelRequest, error) {
	if req.PrepareProtocolTunnel != nil {
		return req.PrepareProtocolTunnel(req.Tunnel, candidate)
	}
	if req.PrepareTunnel != nil {
		return req.PrepareTunnel(req.Tunnel)
	}
	return req.Tunnel, nil
}

func framesFor(body []byte) chan *iop.ProviderTunnelFrame {
	c := make(chan *iop.ProviderTunnelFrame, 3)
	c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
	c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: body}
	c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END}
	close(c)
	return c
}

func providerRequest() singleRequestProviderStageRequest {
	return singleRequestProviderStageRequest{
		StageBinding:     validStageBinding(),
		Limits:           validLimits(),
		Messages:         []chatMessage{{Role: "user", Content: "immutable task"}},
		NodeRef:          "node",
		SessionID:        "session",
		UsageAttribution: "principal",
	}
}

func TestSingleRequestProviderStageUsesFrozenOptionsAndDispatch(t *testing.T) {
	var captured edgeservice.ProviderPoolDispatchRequest
	tunnel := &mockTunnel{frames: framesFor(successBody("ok"))}
	dispatch := matchingDispatch()
	stage := newSingleRequestProviderStage(&mockService{submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
		captured = req
		return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: dispatch}, nil
	}})

	request := providerRequest()
	request.StageBinding.Options["response_format"] = map[string]any{"type": "caller_override_ignored"}
	response, err := stage.submit(context.Background(), request)
	if err != nil {
		t.Fatal(err)
	}
	if response.Output != "ok" || !tunnel.closed {
		t.Fatalf("response=%+v closed=%v", response, tunnel.closed)
	}

	// Verify exact dispatch parameters captured in ProviderPoolDispatchRequest
	if captured.Run.NodeRef != "node" || captured.Run.ModelGroupKey != "plan-model" || captured.Run.ProviderID != "gemini" || captured.Run.UsageAttribution != "principal" || captured.Run.SessionID != "session" || captured.Run.TimeoutSec != 60 || captured.Run.MaxQueue != 10 || captured.Run.QueueTimeoutMS != 5000 || !captured.Run.ProviderPool || captured.Run.RunID != "" || captured.Run.Adapter != "" || captured.Run.Target != "" || captured.Run.Prompt != "" || captured.Run.Input != nil || captured.Run.Background != false || captured.Run.Metadata != nil || captured.Run.EstimatedInputTokens != 0 || captured.Run.ContextClass != "" || captured.Run.ResponseStallTimeoutMS != 0 {
		t.Fatalf("unexpected Run dispatch: %+v", captured.Run)
	}

	if captured.Tunnel.NodeRef != "node" || captured.Tunnel.ModelGroupKey != "plan-model" || captured.Tunnel.ProviderID != "gemini" || captured.Tunnel.UsageAttribution != "principal" || captured.Tunnel.Adapter != "openai_compat" || captured.Tunnel.Target != "gemini-3.6-flash" || captured.Tunnel.SessionID != "session" || captured.Tunnel.Method != http.MethodPost || captured.Tunnel.Path != "/v1/chat/completions" || captured.Tunnel.Operation != "" || captured.Tunnel.Stream || captured.Tunnel.TimeoutSec != 60 || captured.Tunnel.MaxQueue != 10 || captured.Tunnel.QueueTimeoutMS != 5000 || !captured.Tunnel.ProviderPool || captured.Tunnel.RunID != "" || captured.Tunnel.Headers != nil || captured.Tunnel.Body != nil || captured.Tunnel.Metadata != nil || captured.Tunnel.EstimatedInputTokens != 0 || captured.Tunnel.ContextClass != "" || captured.Tunnel.ResponseStallTimeoutMS != 0 {
		t.Fatalf("unexpected Tunnel request: %+v", captured.Tunnel)
	}

	cred := captured.Tunnel.CredentialBinding
	if cred == nil || cred.PrincipalRef != "principal-1" || cred.CredentialSlotRef != "slot-1" || cred.RouteID != "route-1" || cred.ProfileID != "profile-1" || cred.CredentialRevision != 1 || cred.RouteRevision != 1 || cred.ProjectionGeneration != 1 {
		t.Fatalf("unexpected CredentialBinding: %+v", cred)
	}

	body, err := captured.Tunnel.BuildBody("gemini-3.6-flash")
	if err != nil {
		t.Fatal(err)
	}
	var decoded map[string]any
	if err := json.Unmarshal(body, &decoded); err != nil {
		t.Fatal(err)
	}
	if decoded["reasoning_effort"] != "high" || decoded["model"] != "gemini-3.6-flash" || decoded["tools"] != nil || decoded["temperature"] != 0.2 || decoded["response_format"] != nil {
		t.Fatalf("unexpected frozen body: %s", body)
	}

	profile, err := config.ResolveProtocolProfile("gemini", "", config.BuiltInProtocolProfileCatalog())
	if err != nil {
		t.Fatal(err)
	}
	accepted := edgeservice.ProviderPoolCandidate{ProviderID: "gemini", ExecutionPath: string(edgeservice.ProviderPoolPathTunnel), ProtocolProfile: &profile}
	rejected := accepted
	rejected.ProviderID = "other"
	// Check candidate predicate against the same concrete profile facts supplied
	// by the real provider-pool service.
	if captured.AcceptCandidate == nil || !captured.AcceptCandidate(accepted) || captured.AcceptCandidate(rejected) {
		t.Fatalf("AcceptCandidate predicate missing or returned unexpected result")
	}
}

func TestSingleRequestProviderNormalizationSelectsResponsesAndFallsBackEffort(t *testing.T) {
	catalog := config.BuiltInProtocolProfileCatalog()
	openAI := catalog["openai"]
	responses := openAI.Normalization.Effort[string(config.OperationResponses)]
	delete(responses.Levels, "max")
	openAI.Normalization.Effort[string(config.OperationResponses)] = responses
	catalog["openai-without-max"] = openAI
	profile, err := config.ResolveProtocolProfile("openai-without-max", "", catalog)
	if err != nil {
		t.Fatal(err)
	}

	tools := []any{singleRequestWorkToolSchema("workspace_read", map[string]any{
		"type": "object", "properties": map[string]any{}, "additionalProperties": false,
	})}
	requirements := singleRequestProviderRequirements(map[string]any{"reasoning_effort": "max"}, tools, singleRequestPlanResponseFormat())
	plan, err := selectProviderOperation(profile, config.OperationChatCompletions, requirements)
	if err != nil {
		t.Fatal(err)
	}
	if plan.Operation != config.OperationResponses || plan.Effort != "xhigh" || plan.EffortWire != config.ProtocolEffortWireOpenAIResponses {
		t.Fatalf("plan=%+v, want Responses max->xhigh", plan)
	}

	chatBody := func(target string) ([]byte, error) {
		return json.Marshal(map[string]any{
			"model": target, "stream": false, "reasoning_effort": "max", "max_completion_tokens": 256,
			"messages": []any{
				map[string]any{"role": "system", "content": "system"},
				map[string]any{"role": "user", "content": "task"},
			},
			"tools": tools, "tool_choice": "auto", "response_format": singleRequestPlanResponseFormat(),
		})
	}
	prepared, err := singleRequestProviderTunnelPreparer(requirements, chatBody)(
		edgeservice.SubmitProviderTunnelRequest{Path: "/v1/chat/completions"},
		edgeservice.ProviderPoolCandidate{ExecutionPath: string(edgeservice.ProviderPoolPathTunnel), ProtocolProfile: &profile},
	)
	if err != nil {
		t.Fatal(err)
	}
	if prepared.Operation != string(config.OperationResponses) || prepared.Path != "/v1/responses" {
		t.Fatalf("prepared operation/path=%q/%q", prepared.Operation, prepared.Path)
	}
	body, err := prepared.BuildBody("gpt-model")
	if err != nil {
		t.Fatal(err)
	}
	var decoded map[string]any
	if err := json.Unmarshal(body, &decoded); err != nil {
		t.Fatal(err)
	}
	reasoning, _ := decoded["reasoning"].(map[string]any)
	text, _ := decoded["text"].(map[string]any)
	format, _ := text["format"].(map[string]any)
	if decoded["model"] != "gpt-model" || decoded["max_output_tokens"] != float64(256) || reasoning["effort"] != "xhigh" || format["type"] != "json_schema" || len(anySlice(decoded["tools"])) != 1 {
		t.Fatalf("normalized Responses body=%s", body)
	}
}

func TestSingleRequestProviderNormalizationConvertsResponsesResult(t *testing.T) {
	raw := []byte(`{"id":"resp-1","model":"gpt-model","status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{\"goal\":\"g\",\"steps\":[\"a\",\"b\"],\"verification\":[\"v\"]}"}]}]}`)
	body, err := normalizeSingleRequestProviderResponse(raw, edgeservice.RunDispatch{ProfileOperation: string(config.OperationResponses)})
	if err != nil {
		t.Fatal(err)
	}
	response, err := decodeSingleRequestChatResponse(body, edgeservice.RunDispatch{ProfileOperation: string(config.OperationResponses)})
	if err != nil {
		t.Fatal(err)
	}
	if response.Output != `{"goal":"g","steps":["a","b"],"verification":["v"]}` {
		t.Fatalf("output=%q", response.Output)
	}
}

func TestSingleRequestProviderNormalizationCanonicalizesChatResult(t *testing.T) {
	raw := []byte(`{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-model","service_tier":"default","system_fingerprint":"fp","timings":{"prompt_n":2,"predicted_n":1},"choices":[{"index":0,"finish_reason":"stop","logprobs":null,"message":{"role":"assistant","content":"{\"goal\":\"g\",\"steps\":[\"a\",\"b\"],\"verification\":[\"v\"]}","refusal":null,"annotations":[]}}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}`)
	body, err := normalizeSingleRequestProviderResponse(raw, edgeservice.RunDispatch{ProfileOperation: string(config.OperationChatCompletions)})
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Contains(body, []byte("service_tier")) || bytes.Contains(body, []byte("system_fingerprint")) || bytes.Contains(body, []byte("timings")) || bytes.Contains(body, []byte("annotations")) || bytes.Contains(body, []byte("refusal")) {
		t.Fatalf("provider bookkeeping leaked into canonical stage response: %s", body)
	}
	response, err := decodeSingleRequestChatResponse(body, edgeservice.RunDispatch{ProfileOperation: string(config.OperationChatCompletions)})
	if err != nil {
		t.Fatal(err)
	}
	if response.Output != `{"goal":"g","steps":["a","b"],"verification":["v"]}` {
		t.Fatalf("output=%q", response.Output)
	}
}

func TestSingleRequestProviderNormalizationRejectsChatRefusal(t *testing.T) {
	raw := []byte(`{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-model","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":null,"refusal":"cannot comply"}}]}`)
	if _, err := normalizeSingleRequestProviderResponse(raw, edgeservice.RunDispatch{ProfileOperation: string(config.OperationChatCompletions)}); !errors.Is(err, errProviderStageMalformed) {
		t.Fatalf("err=%v, want malformed", err)
	}
}

func TestSingleRequestProviderStageAcceptsDefaultSelectorProviderChoice(t *testing.T) {
	request := providerRequest()
	request.StageBinding.Dispatch.ProviderID = ""
	tunnel := &mockTunnel{frames: framesFor(successBody("ok"))}
	dispatch := matchingDispatch()
	stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
		return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: dispatch}, nil
	}})
	response, err := stage.submit(context.Background(), request)
	if err != nil {
		t.Fatal(err)
	}
	if response.Output != "ok" || response.Dispatch.ProviderID != "gemini" {
		t.Fatalf("response=%+v", response)
	}
}

func TestSingleRequestProviderStageRejectsResponseEnvelope(t *testing.T) {
	tests := []struct {
		name string
		raw  string
	}{
		{
			name: "zero-choices",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[]}`,
		},
		{
			name: "multiple-choices",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"a"}},{"index":1,"finish_reason":"stop","message":{"role":"assistant","content":"b"}}]}`,
		},
		{
			name: "nonzero-index",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":1,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "user-role",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"user","content":"ok"}}]}`,
		},
		{
			name: "empty-role",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"","content":"ok"}}]}`,
		},
		{
			name: "finish-reason-length",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"length","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "finish-reason-tool-calls",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"tool_calls","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "finish-reason-empty",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "unexpected-tool-calls",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","tool_calls":[{"id":"t1","type":"function"}]}}]}`,
		},
		{
			name: "malformed-json",
			raw:  `{"id":"1",`,
		},
		{
			name: "unknown-top-level-field",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}],"extra":1}`,
		},
		{
			name: "unknown-usage-field",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2,"provider_private":1}}`,
		},
		{
			name: "negative-usage",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}],"usage":{"prompt_tokens":-1,"completion_tokens":1,"total_tokens":0}}`,
		},
		{
			name: "trailing-json",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]} {"extra":1}`,
		},
		{
			name: "non-string-content-object",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":{"plan":"A"}}}]}`,
		},
		{
			name: "non-string-content-array",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":["a"]}}]}`,
		},
		{
			name: "null-content",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":null}}]}`,
		},
		{
			name: "unknown-message-field",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","unknown":1}}]}`,
		},
		{
			name: "non-string-reasoning-content-object",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","reasoning_content":{"private":true}}}]}`,
		},
		{
			name: "non-string-reasoning-content-array",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","reasoning_content":["private"]}}]}`,
		},
		{
			name: "null-extra-content",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","extra_content":null}}]}`,
		},
		{
			name: "empty-thought-signature",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","extra_content":{"google":{"thought_signature":""}}}}]}`,
		},
		{
			name: "non-string-thought-signature",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","extra_content":{"google":{"thought_signature":7}}}}]}`,
		},
		{
			name: "unknown-Google-extra-content-field",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","extra_content":{"google":{"thought_signature":"sig","unknown":true}}}}]}`,
		},
		{
			name: "unknown-extra-content-field",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","extra_content":{"google":{"thought_signature":"sig"},"unknown":true}}}]}`,
		},
		{
			name: "duplicate-thought-signature",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","extra_content":{"google":{"thought_signature":"sig","thought_signature":"sig2"}}}}]}`,
		},
		{
			name: "duplicate-response-key",
			raw:  `{"id":"1","id":"2","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "duplicate-choice-key",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "duplicate-message-key",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "case-variant-response-key",
			raw:  `{"Id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "case-folded-duplicate-response-key",
			raw:  `{"id":"1","Id":"2","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "case-variant-choice-key",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"Index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "case-folded-duplicate-choice-key",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"Index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "case-variant-message-key",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"Role":"assistant","content":"ok"}}]}`,
		},
		{
			name: "case-folded-duplicate-message-key",
			raw:  `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok","Content":"ok"}}]}`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tunnel := &mockTunnel{frames: framesFor([]byte(tt.raw))}
			d := matchingDispatch()
			stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
			}})
			_, err := stage.submit(context.Background(), providerRequest())
			if !errors.Is(err, errProviderStageGeneric) {
				t.Fatalf("expected errProviderStageGeneric, got %v", err)
			}
			if !tunnel.closed {
				t.Fatalf("expected tunnel to be closed")
			}
		})
	}

	t.Run("valid-assistant-stop-accepts", func(t *testing.T) {
		validJSON := `{"id":"1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"valid output"}}],"usage":{"prompt_tokens":11,"completion_tokens":2,"total_tokens":13,"prompt_tokens_details":{"cached_tokens":3},"completion_tokens_details":{"reasoning_tokens":1}}} `
		tunnel := &mockTunnel{frames: framesFor([]byte(validJSON))}
		d := matchingDispatch()
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
		}})
		resp, err := stage.submit(context.Background(), providerRequest())
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if resp.Output != "valid output" {
			t.Fatalf("output=%q, expected %q", resp.Output, "valid output")
		}
		if !tunnel.closed {
			t.Fatalf("expected tunnel to be closed")
		}
	})

	t.Run("valid-Gemini-thought-signature-is-discarded", func(t *testing.T) {
		tunnel := &mockTunnel{frames: framesFor(successBodyWithThoughtSignature("visible output"))}
		d := matchingDispatch()
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
		}})
		resp, err := stage.submit(context.Background(), providerRequest())
		if err != nil {
			t.Fatal(err)
		}
		if resp.Output != "visible output" || strings.Contains(resp.Output, "provider-private") {
			t.Fatalf("private thought signature leaked: %+v", resp)
		}
	})
}

func TestSingleRequestProviderStageRejectsFrameFailures(t *testing.T) {
	for name, frames := range map[string]chan *iop.ProviderTunnelFrame{
		"usage-frame": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 2)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE}
			close(c)
			return c
		}(),
		"body-before-start": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 1)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY}
			close(c)
			return c
		}(),
		"duplicate-start": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 2)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			close(c)
			return c
		}(),
		"body-after-terminal": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 4)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: successBody("a")}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: successBody("b")}
			close(c)
			return c
		}(),
		"duplicate-terminal": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 3)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END}
			close(c)
			return c
		}(),
		"non-2xx-status-199": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 1)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 199}
			close(c)
			return c
		}(),
		"non-2xx-status-400": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 1)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 400}
			close(c)
			return c
		}(),
		"non-2xx-status-500": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 1)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 500}
			close(c)
			return c
		}(),
		"provider-error": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 2)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR}
			close(c)
			return c
		}(),
		"missing-terminal": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 1)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			close(c)
			return c
		}(),
		"nil-frame": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 2)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			c <- nil
			close(c)
			return c
		}(),
		"unknown-frame-kind": func() chan *iop.ProviderTunnelFrame {
			c := make(chan *iop.ProviderTunnelFrame, 2)
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			c <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind(999)}
			close(c)
			return c
		}(),
	} {
		t.Run(name, func(t *testing.T) {
			tunnel := &mockTunnel{frames: frames}
			d := matchingDispatch()
			stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
			}})
			_, err := stage.submit(context.Background(), providerRequest())
			if !errors.Is(err, errProviderStageGeneric) || !tunnel.closed {
				t.Fatalf("err=%v closed=%v", err, tunnel.closed)
			}
		})
	}

	t.Run("output-limit-exact-boundary", func(t *testing.T) {
		body := successBody("ok")
		req := providerRequest()
		req.Limits.MaxOutputBytes = len(body)
		tunnel := &mockTunnel{frames: framesFor(body)}
		d := matchingDispatch()
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
		}})
		resp, err := stage.submit(context.Background(), req)
		if err != nil {
			t.Fatalf("unexpected error on exact boundary: %v", err)
		}
		if resp.Output != "ok" || !tunnel.closed {
			t.Fatalf("resp=%+v closed=%v", resp, tunnel.closed)
		}
	})

	t.Run("output-limit-exceeded-boundary", func(t *testing.T) {
		body := successBody("ok")
		req := providerRequest()
		req.Limits.MaxOutputBytes = len(body) - 1
		tunnel := &mockTunnel{frames: framesFor(body)}
		d := matchingDispatch()
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
		}})
		_, err := stage.submit(context.Background(), req)
		if !errors.Is(err, errProviderStageGeneric) || !tunnel.closed {
			t.Fatalf("expected limit failure err=%v closed=%v", err, tunnel.closed)
		}
	})
}

func TestSingleRequestProviderStageRejectsMismatchLimitAndContext(t *testing.T) {
	dispatchFields := []struct {
		name   string
		mutate func(*edgeservice.RunDispatch)
	}{
		{"path-mismatch", func(d *edgeservice.RunDispatch) { d.ExecutionPath = "normalized" }},
		{"driver-mismatch", func(d *edgeservice.RunDispatch) { d.ProfileDriver = "other" }},
		{"modelgroup-mismatch", func(d *edgeservice.RunDispatch) { d.ModelGroupKey = "wrong" }},
		{"provider-mismatch", func(d *edgeservice.RunDispatch) { d.ProviderID = "wrong" }},
		{"target-mismatch", func(d *edgeservice.RunDispatch) { d.Target = "wrong" }},
		{"profile-mismatch", func(d *edgeservice.RunDispatch) { d.ProfileID = "wrong" }},
		{"slot-mismatch", func(d *edgeservice.RunDispatch) { d.CredentialSlotRef = "wrong" }},
		{"revision-mismatch", func(d *edgeservice.RunDispatch) { d.CredentialRevision = 999 }},
	}

	for _, tt := range dispatchFields {
		t.Run(tt.name, func(t *testing.T) {
			d := matchingDispatch()
			tt.mutate(&d)
			tunnel := &mockTunnel{frames: framesFor(successBody("ok"))}
			stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
				return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
			}})
			_, err := stage.submit(context.Background(), providerRequest())
			if !errors.Is(err, errProviderStageGeneric) || !tunnel.closed {
				t.Fatalf("err=%v closed=%v", err, tunnel.closed)
			}
			disposition, ok := singleRequestTerminalDisposition(err)
			if !ok || disposition.Kind != edgeservice.SingleRequestTerminalError || disposition.ErrorClass != edgeservice.SingleRequestTerminalErrorValidation {
				t.Fatalf("dispatch mismatch disposition=%+v ok=%v", disposition, ok)
			}
		})
	}

	t.Run("result-path-not-tunnel", func(t *testing.T) {
		d := matchingDispatch()
		tunnel := &mockTunnel{frames: framesFor(successBody("ok"))}
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathNormalized, Tunnel: tunnel, DispatchInfo: d}, nil
		}})
		_, err := stage.submit(context.Background(), providerRequest())
		if !errors.Is(err, errProviderStageGeneric) || !tunnel.closed {
			t.Fatalf("err=%v closed=%v", err, tunnel.closed)
		}
	})

	t.Run("acquired-tunnel-deadline-timeout", func(t *testing.T) {
		blockedFrames := make(chan *iop.ProviderTunnelFrame)
		tunnel := &mockTunnel{frames: blockedFrames}
		d := matchingDispatch()
		req := providerRequest()
		req.Limits.StageTimeoutMS = 10
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: tunnel, DispatchInfo: d}, nil
		}})
		_, err := stage.submit(context.Background(), req)
		if !errors.Is(err, errProviderStageGeneric) || !tunnel.closed {
			t.Fatalf("expected timeout err=%v closed=%v", err, tunnel.closed)
		}
	})

	t.Run("service-returns-error", func(t *testing.T) {
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return nil, errors.New("service generic failure")
		}})
		_, err := stage.submit(context.Background(), providerRequest())
		if !errors.Is(err, errProviderStageGeneric) {
			t.Fatalf("expected errProviderStageGeneric, got %v", err)
		}
		disposition, ok := singleRequestTerminalDisposition(err)
		if !ok || disposition.Kind != edgeservice.SingleRequestTerminalError || disposition.ErrorClass != edgeservice.SingleRequestTerminalErrorProvider {
			t.Fatalf("provider failure disposition=%+v ok=%v", disposition, ok)
		}
	})

	t.Run("service-returns-nil-result", func(t *testing.T) {
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return nil, nil
		}})
		_, err := stage.submit(context.Background(), providerRequest())
		if !errors.Is(err, errProviderStageGeneric) {
			t.Fatalf("expected errProviderStageGeneric, got %v", err)
		}
	})

	t.Run("service-returns-nil-tunnel", func(t *testing.T) {
		d := matchingDispatch()
		stage := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: nil, DispatchInfo: d}, nil
		}})
		_, err := stage.submit(context.Background(), providerRequest())
		if !errors.Is(err, errProviderStageGeneric) {
			t.Fatalf("expected errProviderStageGeneric, got %v", err)
		}
	})

	t.Run("missing-binding-nil-stage", func(t *testing.T) {
		var stage *singleRequestProviderStage
		_, err := stage.submit(context.Background(), providerRequest())
		if !errors.Is(err, errProviderStageMissingBinding) {
			t.Fatalf("expected errProviderStageMissingBinding, got %v", err)
		}
	})

	t.Run("missing-binding-nil-service", func(t *testing.T) {
		stage := newSingleRequestProviderStage(nil)
		_, err := stage.submit(context.Background(), providerRequest())
		if !errors.Is(err, errProviderStageMissingBinding) {
			t.Fatalf("expected errProviderStageMissingBinding, got %v", err)
		}
	})

	t.Run("missing-binding-nil-dispatch", func(t *testing.T) {
		req := providerRequest()
		req.StageBinding.Dispatch = nil
		stage := newSingleRequestProviderStage(&mockService{})
		_, err := stage.submit(context.Background(), req)
		if !errors.Is(err, errProviderStageMissingBinding) {
			t.Fatalf("expected errProviderStageMissingBinding, got %v", err)
		}
	})

	t.Run("missing-binding-zero-timeout", func(t *testing.T) {
		req := providerRequest()
		req.Limits.StageTimeoutMS = 0
		stage := newSingleRequestProviderStage(&mockService{})
		_, err := stage.submit(context.Background(), req)
		if !errors.Is(err, errProviderStageMissingBinding) {
			t.Fatalf("expected errProviderStageMissingBinding, got %v", err)
		}
	})

	t.Run("missing-binding-zero-max-bytes", func(t *testing.T) {
		req := providerRequest()
		req.Limits.MaxOutputBytes = 0
		stage := newSingleRequestProviderStage(&mockService{})
		_, err := stage.submit(context.Background(), req)
		if !errors.Is(err, errProviderStageMissingBinding) {
			t.Fatalf("expected errProviderStageMissingBinding, got %v", err)
		}
	})

	t.Run("missing-input-empty-messages", func(t *testing.T) {
		req := providerRequest()
		req.Messages = nil
		stage := newSingleRequestProviderStage(&mockService{})
		_, err := stage.submit(context.Background(), req)
		if !errors.Is(err, errProviderStageMissingInput) {
			t.Fatalf("expected errProviderStageMissingInput, got %v", err)
		}
	})

	t.Run("cancelled-context", func(t *testing.T) {
		ctx, cancel := context.WithCancel(context.Background())
		cancel()
		stage := newSingleRequestProviderStage(&mockService{submit: func(ctx context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			<-ctx.Done()
			return nil, ctx.Err()
		}})
		_, err := stage.submit(ctx, providerRequest())
		if !errors.Is(err, errProviderStageGeneric) {
			t.Fatalf("expected errProviderStageGeneric, got %v", err)
		}
	})
}
