package openai

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

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

	edgeservice "iop/apps/edge/internal/service"
)

type singleRequestAnthropicSSEEvent struct {
	Name string
	Data map[string]any
}

func parseSingleRequestAnthropicSSE(t *testing.T, wire string) []singleRequestAnthropicSSEEvent {
	t.Helper()
	var events []singleRequestAnthropicSSEEvent
	for _, frame := range strings.Split(strings.TrimSpace(wire), "\n\n") {
		var event singleRequestAnthropicSSEEvent
		for _, line := range strings.Split(frame, "\n") {
			switch {
			case strings.HasPrefix(line, "event: "):
				event.Name = strings.TrimPrefix(line, "event: ")
			case strings.HasPrefix(line, "data: "):
				if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &event.Data); err != nil {
					t.Fatalf("decode SSE data %q: %v", line, err)
				}
			}
		}
		if event.Name == "" || event.Data == nil {
			t.Fatalf("malformed SSE frame: %q", frame)
		}
		events = append(events, event)
	}
	return events
}

func countSingleRequestAnthropicEvents(events []singleRequestAnthropicSSEEvent, name string) int {
	count := 0
	for _, event := range events {
		if event.Name == name {
			count++
		}
	}
	return count
}

func singleRequestAnthropicDeltaTexts(events []singleRequestAnthropicSSEEvent) []string {
	var texts []string
	for _, event := range events {
		if event.Name != "content_block_delta" {
			continue
		}
		delta, _ := event.Data["delta"].(map[string]any)
		if text, ok := delta["text"].(string); ok {
			texts = append(texts, text)
		}
	}
	return texts
}

func TestSingleRequestAnthropicStreamOneEnvelopeOneTerminal(t *testing.T) {
	w := httptest.NewRecorder()
	stream, err := newSingleRequestAnthropicStream(w, "req_public", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Start(); err != nil {
		t.Fatal(err)
	}
	progress := []edgeservice.SingleRequestProgress{
		{Stage: edgeservice.SingleRequestStateAccepted, Message: "PRIVATE_ACCEPTED"},
		{Stage: edgeservice.SingleRequestStatePlanning, Message: "PRIVATE_PLAN"},
		{Stage: edgeservice.SingleRequestStateInternalTool, Message: "PRIVATE_TOOL"},
		{Stage: edgeservice.SingleRequestStatePlanning, Message: "PRIVATE_PLAN_AGAIN"},
		{Stage: edgeservice.SingleRequestStateWorking, Message: "PRIVATE_WORK"},
		{Stage: edgeservice.SingleRequestStateReviewing, Message: "PRIVATE_REVIEW"},
	}
	for _, item := range progress {
		if err := stream.Progress(item); err != nil {
			t.Fatal(err)
		}
	}
	if err := stream.Final(edgeservice.SingleRequestResult{Output: "safe final result"}); err != nil {
		t.Fatal(err)
	}

	wireAtTerminal := w.Body.String()
	if err := stream.Ping(); err != nil {
		t.Fatalf("post-terminal ping returned established success as error: %v", err)
	}
	if err := stream.Final(edgeservice.SingleRequestResult{Output: "duplicate"}); err != nil {
		t.Fatalf("post-terminal final returned established success as error: %v", err)
	}
	if got := w.Body.String(); got != wireAtTerminal {
		t.Fatalf("post-terminal call changed wire:\n%s", got)
	}

	events := parseSingleRequestAnthropicSSE(t, wireAtTerminal)
	if countSingleRequestAnthropicEvents(events, "message_start") != 1 ||
		countSingleRequestAnthropicEvents(events, "message_delta") != 1 ||
		countSingleRequestAnthropicEvents(events, "message_stop") != 1 ||
		countSingleRequestAnthropicEvents(events, "error") != 0 {
		t.Fatalf("unexpected envelope/terminal events: %+v", events)
	}
	if events[0].Name != "message_start" || events[len(events)-1].Name != "message_stop" {
		t.Fatalf("stream endpoints=%s/%s, want message_start/message_stop", events[0].Name, events[len(events)-1].Name)
	}
	message, _ := events[0].Data["message"].(map[string]any)
	if message["id"] != "msg_iop_public" || message["model"] != "virtual-model" {
		t.Fatalf("message identity=%+v", message)
	}

	wantTexts := []string{
		"Planning the requested work.",
		"Executing the requested work.",
		"Reviewing the completed work.",
		"safe final result",
	}
	if got := singleRequestAnthropicDeltaTexts(events); fmt.Sprint(got) != fmt.Sprint(wantTexts) {
		t.Fatalf("text deltas=%q, want %q", got, wantTexts)
	}
	block := 0
	for _, event := range events {
		if event.Name != "content_block_start" {
			continue
		}
		if got := int(event.Data["index"].(float64)); got != block {
			t.Fatalf("content block index=%d, want %d", got, block)
		}
		block++
	}
	if block != len(wantTexts) {
		t.Fatalf("content blocks=%d, want %d", block, len(wantTexts))
	}
}

func TestSingleRequestAnthropicStreamTerminalDispositionMatrix(t *testing.T) {
	const privatePartial = "PRIVATE_STREAM_PARTIAL_OUTPUT"
	tests := []struct {
		name        string
		disposition edgeservice.SingleRequestTerminalDisposition
		wantStop    string
		wantType    string
		wantMessage string
		silent      bool
	}{
		{name: "end turn", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalEndTurn}, wantStop: "end_turn"},
		{name: "length", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalLength}, wantStop: "max_tokens"},
		{name: "cancelled", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalCancelled}, silent: true},
		{name: "provider", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}, wantType: "api_error", wantMessage: "single-request execution failed"},
		{name: "validation", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorValidation}, wantType: "invalid_request_error", wantMessage: "single-request execution was rejected"},
		{name: "timeout", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorTimeout}, wantType: "api_error", wantMessage: "single-request execution timed out"},
		{name: "budget", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorBudget}, wantType: "api_error", wantMessage: "single-request execution failed"},
		{name: "repetition", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorRepetition}, wantType: "api_error", wantMessage: "single-request execution failed"},
		{name: "malformed", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorMalformed}, wantType: "api_error", wantMessage: "single-request execution failed"},
		{name: "context", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorContext}, wantType: "invalid_request_error", wantMessage: "single-request context limit exceeded"},
		{name: "internal tool", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorInternalTool}, wantType: "api_error", wantMessage: "single-request execution failed"},
		{name: "workspace cleanup", disposition: edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorWorkspaceCleanup}, wantType: "api_error", wantMessage: "single-request execution failed"},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			w := httptest.NewRecorder()
			stream, err := newSingleRequestAnthropicStream(w, "req_terminal", "virtual-model")
			if err != nil {
				t.Fatal(err)
			}
			var observed []edgeservice.SingleRequestTerminalDisposition
			stream.setTerminalRejectionObserver(func(disposition edgeservice.SingleRequestTerminalDisposition) {
				observed = append(observed, disposition)
			})
			if err := stream.Start(); err != nil {
				t.Fatal(err)
			}
			switch tc.disposition.Kind {
			case edgeservice.SingleRequestTerminalEndTurn:
				err = stream.Final(edgeservice.SingleRequestResult{Output: "safe final result", Terminal: tc.disposition})
			case edgeservice.SingleRequestTerminalLength:
				err = stream.Final(edgeservice.SingleRequestResult{Output: privatePartial, Terminal: tc.disposition})
			default:
				err = stream.TerminalError(tc.disposition)
			}
			if err != nil {
				t.Fatalf("terminal: %v", err)
			}
			wireAtTerminal := w.Body.String()
			if err := stream.TerminalError(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}); err != nil {
				t.Fatalf("post-terminal error: %v", err)
			}
			if err := stream.Final(edgeservice.SingleRequestResult{Output: "duplicate"}); err != nil {
				t.Fatalf("post-terminal final: %v", err)
			}
			if w.Body.String() != wireAtTerminal {
				t.Fatalf("second terminal changed wire: %q", w.Body.String())
			}
			if tc.disposition.Kind == edgeservice.SingleRequestTerminalError {
				if len(observed) != 1 || observed[0] != tc.disposition {
					t.Fatalf("terminal rejection observations=%+v, want exactly %+v", observed, tc.disposition)
				}
			} else if len(observed) != 0 {
				t.Fatalf("terminal rejection observations=%+v, want none", observed)
			}
			if strings.Contains(wireAtTerminal, privatePartial) {
				t.Fatalf("private partial output reached SSE: %q", wireAtTerminal)
			}

			events := parseSingleRequestAnthropicSSE(t, wireAtTerminal)
			if countSingleRequestAnthropicEvents(events, "message_start") != 1 {
				t.Fatalf("message_start count in %+v", events)
			}
			if tc.silent {
				if len(events) != 1 || countSingleRequestAnthropicEvents(events, "error") != 0 {
					t.Fatalf("cancel events=%+v, want silent after message_start", events)
				}
				return
			}
			if tc.wantStop != "" {
				if countSingleRequestAnthropicEvents(events, "message_delta") != 1 || countSingleRequestAnthropicEvents(events, "message_stop") != 1 || countSingleRequestAnthropicEvents(events, "error") != 0 {
					t.Fatalf("message terminal events=%+v", events)
				}
				delta := events[len(events)-2].Data["delta"].(map[string]any)
				if delta["stop_reason"] != tc.wantStop {
					t.Fatalf("stop_reason=%v, want %q", delta["stop_reason"], tc.wantStop)
				}
				if tc.disposition.Kind == edgeservice.SingleRequestTerminalLength && len(singleRequestAnthropicDeltaTexts(events)) != 0 {
					t.Fatalf("length text deltas=%q, want none", singleRequestAnthropicDeltaTexts(events))
				}
				return
			}
			if countSingleRequestAnthropicEvents(events, "error") != 1 || countSingleRequestAnthropicEvents(events, "message_stop") != 0 {
				t.Fatalf("error terminal events=%+v", events)
			}
			errorBody, _ := events[len(events)-1].Data["error"].(map[string]any)
			if errorBody["type"] != tc.wantType || errorBody["message"] != tc.wantMessage {
				t.Fatalf("error=%+v, want %s/%q", errorBody, tc.wantType, tc.wantMessage)
			}
		})
	}
}

func TestSingleRequestAnthropicStreamLiveContextProviderCancellation(t *testing.T) {
	var providerCalls atomic.Int32
	var providerContextDone atomic.Bool
	runner := &mockService{submit: func(ctx context.Context, _ edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
		if ctx.Err() != nil {
			providerContextDone.Store(true)
		}
		providerCalls.Add(1)
		return nil, context.Canceled
	}}
	executor := NewSingleRequestExecutor(runner)
	svc := newAdmittedAnthropicSingleRequestService(t, executor, "ws-opaque-ref")
	srv := newAnthropicSingleRequestServer(t, svc)
	w := httptest.NewRecorder()
	before := testutil.ToFloat64(singleRequestIngressTotal)
	body := `{"model":"` + testSingleRequestModel + `","max_tokens":128,"stream":true,"messages":[{"role":"user","content":"run"}]}`

	serveAnthropicSingleRequest(t, srv, context.Background(), "/v1/messages", body, w)

	if got := providerCalls.Load(); got != 1 {
		t.Fatalf("provider calls=%d, want 1", got)
	}
	if providerContextDone.Load() {
		t.Fatal("provider context was already done before the provider-owned cancellation")
	}
	if got := testutil.ToFloat64(singleRequestIngressTotal) - before; got != 1 {
		t.Fatalf("ingress delta=%v, want 1", got)
	}
	if w.Code != http.StatusOK || !strings.HasPrefix(w.Header().Get("Content-Type"), "text/event-stream") {
		t.Fatalf("status=%d content-type=%q body=%s", w.Code, w.Header().Get("Content-Type"), w.Body.String())
	}
	events := parseSingleRequestAnthropicSSE(t, w.Body.String())
	if countSingleRequestAnthropicEvents(events, "message_start") != 1 ||
		countSingleRequestAnthropicEvents(events, "error") != 1 ||
		countSingleRequestAnthropicEvents(events, "message_stop") != 0 {
		t.Fatalf("events=%+v, want one sanitized error terminal", events)
	}
	errorBody, _ := events[len(events)-1].Data["error"].(map[string]any)
	if errorBody["type"] != "api_error" || errorBody["message"] != "single-request execution failed" {
		t.Fatalf("error=%+v, want sanitized provider api_error", errorBody)
	}
	for _, forbidden := range []string{"context canceled", "cancelled", "ws-opaque-ref", "plan-model", "provider-plan"} {
		if strings.Contains(w.Body.String(), forbidden) {
			t.Fatalf("streaming error leaked %q: %s", forbidden, w.Body.String())
		}
	}
}

func TestSingleRequestAnthropicStreamPingAndProgressOrdering(t *testing.T) {
	w := httptest.NewRecorder()
	stream, err := newSingleRequestAnthropicStream(w, "req_ping", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Progress(edgeservice.SingleRequestProgress{Stage: edgeservice.SingleRequestStatePlanning}); err != nil {
		t.Fatal(err)
	}
	if err := stream.Ping(); err != nil {
		t.Fatal(err)
	}
	if err := stream.Progress(edgeservice.SingleRequestProgress{Stage: edgeservice.SingleRequestStateWorking}); err != nil {
		t.Fatal(err)
	}
	if err := stream.Ping(); err != nil {
		t.Fatal(err)
	}
	if err := stream.Final(edgeservice.SingleRequestResult{Output: "done"}); err != nil {
		t.Fatal(err)
	}

	events := parseSingleRequestAnthropicSSE(t, w.Body.String())
	var names []string
	for _, event := range events {
		names = append(names, event.Name)
	}
	want := []string{
		"message_start",
		"content_block_start", "content_block_delta", "content_block_stop",
		"ping",
		"content_block_start", "content_block_delta", "content_block_stop",
		"ping",
		"content_block_start", "content_block_delta", "content_block_stop",
		"message_delta", "message_stop",
	}
	if fmt.Sprint(names) != fmt.Sprint(want) {
		t.Fatalf("event order=%v, want %v", names, want)
	}
}

func TestSingleRequestAnthropicStreamRepairSummary(t *testing.T) {
	w := httptest.NewRecorder()
	stream, err := newSingleRequestAnthropicStream(w, "req_repair", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	for _, stage := range []edgeservice.SingleRequestState{
		edgeservice.SingleRequestStateReviewing,
		edgeservice.SingleRequestStateRepairing,
		edgeservice.SingleRequestStateInternalTool,
		edgeservice.SingleRequestStateRepairing,
	} {
		if err := stream.Progress(edgeservice.SingleRequestProgress{Stage: stage}); err != nil {
			t.Fatal(err)
		}
	}
	if err := stream.Final(edgeservice.SingleRequestResult{Output: "repaired"}); err != nil {
		t.Fatal(err)
	}
	texts := singleRequestAnthropicDeltaTexts(parseSingleRequestAnthropicSSE(t, w.Body.String()))
	if got := strings.Join(texts, "|"); got != "Reviewing the completed work.|Repairing issues found during review.|repaired" {
		t.Fatalf("repair projection=%q", got)
	}
}

func TestSingleRequestAnthropicStreamRedactsPrivateEvents(t *testing.T) {
	const private = "PRIVATE_PROVIDER_ROUTE_CREDENTIAL_WORKSPACE_COMMAND_TOOL_SENTINEL"
	w := httptest.NewRecorder()
	stream, err := newSingleRequestAnthropicStream(w, "req_private", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Progress(edgeservice.SingleRequestProgress{
		Stage:   edgeservice.SingleRequestStatePlanning,
		Message: private,
		Result:  &edgeservice.SingleRequestResult{Output: private},
		Err:     errors.New(private),
	}); err != nil {
		t.Fatal(err)
	}
	if err := stream.Progress(edgeservice.SingleRequestProgress{Stage: edgeservice.SingleRequestState(private)}); !errors.Is(err, errSingleRequestAnthropicUnknownProgress) {
		t.Fatalf("unknown stage error=%v", err)
	}
	if err := stream.Error(singleRequestAnthropicTerminalFailure); err != nil {
		t.Fatal(err)
	}
	wire := w.Body.String()
	for _, forbidden := range []string{private, "tool_use", "thinking_delta", "input_json_delta"} {
		if strings.Contains(wire, forbidden) {
			t.Fatalf("wire leaked %q:\n%s", forbidden, wire)
		}
	}
	events := parseSingleRequestAnthropicSSE(t, wire)
	if countSingleRequestAnthropicEvents(events, "error") != 1 || countSingleRequestAnthropicEvents(events, "message_stop") != 0 {
		t.Fatalf("error terminal events=%+v", events)
	}
}

func TestSingleRequestAnthropicStreamErrorTerminalRace(t *testing.T) {
	w := httptest.NewRecorder()
	stream, err := newSingleRequestAnthropicStream(w, "req_race", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	if err := stream.Start(); err != nil {
		t.Fatal(err)
	}
	start := make(chan struct{})
	var wg sync.WaitGroup
	for i := 0; i < 48; i++ {
		wg.Add(1)
		go func(index int) {
			defer wg.Done()
			<-start
			switch index % 3 {
			case 0:
				_ = stream.Ping()
			case 1:
				_ = stream.Final(edgeservice.SingleRequestResult{Output: "safe"})
			case 2:
				_ = stream.Error(singleRequestAnthropicTerminalFailure)
			}
		}(i)
	}
	close(start)
	wg.Wait()

	events := parseSingleRequestAnthropicSSE(t, w.Body.String())
	terminalCount := countSingleRequestAnthropicEvents(events, "message_stop") + countSingleRequestAnthropicEvents(events, "error")
	if terminalCount != 1 {
		t.Fatalf("terminal count=%d events=%+v", terminalCount, events)
	}
	terminalIndex := -1
	for index, event := range events {
		if event.Name == "message_stop" || event.Name == "error" {
			terminalIndex = index
		}
	}
	if terminalIndex != len(events)-1 {
		t.Fatalf("events followed terminal: %+v", events[terminalIndex+1:])
	}
}

type manualSingleRequestAnthropicTicker struct {
	ticks   chan time.Time
	stopped chan struct{}
	once    sync.Once
}

func newManualSingleRequestAnthropicTicker() *manualSingleRequestAnthropicTicker {
	return &manualSingleRequestAnthropicTicker{
		ticks: make(chan time.Time, 8), stopped: make(chan struct{}),
	}
}

func (t *manualSingleRequestAnthropicTicker) Ticks() <-chan time.Time { return t.ticks }
func (t *manualSingleRequestAnthropicTicker) Stop() {
	t.once.Do(func() { close(t.stopped) })
}

type observedSingleRequestAnthropicWriter struct {
	mu              sync.Mutex
	header          http.Header
	status          int
	body            bytes.Buffer
	events          chan string
	failEvent       string
	flushErrorEvent string
	lastEvent       string
	onEvent         func(string)
}

func newObservedSingleRequestAnthropicWriter() *observedSingleRequestAnthropicWriter {
	return &observedSingleRequestAnthropicWriter{
		header: make(http.Header), events: make(chan string, 64),
	}
}

func (w *observedSingleRequestAnthropicWriter) Header() http.Header { return w.header }
func (w *observedSingleRequestAnthropicWriter) WriteHeader(status int) {
	w.mu.Lock()
	w.status = status
	w.mu.Unlock()
}
func (w *observedSingleRequestAnthropicWriter) Write(p []byte) (int, error) {
	name := ""
	if line, _, ok := strings.Cut(string(p), "\n"); ok && strings.HasPrefix(line, "event: ") {
		name = strings.TrimPrefix(line, "event: ")
	}
	if name == w.failEvent {
		return 0, io.ErrClosedPipe
	}
	w.mu.Lock()
	n, err := w.body.Write(p)
	w.lastEvent = name
	w.mu.Unlock()
	if name != "" {
		if w.onEvent != nil {
			w.onEvent(name)
		}
		w.events <- name
	}
	return n, err
}
func (w *observedSingleRequestAnthropicWriter) Flush() {}
func (w *observedSingleRequestAnthropicWriter) FlushError() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.lastEvent == w.flushErrorEvent {
		return io.ErrClosedPipe
	}
	return nil
}
func (w *observedSingleRequestAnthropicWriter) String() string {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.body.String()
}

func waitForSingleRequestAnthropicEvent(t *testing.T, events <-chan string, want string) {
	t.Helper()
	timer := time.NewTimer(2 * time.Second)
	defer timer.Stop()
	for {
		select {
		case event := <-events:
			if event == want {
				return
			}
		case <-timer.C:
			t.Fatalf("timed out waiting for %s", want)
		}
	}
}

func newSingleRequestAnthropicTestBinding(t *testing.T) *edgeservice.SingleRequestBinding {
	t.Helper()
	binding, err := edgeservice.NewSingleRequestBinding(
		"virtual-model",
		"opaque-workspace",
		edgeservice.SingleRequestStageBinding{Model: "plan"},
		edgeservice.SingleRequestStageBinding{Model: "work"},
		edgeservice.SingleRequestStageBinding{Model: "review"},
		edgeservice.SingleRequestLimits{
			WallClockMS: 10_000, StageTimeoutMS: 5_000, MaxToolIterations: 4, MaxOutputBytes: 4096,
		},
	)
	if err != nil {
		t.Fatal(err)
	}
	return binding
}

func startSingleRequestAnthropicTestExecution(
	t *testing.T,
	executor anthropicSingleRequestExecutorFunc,
) edgeservice.SingleRequestExecution {
	t.Helper()
	svc := newAdmittedAnthropicSingleRequestService(t, executor, "opaque-workspace")
	execution, err := svc.StartSingleRequest(context.Background(), edgeservice.SingleRequestRequest{
		RequestID: "req_pump", Binding: newSingleRequestAnthropicTestBinding(t), Prompt: "private prompt",
	})
	if err != nil {
		t.Fatal(err)
	}
	return execution
}

func TestSingleRequestAnthropicStreamPumpStopsPingBeforeTerminal(t *testing.T) {
	release := make(chan struct{})
	execution := startSingleRequestAnthropicTestExecution(t, func(
		_ context.Context,
		req edgeservice.SingleRequestRequest,
		ctrl edgeservice.SingleRequestController,
	) error {
		if err := ctrl.SubmitEnvelope(edgeservice.SingleRequestEnvelope{
			RequestID: req.RequestID, Sequence: 1, Stage: edgeservice.SingleRequestStatePlanning,
		}); err != nil {
			return err
		}
		<-release
		for index, stage := range []edgeservice.SingleRequestState{
			edgeservice.SingleRequestStateWorking,
			edgeservice.SingleRequestStateReviewing,
			edgeservice.SingleRequestStateFinalizing,
		} {
			envelope := edgeservice.SingleRequestEnvelope{
				RequestID: req.RequestID, Sequence: uint64(index + 2), Stage: stage,
			}
			if stage == edgeservice.SingleRequestStateFinalizing {
				envelope.Result = &edgeservice.SingleRequestResult{Output: "safe final"}
			}
			if err := ctrl.SubmitEnvelope(envelope); err != nil {
				return err
			}
		}
		return nil
	})

	w := newObservedSingleRequestAnthropicWriter()
	stateAtStop := make(chan edgeservice.SingleRequestState, 1)
	w.onEvent = func(event string) {
		if event == "message_stop" {
			stateAtStop <- execution.State()
		}
	}
	stream, err := newSingleRequestAnthropicStream(w, "req_pump", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	ticker := newManualSingleRequestAnthropicTicker()
	done := make(chan error, 1)
	go func() {
		done <- pumpSingleRequestAnthropicStream(context.Background(), execution, stream, func() singleRequestAnthropicTicker {
			return ticker
		})
	}()

	waitForSingleRequestAnthropicEvent(t, w.events, "content_block_stop")
	ticker.ticks <- time.Unix(1, 0)
	waitForSingleRequestAnthropicEvent(t, w.events, "ping")
	close(release)
	if err := <-done; err != nil {
		t.Fatal(err)
	}
	if got := <-stateAtStop; got != edgeservice.SingleRequestStateFinalizing {
		t.Fatalf("state at message_stop=%s, want finalizing", got)
	}
	if got := execution.State(); got != edgeservice.SingleRequestStateCompleted {
		t.Fatalf("state after terminal acknowledgement=%s, want completed", got)
	}
	select {
	case <-ticker.stopped:
	default:
		t.Fatal("ticker was not stopped before pump return")
	}

	wireAtTerminal := w.String()
	ticker.ticks <- time.Unix(2, 0)
	if got := w.String(); got != wireAtTerminal {
		t.Fatalf("manual tick wrote after terminal:\n%s", got)
	}
	events := parseSingleRequestAnthropicSSE(t, wireAtTerminal)
	if countSingleRequestAnthropicEvents(events, "ping") != 1 || events[len(events)-1].Name != "message_stop" {
		t.Fatalf("ping/terminal events=%+v", events)
	}
}

func TestSingleRequestAnthropicStreamTerminalWriteFailureDoesNotComplete(t *testing.T) {
	execution := startSingleRequestAnthropicTestExecution(t, func(
		_ context.Context,
		req edgeservice.SingleRequestRequest,
		ctrl edgeservice.SingleRequestController,
	) error {
		for index, stage := range []edgeservice.SingleRequestState{
			edgeservice.SingleRequestStatePlanning,
			edgeservice.SingleRequestStateWorking,
			edgeservice.SingleRequestStateReviewing,
			edgeservice.SingleRequestStateFinalizing,
		} {
			envelope := edgeservice.SingleRequestEnvelope{
				RequestID: req.RequestID, Sequence: uint64(index + 1), Stage: stage,
			}
			if stage == edgeservice.SingleRequestStateFinalizing {
				envelope.Result = &edgeservice.SingleRequestResult{Output: "safe final"}
			}
			if err := ctrl.SubmitEnvelope(envelope); err != nil {
				return err
			}
		}
		return nil
	})
	w := newObservedSingleRequestAnthropicWriter()
	w.failEvent = "message_stop"
	stream, err := newSingleRequestAnthropicStream(w, "req_pump", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	ticker := newManualSingleRequestAnthropicTicker()
	err = pumpSingleRequestAnthropicStream(context.Background(), execution, stream, func() singleRequestAnthropicTicker {
		return ticker
	})
	if !errors.Is(err, io.ErrClosedPipe) {
		t.Fatalf("pump error=%v, want write failure; wire=%s", err, w.String())
	}
	if got := execution.State(); got != edgeservice.SingleRequestStateFailed {
		t.Fatalf("state=%s, want failed", got)
	}
	if strings.Contains(w.String(), "event: message_stop") {
		t.Fatalf("failed message_stop unexpectedly reached wire:\n%s", w.String())
	}
}

func TestSingleRequestAnthropicStreamTerminalFlushFailureDoesNotComplete(t *testing.T) {
	execution := startSingleRequestAnthropicTestExecution(t, func(
		_ context.Context,
		req edgeservice.SingleRequestRequest,
		ctrl edgeservice.SingleRequestController,
	) error {
		for index, stage := range []edgeservice.SingleRequestState{
			edgeservice.SingleRequestStatePlanning,
			edgeservice.SingleRequestStateWorking,
			edgeservice.SingleRequestStateReviewing,
			edgeservice.SingleRequestStateFinalizing,
		} {
			envelope := edgeservice.SingleRequestEnvelope{
				RequestID: req.RequestID, Sequence: uint64(index + 1), Stage: stage,
			}
			if stage == edgeservice.SingleRequestStateFinalizing {
				envelope.Result = &edgeservice.SingleRequestResult{Output: "safe final"}
			}
			if err := ctrl.SubmitEnvelope(envelope); err != nil {
				return err
			}
		}
		return nil
	})
	w := newObservedSingleRequestAnthropicWriter()
	w.flushErrorEvent = "message_stop"
	stream, err := newSingleRequestAnthropicStream(w, "req_pump", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	ticker := newManualSingleRequestAnthropicTicker()
	err = pumpSingleRequestAnthropicStream(context.Background(), execution, stream, func() singleRequestAnthropicTicker {
		return ticker
	})
	if !errors.Is(err, io.ErrClosedPipe) {
		t.Fatalf("pump error=%v, want flush failure; wire=%s", err, w.String())
	}
	if got := execution.State(); got != edgeservice.SingleRequestStateFailed {
		t.Fatalf("state=%s, want failed", got)
	}
	if !strings.Contains(w.String(), "event: message_stop") {
		t.Fatalf("message_stop bytes did not reach writer before flush failure:\n%s", w.String())
	}
}

func TestSingleRequestAnthropicStreamDisconnectStopsWriter(t *testing.T) {
	executorStarted := make(chan struct{})
	execution := startSingleRequestAnthropicTestExecution(t, func(
		ctx context.Context,
		_ edgeservice.SingleRequestRequest,
		_ edgeservice.SingleRequestController,
	) error {
		close(executorStarted)
		<-ctx.Done()
		return ctx.Err()
	})
	<-executorStarted
	w := newObservedSingleRequestAnthropicWriter()
	stream, err := newSingleRequestAnthropicStream(w, "req_pump", "virtual-model")
	if err != nil {
		t.Fatal(err)
	}
	ticker := newManualSingleRequestAnthropicTicker()
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		done <- pumpSingleRequestAnthropicStream(ctx, execution, stream, func() singleRequestAnthropicTicker {
			return ticker
		})
	}()
	waitForSingleRequestAnthropicEvent(t, w.events, "message_start")
	cancel()
	if err := <-done; !errors.Is(err, context.Canceled) {
		t.Fatalf("pump error=%v, want context cancellation", err)
	}
	wireAtReturn := w.String()
	ticker.ticks <- time.Unix(3, 0)
	if got := w.String(); got != wireAtReturn {
		t.Fatalf("tick wrote after disconnect:\n%s", got)
	}
	if strings.Contains(wireAtReturn, "event: message_stop") || strings.Contains(wireAtReturn, "event: error") {
		t.Fatalf("disconnect synthesized terminal after caller cancellation:\n%s", wireAtReturn)
	}
}

func TestAnthropicSingleRequestStreamingUsesOnePost(t *testing.T) {
	const (
		privatePrompt = "PRIVATE_STREAMING_CALLER_PROMPT_SENTINEL"
		finalOutput   = "streaming workspace task completed"
	)
	var calls atomic.Int32
	controllerCh := make(chan edgeservice.SingleRequestController, 1)
	executor := anthropicSingleRequestExecutorFunc(func(
		_ context.Context,
		req edgeservice.SingleRequestRequest,
		ctrl edgeservice.SingleRequestController,
	) error {
		calls.Add(1)
		if err := submitAnthropicSingleRequestLifecycle(req, ctrl, finalOutput); err != nil {
			return err
		}
		controllerCh <- ctrl
		return nil
	})
	svc := newAdmittedAnthropicSingleRequestService(t, executor, "ws-opaque-ref")
	srv := newAnthropicSingleRequestServer(t, svc)
	httpServer := httptest.NewServer(srv.routes())
	defer httpServer.Close()

	before := testutil.ToFloat64(singleRequestIngressTotal)
	body := `{"model":"` + testSingleRequestModel + `","max_tokens":128,"stream":true,"messages":[{"role":"user","content":"` + privatePrompt + `"}]}`
	req := newAnthropicSingleRequestHTTPReq(t, context.Background(), httpServer.URL, "/v1/messages", body)
	response, err := httpServer.Client().Do(req)
	if err != nil {
		t.Fatalf("POST /v1/messages: %v", err)
	}
	defer response.Body.Close()
	wire, err := io.ReadAll(response.Body)
	if err != nil {
		t.Fatal(err)
	}
	if response.StatusCode != http.StatusOK || !strings.HasPrefix(response.Header.Get("Content-Type"), "text/event-stream") {
		t.Fatalf("status=%d content-type=%q body=%s", response.StatusCode, response.Header.Get("Content-Type"), wire)
	}
	events := parseSingleRequestAnthropicSSE(t, string(wire))
	if countSingleRequestAnthropicEvents(events, "message_start") != 1 ||
		countSingleRequestAnthropicEvents(events, "message_stop") != 1 ||
		countSingleRequestAnthropicEvents(events, "error") != 0 {
		t.Fatalf("unexpected streaming terminal: %+v", events)
	}
	if !strings.Contains(string(wire), finalOutput) {
		t.Fatalf("final output missing from wire: %s", wire)
	}
	for _, forbidden := range []string{
		privatePrompt, "PRIVATE_STAGE_SENTINEL", "tool_use", "plan-model", "provider-plan", "slot-plan", "ws-opaque-ref",
	} {
		if strings.Contains(string(wire), forbidden) {
			t.Fatalf("stream leaked %q: %s", forbidden, wire)
		}
	}
	if got := calls.Load(); got != 1 {
		t.Fatalf("executor calls=%d, want 1", got)
	}
	if got := testutil.ToFloat64(singleRequestIngressTotal) - before; got != 1 {
		t.Fatalf("single-request ingress counter delta=%v, want 1", got)
	}
	if got := (<-controllerCh).State(); got != edgeservice.SingleRequestStateCompleted {
		t.Fatalf("terminal acknowledgement state=%s, want completed", got)
	}
}
