package openai

import (
	"context"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"

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

// TestProviderSwitchEmitsUsagePerActualAttemptAndOneRequest drives the real
// provider-pool request runtime through a tunnel-to-normalized recovery. The
// replaced attempt and the committed attempt report distinct provider usage;
// both must be retained while the HTTP request terminal advances only once.
func TestProviderSwitchEmitsUsagePerActualAttemptAndOneRequest(t *testing.T) {
	const (
		edgeID     = "edge-provider-switch-usage"
		routeModel = "client-model"
	)
	initialFrames := bufferedTunnelFrames(
		&iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: http.StatusOK, Headers: map[string]string{"Content-Type": "application/json"}},
		&iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(`{"choices":[{"message":{"content":"discarded"}}],"usage":{"prompt_tokens":11,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":1},"completion_tokens_details":{"reasoning_tokens":2}}}`)},
		&iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true},
	)
	service := newScriptedPoolRunService(
		scriptedPoolAttempt{
			path: string(edgeservice.ProviderPoolPathTunnel), runID: "usage-attempt-a",
			provider: "provider-a", target: "served-a", frames: initialFrames,
		},
		scriptedPoolAttempt{
			path: string(edgeservice.ProviderPoolPathNormalized), runID: "usage-attempt-b",
			provider: "provider-b", target: "served-b",
			runEvents: bufferedRunEvents(
				&iop.RunEvent{Type: "delta", Delta: "accepted"},
				&iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 17, OutputTokens: 5, ReasoningTokens: 4, CachedInputTokens: 6}},
			),
		},
	)
	rawBody := []byte(`{"model":"client-model","messages":[{"role":"user","content":"hi"}]}`)
	srv, dc := buildStreamGatePoolChatFixture(t, service, rawBody, 3)
	srv.SetEdgeID(edgeID)
	dc.usage = &openAIUsageRecorder{
		request:  usageRequestLabels{edgeID: edgeID, routeModel: routeModel, endpoint: usageEndpointChatCompletions},
		attempts: make(map[string]usageDispatchBinding),
	}

	labelsA := usageAttemptLabels{
		usageRequestLabels: dc.usage.request, usageAttribution: config.UsageAttributionProvider,
		providerID: "provider-a", servedModel: "served-a", responseMode: responseModePassthrough,
	}
	labelsB := usageAttemptLabels{
		usageRequestLabels: dc.usage.request, usageAttribution: config.UsageAttributionProvider,
		providerID: "provider-b", servedModel: "served-b", responseMode: responseModeNormalized,
	}
	aInputBefore := requestTokenValue(t, labelsA, tokenTypeInput)
	aOutputBefore := requestTokenValue(t, labelsA, tokenTypeOutput)
	aReasoningBefore := requestTokenValue(t, labelsA, tokenTypeReasoning)
	aCachedBefore := requestTokenValue(t, labelsA, tokenTypeCachedInput)
	bInputBefore := requestTokenValue(t, labelsB, tokenTypeInput)
	bOutputBefore := requestTokenValue(t, labelsB, tokenTypeOutput)
	bReasoningBefore := requestTokenValue(t, labelsB, tokenTypeReasoning)
	bCachedBefore := requestTokenValue(t, labelsB, tokenTypeCachedInput)
	requestBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
		edgeID, "", "", "", routeModel, usageEndpointChatCompletions,
		responseModeNormalized, usageStatusSuccess, usageSourceProviderReported,
	))

	w, sink, runErr := runPoolChatStreamGate(t, srv, dc, newInjectedViolationRegistration(t, dc, "test.usage_switch", 1))
	if runErr != nil {
		t.Fatalf("runtime run: %v", runErr)
	}
	committed, success := sink.terminalStatus()
	if !committed || !success || !strings.Contains(w.body.String(), "accepted") {
		t.Fatalf("replacement response did not commit successfully: committed=%v success=%v body=%q", committed, success, w.body.String())
	}
	dc.finishUsageRequest(usageStatusSuccess, responseModeNormalized)

	for name, gotWant := range map[string][2]float64{
		"provider-a input":     {requestTokenValue(t, labelsA, tokenTypeInput) - aInputBefore, 11},
		"provider-a output":    {requestTokenValue(t, labelsA, tokenTypeOutput) - aOutputBefore, 3},
		"provider-a reasoning": {requestTokenValue(t, labelsA, tokenTypeReasoning) - aReasoningBefore, 2},
		"provider-a cached":    {requestTokenValue(t, labelsA, tokenTypeCachedInput) - aCachedBefore, 1},
		"provider-b input":     {requestTokenValue(t, labelsB, tokenTypeInput) - bInputBefore, 17},
		"provider-b output":    {requestTokenValue(t, labelsB, tokenTypeOutput) - bOutputBefore, 5},
		"provider-b reasoning": {requestTokenValue(t, labelsB, tokenTypeReasoning) - bReasoningBefore, 4},
		"provider-b cached":    {requestTokenValue(t, labelsB, tokenTypeCachedInput) - bCachedBefore, 6},
	} {
		if gotWant[0] != gotWant[1] {
			t.Fatalf("%s tokens = %v, want %v", name, gotWant[0], gotWant[1])
		}
	}
	requestAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
		edgeID, "", "", "", routeModel, usageEndpointChatCompletions,
		responseModeNormalized, usageStatusSuccess, usageSourceProviderReported,
	))
	if got := requestAfter - requestBefore; got != 1 {
		t.Fatalf("request terminal count = %v, want 1", got)
	}
}

func TestNonStreamingTunnelAttemptFinalizesBodyUsage(t *testing.T) {
	tests := []struct {
		name                 string
		endpoint             string
		status               int
		contentType          string
		body                 string
		endKind              iop.ProviderTunnelFrameKind
		wantInput            int
		wantOutput           int
		wantReasoning        int
		wantCached           int
		wantProviderReported bool
		wantHolderSet        bool
	}{
		{
			name:                 "chat completions success body",
			endpoint:             openAIRebuildEndpointChat,
			status:               http.StatusOK,
			contentType:          "application/json",
			body:                 `{"choices":[{"message":{"content":"hello"}}],"usage":{"prompt_tokens":9,"completion_tokens":4,"prompt_tokens_details":{"cached_tokens":1},"completion_tokens_details":{"reasoning_tokens":2}}}`,
			endKind:              iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
			wantInput:            9,
			wantOutput:           4,
			wantReasoning:        2,
			wantCached:           1,
			wantProviderReported: true,
			wantHolderSet:        true,
		},
		{
			name:                 "responses API success body",
			endpoint:             openAIRebuildEndpointResponses,
			status:               http.StatusOK,
			contentType:          "application/json",
			body:                 `{"output_text":"hello","usage":{"input_tokens":14,"output_tokens":7,"input_tokens_details":{"cached_tokens":3},"output_tokens_details":{"reasoning_tokens":5}}}`,
			endKind:              iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
			wantInput:            14,
			wantOutput:           7,
			wantReasoning:        5,
			wantCached:           3,
			wantProviderReported: true,
			wantHolderSet:        true,
		},
		{
			name:                 "chat completions provider error boundary",
			endpoint:             openAIRebuildEndpointChat,
			status:               http.StatusOK,
			contentType:          "application/json",
			body:                 `{"error":{"message":"internal error"},"usage":{"prompt_tokens":5,"completion_tokens":2,"prompt_tokens_details":{"cached_tokens":0},"completion_tokens_details":{"reasoning_tokens":1}}}`,
			endKind:              iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
			wantInput:            5,
			wantOutput:           2,
			wantReasoning:        1,
			wantCached:           0,
			wantProviderReported: true,
			wantHolderSet:        false,
		},
		{
			name:                 "responses API provider error boundary",
			endpoint:             openAIRebuildEndpointResponses,
			status:               http.StatusOK,
			contentType:          "application/json",
			body:                 `{"error":{"message":"internal error"},"usage":{"input_tokens":8,"output_tokens":3,"input_tokens_details":{"cached_tokens":2},"output_tokens_details":{"reasoning_tokens":0}}}`,
			endKind:              iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
			wantInput:            8,
			wantOutput:           3,
			wantReasoning:        0,
			wantCached:           2,
			wantProviderReported: true,
			wantHolderSet:        false,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			frames := bufferedTunnelFrames(
				&iop.ProviderTunnelFrame{
					Kind:       iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
					StatusCode: int32(tc.status),
					Headers:    map[string]string{"Content-Type": tc.contentType},
				},
				&iop.ProviderTunnelFrame{
					Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
					Body: []byte(tc.body),
				},
				&iop.ProviderTunnelFrame{
					Kind: tc.endKind,
					End:  tc.endKind == iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
				},
			)

			stream := edgeservice.ProviderTunnelStream{Frames: frames}
			assembler := &providerChatAssembler{streaming: false}
			rewriter := newProviderModelRewriter(false, "request-model")
			state := &openAITunnelCodecState{}
			state.bindEndpoint(tc.endpoint)
			src := newOpenAITunnelEndpointEventSource(stream, 5*time.Second, rewriter, assembler, tc.endpoint, state)

			holder := &openAIStreamGateUsageHolder{}
			attempt := &openAIAttemptUsage{}
			tracking := &openAIStreamGateUsageTrackingTunnelSource{
				openAITunnelEventSource: src,
				usage:                   holder,
				attempt:                 attempt,
			}

			ctx := context.Background()
			for {
				ev, err := tracking.NextEvent(ctx)
				if err != nil || ev.Kind() == streamgate.EventKindTerminal || ev.Kind() == streamgate.EventKindProviderError {
					break
				}
			}

			attObs := attempt.snapshot()
			if attObs.providerReported != tc.wantProviderReported {
				t.Fatalf("attempt providerReported = %v, want %v", attObs.providerReported, tc.wantProviderReported)
			}
			if attObs.inputTokens != tc.wantInput || attObs.outputTokens != tc.wantOutput ||
				attObs.reasoningTokens != tc.wantReasoning || attObs.cachedInputTokens != tc.wantCached {
				t.Fatalf("attempt usage = %+v, want input=%d output=%d reasoning=%d cached=%d",
					attObs, tc.wantInput, tc.wantOutput, tc.wantReasoning, tc.wantCached)
			}

			hObs := holder.get()
			if tc.wantHolderSet {
				if !hObs.providerReported {
					t.Fatalf("holder providerReported = false, want true")
				}
				if hObs.inputTokens != tc.wantInput || hObs.outputTokens != tc.wantOutput ||
					hObs.reasoningTokens != tc.wantReasoning || hObs.cachedInputTokens != tc.wantCached {
					t.Fatalf("holder usage = %+v, want input=%d output=%d reasoning=%d cached=%d",
						hObs, tc.wantInput, tc.wantOutput, tc.wantReasoning, tc.wantCached)
				}
			} else {
				if hObs.providerReported {
					t.Fatalf("holder providerReported = true, want false for non-terminal error")
				}
			}
		})
	}
}