package openai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
	"unicode"

	"go.uber.org/zap"

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

const (
	streamTraceMetadataKey = "iop_trace_stream"
	streamTraceEnvKey      = "IOP_OPENAI_COMPAT_TRACE_STREAM"
)

func (s *Server) streamChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, submitReq edgeservice.SubmitRunRequest, handle edgeservice.RunResult, outputPolicy strictOutputPolicy, validation toolValidationContract, retrySubmit func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error)) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		handle.Close()
		writeError(w, http.StatusInternalServerError, "streaming_not_supported", "response writer does not support streaming")
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	// Buffered streams collect and validate the full response before emitting
	// user-visible chunks. Strict buffered output opts in explicitly, and
	// tool-bearing streams use the same path so malformed tool calls can be
	// retried or rejected before reaching a client tool runner.
	if (outputPolicy.Strict && outputPolicy.StreamBuffer) || len(req.Tools) > 0 {
		s.streamBufferedChatCompletion(w, r, req, submitReq, handle, flusher, outputPolicy, validation, retrySubmit)
		return
	}

	// Live SSE may emit content deltas before the terminal event, so runtime
	// tool validation is excluded upstream; write the role chunk immediately.
	defer handle.Close()
	created := time.Now().Unix()
	id := "chatcmpl-" + handle.Dispatch().RunID
	model := responseModel(req.Model, handle.Dispatch().Target)
	metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModeNormalized)
	traceStream := openAICompatTraceStreamEnabled(submitReq.Metadata)
	traceSeq := 0
	shouldExposeReasoning := req.includeReasoning() && (!outputPolicy.Strict || req.explicitlyIncludesReasoning())
	writeTracedContentDelta := func(content, source string, filtered bool) {
		if content == "" {
			return
		}
		traceSeq++
		if traceStream {
			logOpenAICompatStreamOutput(s.logger, handle.Dispatch().RunID, traceSeq, "content", content,
				zap.Int("source_len", len(source)),
				zap.Bool("filtered", filtered),
				zap.Bool("source_has_open_think", hasOpenThinkTag(source)),
				zap.Bool("source_has_close_think", hasCloseThinkTag(source)),
			)
		}
		writeContentDeltaSSE(w, flusher, id, created, model, content)
	}
	writeTracedReasoningDelta := func(reasoning string) {
		if reasoning == "" {
			return
		}
		traceSeq++
		if traceStream {
			logOpenAICompatStreamOutput(s.logger, handle.Dispatch().RunID, traceSeq, "reasoning", reasoning)
		}
		writeSSE(w, flusher, chatCompletionChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []chatCompletionChunkChoice{{
				Index: 0,
				Delta: chatDelta{ReasoningContent: reasoning},
			}},
		})
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})

	var contentBuilder strings.Builder
	var reasoningBuilder strings.Builder
	var emittedContent strings.Builder
	var toolTextFilter *streamToolTextFilter
	contentSentinelFilter := &streamSentinelFilter{}
	reasoningSentinelFilter := &streamSentinelFilter{}
	if len(req.Tools) > 0 {
		toolTextFilter = &streamToolTextFilter{}
	}
	defer func() {
		s.logger.Info("openai chat completion stream closed",
			zap.String("run_id", handle.Dispatch().RunID),
			zap.Bool("strict_output", outputPolicy.Strict),
			zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
			zap.Int("content_len", contentBuilder.Len()),
			zap.Int("reasoning_len", reasoningBuilder.Len()),
		)
	}()

	stream := handle.Stream()
	if stream.Events == nil {
		writeSSEError(w, flusher, "run stream unavailable")
		return
	}
	processContentDelta := func(sourceDelta string) {
		if sourceDelta == "" {
			return
		}
		delta := contentSentinelFilter.Append(sourceDelta)
		if delta == "" {
			return
		}
		contentBuilder.WriteString(delta)
		filteredDelta := delta
		if toolTextFilter != nil {
			filteredDelta = toolTextFilter.Append(filteredDelta)
		}
		if filteredDelta == "" {
			if traceStream {
				s.logger.Info("openai chat completion stream output suppressed",
					zap.String("run_id", handle.Dispatch().RunID),
					zap.String("type", "content"),
					zap.Int("source_len", len(sourceDelta)),
					zap.Bool("source_has_open_think", hasOpenThinkTag(sourceDelta)),
					zap.Bool("source_has_close_think", hasCloseThinkTag(sourceDelta)),
				)
			}
			return
		}
		emittedContent.WriteString(filteredDelta)
		writeTracedContentDelta(filteredDelta, sourceDelta, filteredDelta != sourceDelta)
	}
	for {
		select {
		case <-r.Context().Done():
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
			emitUsageMetrics(metricLabels, usageStatusForError(r.Context().Err()), usageObservation{reasoningChars: reasoningBuilder.Len()})
			return
		case nodeEvent, ok := <-stream.NodeEvents:
			if !ok {
				stream.NodeEvents = nil
				continue
			}
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				emitUsageMetrics(metricLabels, usageStatusError, usageObservation{reasoningChars: reasoningBuilder.Len()})
				writeSSEError(w, flusher, "node disconnected")
				return
			}
		case event, ok := <-stream.Events:
			if !ok {
				writeSSEError(w, flusher, "run stream closed")
				return
			}
			if event == nil {
				continue
			}
			switch event.GetType() {
			case "delta":
				processContentDelta(event.GetDelta())
			case "reasoning_delta":
				if event.GetDelta() == "" {
					continue
				}
				reasoning := reasoningSentinelFilter.Append(event.GetDelta())
				if reasoning == "" {
					continue
				}
				reasoningBuilder.WriteString(reasoning)
				if !shouldExposeReasoning {
					continue
				}
				writeTracedReasoningDelta(reasoning)
			case "complete":
				processContentDelta(contentSentinelFilter.Flush())
				if reasoning := reasoningSentinelFilter.Flush(); reasoning != "" {
					reasoningBuilder.WriteString(reasoning)
					if shouldExposeReasoning {
						writeTracedReasoningDelta(reasoning)
					}
				}
				metadata := event.GetMetadata()
				streamUsage := usageObservation{reasoningChars: reasoningBuilder.Len()}
				if u := event.GetUsage(); u != nil {
					streamUsage.inputTokens = int(u.GetInputTokens())
					streamUsage.outputTokens = int(u.GetOutputTokens())
					streamUsage.reasoningTokens = int(u.GetReasoningTokens())
					streamUsage.cachedInputTokens = int(u.GetCachedInputTokens())
				}
				finishReason := metadata["finish_reason"]
				if finishReason == "" {
					finishReason = "stop"
				}
				nativeToolCalls := toolCallsFromRunMetadata(metadata)
				if len(nativeToolCalls) > 0 {
					if toolTextFilter != nil {
						pending := toolTextFilter.FlushNonCandidate()
						if pending != "" {
							emittedContent.WriteString(pending)
							writeTracedContentDelta(pending, pending, false)
						}
					}
					nativeToolCalls = normalizeNativeToolCallsForResponse(nativeToolCalls, req.Tools)
					if traceStream {
						s.logger.Info("openai chat completion stream output chunk",
							zap.String("run_id", handle.Dispatch().RunID),
							zap.Int("seq", traceSeq+1),
							zap.String("type", "tool_calls"),
							zap.Int("tool_call_count", len(nativeToolCalls)),
						)
					}
					writeToolCallsDeltaSSE(w, flusher, id, created, model, nativeToolCalls)
					finishReason = "tool_calls"
				} else if len(req.Tools) > 0 {
					res := synthesizeToolCallsFromTextResult(contentBuilder.String(), req.Tools, handle.Dispatch().RunID)
					if res.validationErr != nil {
						emitUsageMetrics(metricLabels, usageStatusError, usageObservation{reasoningChars: reasoningBuilder.Len()})
						writeSSEErrorWithType(w, flusher, "tool_validation_error", res.validationErr.Error())
						return
					}
					if len(res.toolCalls) > 0 {
						if remaining := unstreamedCleanedText(res.cleaned, emittedContent.String()); remaining != "" {
							emittedContent.WriteString(remaining)
							writeTracedContentDelta(remaining, remaining, false)
						}
						if traceStream {
							s.logger.Info("openai chat completion stream output chunk",
								zap.String("run_id", handle.Dispatch().RunID),
								zap.Int("seq", traceSeq+1),
								zap.String("type", "tool_calls"),
								zap.Int("tool_call_count", len(res.toolCalls)),
							)
						}
						writeToolCallsDeltaSSE(w, flusher, id, created, model, res.toolCalls)
						finishReason = "tool_calls"
					} else if !res.candidateFound {
						if toolTextFilter != nil {
							if pending := toolTextFilter.Flush(); pending != "" {
								emittedContent.WriteString(pending)
								writeTracedContentDelta(pending, pending, false)
							}
						}
					}
				} else if toolTextFilter != nil {
					if pending := toolTextFilter.Flush(); pending != "" {
						emittedContent.WriteString(pending)
						writeTracedContentDelta(pending, pending, false)
					}
				}
				if !outputPolicy.Strict &&
					strings.TrimSpace(contentBuilder.String()) == "" &&
					strings.TrimSpace(emittedContent.String()) == "" &&
					strings.TrimSpace(reasoningBuilder.String()) != "" &&
					len(nativeToolCalls) == 0 {
					fallback := hiddenReasoningFallbackContent(finishReason)
					if req.includeReasoning() {
						fallback = reasoningOnlyFallbackContent(reasoningBuilder.String(), finishReason)
					}
					emittedContent.WriteString(fallback)
					writeTracedContentDelta(fallback, reasoningBuilder.String(), false)
				}
				if traceStream {
					s.logger.Info("openai chat completion stream output done",
						zap.String("run_id", handle.Dispatch().RunID),
						zap.Int("seq", traceSeq+1),
						zap.String("finish_reason", finishReason),
						zap.Int("content_len", contentBuilder.Len()),
						zap.Int("reasoning_len", reasoningBuilder.Len()),
					)
				}
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index:        0,
						Delta:        chatDelta{},
						FinishReason: finishReason,
					}},
				})
				fmt.Fprint(w, "data: [DONE]\n\n")
				flusher.Flush()
				emitUsageMetrics(metricLabels, usageStatusSuccess, streamUsage)
				return
			case "error", "cancelled":
				msg := event.GetError()
				if msg == "" {
					msg = event.GetMessage()
				}
				if msg == "" {
					msg = "run failed"
				}
				emitUsageMetrics(metricLabels, usageStatusError, usageObservation{reasoningChars: reasoningBuilder.Len()})
				writeSSEError(w, flusher, msg)
				return
			}
		case <-time.After(handle.WaitTimeout()):
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
			emitUsageMetrics(metricLabels, usageStatusCancel, usageObservation{reasoningChars: reasoningBuilder.Len()})
			writeSSEError(w, flusher, errRunTimedOut.Error())
			return
		}
	}
}

// tunnelChatCompletionPassthrough serves a Chat Completions request over the
// provider tunnel (SDD S04/S06): provider status, headers, and content are
// relayed without IOP sideband fields or events. Chat Completions model echoes
// are normalized back to the caller-facing model alias.
func (s *Server) tunnelChatCompletionPassthrough(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, dispatch routeDispatch, runMeta map[string]string, rawBody []byte, estimate int, contextClass string) {
	handle, ok := s.submitChatCompletionTunnel(w, r, req, dispatch, runMeta, rawBody, estimate, contextClass)
	if !ok {
		return
	}
	defer handle.Close()
	metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModePassthrough)
	s.writeProviderTunnelResponse(w, r, handle, req.Stream, req.Model, metricLabels)
}

// errProviderAuthRequired signals that provider auth forwarding is configured
// as required but the caller did not supply the configured provider token
// header. The raw token is never part of this error.
var errProviderAuthRequired = errors.New("provider auth token is required")

// providerTunnelAuthHeaders builds the provider tunnel request headers that
// forward a request-time raw user token to the provider (SDD S03). It reads the
// token only from the configured `openai.provider_auth.from_header` and never
// reuses the inbound IOP `Authorization` header. Returns nil when provider auth
// is disabled or the (non-required) token header is absent, and
// errProviderAuthRequired when a required token header is missing.
func (s *Server) providerTunnelAuthHeaders(r *http.Request) (map[string]string, error) {
	auth := s.cfg.ProviderAuth
	if !auth.Enabled {
		return nil, nil
	}
	raw := strings.TrimSpace(r.Header.Get(auth.FromHeader))
	if raw == "" {
		if auth.Required {
			return nil, errProviderAuthRequired
		}
		return nil, nil
	}
	value := raw
	if scheme := strings.TrimSpace(auth.Scheme); scheme != "" &&
		!strings.HasPrefix(strings.ToLower(raw), strings.ToLower(scheme)+" ") {
		value = scheme + " " + raw
	}
	return map[string]string{auth.TargetHeader: value}, nil
}

// submitChatCompletionTunnel dispatches a Chat Completions provider tunnel
// request for pure passthrough. On failure it writes the dispatch error to the
// caller and returns ok=false.
func (s *Server) submitChatCompletionTunnel(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, dispatch routeDispatch, runMeta map[string]string, rawBody []byte, estimate int, contextClass string) (edgeservice.ProviderTunnelResult, bool) {
	providerAuthHeaders, err := s.providerTunnelAuthHeaders(r)
	if err != nil {
		// Missing required provider auth is rejected before dispatch; the raw
		// token is never echoed into the error surface.
		writeError(w, http.StatusBadRequest, "invalid_request_error", "provider auth token is required")
		return nil, false
	}

	metadata := make(map[string]string, len(runMeta)+5)
	for k, v := range runMeta {
		metadata[k] = v
	}
	metadata["openai_model"] = req.Model
	metadata["openai_stream"] = strconv.FormatBool(req.Stream)
	metadata["estimated_input_tokens"] = strconv.Itoa(estimate)
	metadata["context_class"] = contextClass

	tunnelReq := edgeservice.SubmitProviderTunnelRequest{
		NodeRef:       dispatch.NodeRef,
		ModelGroupKey: strings.TrimSpace(req.Model),
		Adapter:       dispatch.Adapter,
		Target:        dispatch.Target,
		SessionID:     dispatch.SessionID,
		Method:        http.MethodPost,
		Path:          "/v1/chat/completions",
		Headers:       providerAuthHeaders,
		BuildBody: func(target string) ([]byte, error) {
			return rewriteChatCompletionModel(rawBody, target, req)
		},
		Stream:               req.Stream,
		TimeoutSec:           dispatch.TimeoutSec,
		MaxQueue:             dispatch.MaxQueue,
		QueueTimeoutMS:       dispatch.QueueTimeoutMS,
		Metadata:             metadata,
		EstimatedInputTokens: estimate,
		ContextClass:         contextClass,
		ProviderPool:         dispatch.ProviderPool,
	}

	metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModePassthrough)
	handle, err := s.service.SubmitProviderTunnel(r.Context(), tunnelReq)
	if err != nil {
		emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
		writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
		return nil, false
	}

	s.logger.Info("openai chat completion passthrough dispatch",
		zap.String("run_id", handle.Dispatch().RunID),
		zap.String("node_id", handle.Dispatch().NodeID),
		zap.String("provider_id", handle.Dispatch().ProviderID),
		zap.String("provider_type", handle.Dispatch().ProviderType),
		zap.String("execution_path", handle.Dispatch().ExecutionPath),
		zap.String("model_group", handle.Dispatch().ModelGroupKey),
		zap.String("adapter", handle.Dispatch().Adapter),
		zap.String("target", handle.Dispatch().Target),
		zap.Bool("stream", req.Stream),
		zap.String("response_mode", responseModePassthrough),
		zap.Int("estimated_input_tokens", handle.Dispatch().EstimatedInputTokens),
		zap.String("context_class", handle.Dispatch().ContextClass),
		zap.String("queue_reason", handle.Dispatch().QueueReason),
	)
	return handle, true
}

// writeProviderTunnelResponse relays ordered tunnel frames to the HTTP caller.
// The response-start frame sets status/headers, body frames are written and
// flushed in order, and END terminates the response. Caller disconnect and
// wait timeout propagate cancellation to the Node cancel path. The caller
// supplies metricLabels so success/error/cancel usage is attributed to the
// right endpoint (Chat vs Responses) and model_group; the writer never
// recomputes labels from requestModel, which is reserved for model-echo
// rewriting only.
func (s *Server) writeProviderTunnelResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult, reqStream bool, requestModel string, metricLabels usageLabels) {
	frames := handle.Stream().Frames
	if frames == nil {
		writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream unavailable")
		return
	}
	flusher, _ := w.(http.Flusher)
	timer := time.NewTimer(handle.WaitTimeout())
	defer timer.Stop()

	assembler := &providerChatAssembler{streaming: reqStream}
	modelRewriter := newProviderModelRewriter(reqStream, requestModel)

	wroteHeader := false
	bodyBytes := 0
	var bufferedBody []byte
	rewriteModelEcho := modelRewriter != nil
	// metricStatus is the terminal status reported to usage metrics. It defaults
	// to error and is upgraded to success on END or downgraded to cancel on
	// caller give-up; provider usage observed from the passthrough body is
	// reported without ever mutating that body (SDD S05/D04).
	metricStatus := usageStatusError
	defer func() {
		obs := assembler.observation()
		s.logger.Info("openai chat completion passthrough closed",
			zap.String("run_id", handle.Dispatch().RunID),
			zap.Bool("wrote_header", wroteHeader),
			zap.Int("body_bytes", bodyBytes),
			zap.String("assembled_content", obs.Content),
			zap.String("assembled_reasoning", obs.Reasoning),
			zap.Strings("assembled_tool_calls", obs.ToolCallNames),
			zap.Int("assembled_tool_call_count", len(obs.ToolCallNames)),
		)
		emitUsageMetrics(metricLabels, metricStatus, assembler.usageObservation())
	}()

	for {
		select {
		case <-r.Context().Done():
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
			metricStatus = usageStatusForError(r.Context().Err())
			return
		case <-timer.C:
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
			metricStatus = usageStatusCancel
			if !wroteHeader {
				writeError(w, http.StatusBadGateway, "run_error", errRunTimedOut.Error())
			}
			return
		case frame, ok := <-frames:
			if !ok {
				if !wroteHeader {
					writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream closed before provider response")
				}
				return
			}
			switch frame.GetKind() {
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
				if wroteHeader {
					continue
				}
				copyProviderResponseHeaders(w.Header(), frame.GetHeaders())
				status := int(frame.GetStatusCode())
				if status == 0 {
					status = http.StatusOK
				}
				if status >= http.StatusBadRequest {
					rewriteModelEcho = false
				}
				if rewriteModelEcho {
					// The response body may change length when provider-served
					// model names are rewritten to caller-facing model aliases.
					w.Header().Del("Content-Length")
				}
				w.WriteHeader(status)
				wroteHeader = true
				if flusher != nil {
					flusher.Flush()
				}
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
				body := frame.GetBody()
				if len(body) == 0 {
					continue
				}
				if !wroteHeader {
					// Defensive: a body frame before response-start still reaches the
					// caller instead of being dropped.
					w.WriteHeader(http.StatusOK)
					wroteHeader = true
				}
				bodyBytes += len(body)
				assembler.Write(body)
				if rewriteModelEcho {
					if reqStream {
						body = modelRewriter.AppendStream(body)
					} else {
						bufferedBody = append(bufferedBody, body...)
						continue
					}
				}
				if len(body) == 0 {
					continue
				}
				if _, err := w.Write(body); err != nil {
					// The caller is gone mid-stream; propagate cancel to the Node.
					s.sendCancelRun(handle.Dispatch())
					metricStatus = usageStatusCancel
					return
				}
				if flusher != nil {
					flusher.Flush()
				}
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
				msg := frame.GetError()
				if msg == "" {
					msg = "provider tunnel failed"
				}
				if !wroteHeader {
					writeError(w, http.StatusBadGateway, "provider_tunnel_error", msg)
					return
				}
				// Status/headers are already committed; the response is truncated
				// and the caller observes the broken stream.
				s.logger.Warn("openai passthrough tunnel error after response start",
					zap.String("run_id", handle.Dispatch().RunID),
					zap.String("error", msg),
				)
				if rewriteModelEcho && !reqStream && len(bufferedBody) > 0 {
					body := modelRewriter.RewriteComplete(bufferedBody)
					_, _ = w.Write(body)
					if flusher != nil {
						flusher.Flush()
					}
				}
				return
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
				if !wroteHeader {
					writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel ended before provider response")
					return
				}
				if rewriteModelEcho {
					var body []byte
					if reqStream {
						body = modelRewriter.FlushStream()
					} else if len(bufferedBody) > 0 {
						body = modelRewriter.RewriteComplete(bufferedBody)
					}
					if len(body) > 0 {
						if _, err := w.Write(body); err != nil {
							s.sendCancelRun(handle.Dispatch())
							metricStatus = usageStatusCancel
							return
						}
						if flusher != nil {
							flusher.Flush()
						}
					}
				}
				metricStatus = usageStatusSuccess
				return
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
				// Usage frames are Edge-internal observation candidates; pure
				// passthrough records them for metrics but never merges them into
				// the response body (SDD D04).
				assembler.recordProtoUsage(frame.GetUsage())
				continue
			}
		}
	}
}

type providerModelRewriter struct {
	streaming bool
	model     string
	pending   []byte
}

func newProviderModelRewriter(streaming bool, model string) *providerModelRewriter {
	model = strings.TrimSpace(model)
	if model == "" {
		return nil
	}
	return &providerModelRewriter{streaming: streaming, model: model}
}

func (r *providerModelRewriter) AppendStream(chunk []byte) []byte {
	if r == nil || !r.streaming || len(chunk) == 0 {
		return chunk
	}
	r.pending = append(r.pending, chunk...)
	var out bytes.Buffer
	for {
		idx := bytes.IndexByte(r.pending, '\n')
		if idx < 0 {
			break
		}
		line := r.pending[:idx+1]
		out.Write(rewriteProviderSSEModelLine(line, r.model))
		r.pending = r.pending[idx+1:]
	}
	return out.Bytes()
}

func (r *providerModelRewriter) FlushStream() []byte {
	if r == nil || len(r.pending) == 0 {
		return nil
	}
	pending := r.pending
	r.pending = nil
	return rewriteProviderSSEModelLine(pending, r.model)
}

func (r *providerModelRewriter) RewriteComplete(body []byte) []byte {
	if r == nil || len(body) == 0 {
		return body
	}
	return rewriteProviderJSONModel(body, r.model)
}

func rewriteProviderSSEModelLine(line []byte, model string) []byte {
	body, ending := splitLineEnding(line)
	prefix, payload, ok := bytes.Cut(body, []byte(":"))
	if !ok || strings.TrimSpace(string(prefix)) != "data" {
		return line
	}
	payload = bytes.TrimSpace(payload)
	if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
		return line
	}
	rewritten := rewriteProviderJSONModel(payload, model)
	if bytes.Equal(rewritten, payload) {
		return line
	}
	out := make([]byte, 0, len("data: ")+len(rewritten)+len(ending))
	out = append(out, "data: "...)
	out = append(out, rewritten...)
	out = append(out, ending...)
	return out
}

func splitLineEnding(line []byte) ([]byte, []byte) {
	if len(line) == 0 || line[len(line)-1] != '\n' {
		return line, nil
	}
	if len(line) >= 2 && line[len(line)-2] == '\r' {
		return line[:len(line)-2], []byte("\r\n")
	}
	return line[:len(line)-1], []byte("\n")
}

func rewriteProviderJSONModel(body []byte, model string) []byte {
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(body, &raw); err != nil {
		return body
	}
	if _, ok := raw["model"]; !ok {
		return body
	}
	modelJSON, err := json.Marshal(model)
	if err != nil {
		return body
	}
	raw["model"] = modelJSON
	rewritten, err := json.Marshal(raw)
	if err != nil {
		return body
	}
	return rewritten
}

// providerAssembledObservation is the Edge-internal human-readable view of a
// provider Chat Completions/Responses body assembled by providerChatAssembler.
// It feeds Edge-local logging and usage metrics only and is never written to
// the caller response; pure passthrough relays provider bytes unmodified.
type providerAssembledObservation struct {
	Content       string   `json:"content,omitempty"`
	Reasoning     string   `json:"reasoning,omitempty"`
	ToolCallNames []string `json:"tool_call_names,omitempty"`
	BodyBytes     int      `json:"body_bytes"`
}

// providerChatAssembler accumulates a human-readable view of the provider
// Chat Completions response for the assembled observation. It
// tolerates non-JSON and partial payloads: whatever cannot be parsed simply
// yields an empty summary.
type providerChatAssembler struct {
	streaming          bool
	nonStreamingParsed bool
	bodyBytes          int
	pending            []byte
	content            strings.Builder
	reasoning          strings.Builder
	toolCallNames      []string
	// usage holds provider-reported token counts observed in the passthrough
	// body. It feeds Edge-internal metrics only.
	usage usageObservation
}

// providerChatDeltaEnvelope matches the provider fields the assembler
// observes; streaming uses delta, non-streaming uses message.
type providerChatDeltaEnvelope struct {
	Content          string `json:"content"`
	ReasoningContent string `json:"reasoning_content"`
	Reasoning        string `json:"reasoning"`
	ToolCalls        []struct {
		Function struct {
			Name string `json:"name"`
		} `json:"function"`
	} `json:"tool_calls"`
}

// providerUsageEnvelope matches the OpenAI-compatible usage object from
// both Chat Completions and Responses. Chat Completions uses prompt_tokens/
// completion_tokens; Responses uses input_tokens/output_tokens. The optional
// detail objects carry reasoning and cached-input tokens for both formats.
type providerUsageEnvelope struct {
	PromptTokens        int `json:"prompt_tokens"`
	CompletionTokens    int `json:"completion_tokens"`
	InputTokens         int `json:"input_tokens"`
	OutputTokens        int `json:"output_tokens"`
	PromptTokensDetails *struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`
	CompletionTokensDetails *struct {
		ReasoningTokens int `json:"reasoning_tokens"`
	} `json:"completion_tokens_details"`
	InputTokensDetails *struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"input_tokens_details"`
	OutputTokensDetails *struct {
		ReasoningTokens int `json:"reasoning_tokens"`
	} `json:"output_tokens_details"`
}

// recordUsage stores the latest provider-reported usage. The provider tunnel
// body is never mutated; this observation feeds Edge-internal metrics only
// (SDD S05/D04). Chat Completions keys (prompt_tokens/completion_tokens) take
// precedence; if absent, Responses keys (input_tokens/output_tokens) are used.
func (a *providerChatAssembler) recordUsage(u *providerUsageEnvelope) {
	if u == nil {
		return
	}
	// Prefer Chat Completions keys; fall back to Responses keys.
	if u.PromptTokens != 0 {
		a.usage.inputTokens = u.PromptTokens
		a.usage.outputTokens = u.CompletionTokens
		if d := u.PromptTokensDetails; d != nil {
			a.usage.cachedInputTokens = d.CachedTokens
		}
		if d := u.CompletionTokensDetails; d != nil {
			a.usage.reasoningTokens = d.ReasoningTokens
		}
		return
	}
	if u.InputTokens != 0 {
		a.usage.inputTokens = u.InputTokens
		a.usage.outputTokens = u.OutputTokens
		if d := u.InputTokensDetails; d != nil {
			a.usage.cachedInputTokens = d.CachedTokens
		}
		if d := u.OutputTokensDetails; d != nil {
			a.usage.reasoningTokens = d.ReasoningTokens
		}
	}
}

// recordProtoUsage stores usage carried on a provider tunnel USAGE frame. Like
// recordUsage it observes for metrics only and never alters the response body.
func (a *providerChatAssembler) recordProtoUsage(u *iop.Usage) {
	if u == nil {
		return
	}
	a.usage.inputTokens = int(u.GetInputTokens())
	a.usage.outputTokens = int(u.GetOutputTokens())
	a.usage.reasoningTokens = int(u.GetReasoningTokens())
	a.usage.cachedInputTokens = int(u.GetCachedInputTokens())
}

// usageObservation returns the observed token usage plus the observed reasoning
// character count used by auxiliary and estimated-token reasoning metrics.
func (a *providerChatAssembler) usageObservation() usageObservation {
	obs := a.usage
	obs.reasoningChars = a.reasoning.Len()
	return obs
}

func (a *providerChatAssembler) Write(chunk []byte) {
	a.bodyBytes += len(chunk)
	a.pending = append(a.pending, chunk...)
	if !a.streaming {
		return
	}
	for {
		idx := bytes.IndexByte(a.pending, '\n')
		if idx < 0 {
			return
		}
		line := strings.TrimRight(string(a.pending[:idx]), "\r")
		a.pending = a.pending[idx+1:]
		a.consumeSSELine(line)
	}
}

func (a *providerChatAssembler) consumeSSELine(line string) {
	payload, ok := strings.CutPrefix(line, "data:")
	if !ok {
		return
	}
	payload = strings.TrimSpace(payload)
	if payload == "" || payload == "[DONE]" {
		return
	}
	var chunk struct {
		Choices []struct {
			Delta providerChatDeltaEnvelope `json:"delta"`
		} `json:"choices"`
		Usage *providerUsageEnvelope `json:"usage"`
	}
	if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
		return
	}
	for _, choice := range chunk.Choices {
		a.consumeDelta(choice.Delta)
	}
	a.recordUsage(chunk.Usage)

	// Responses streaming: nested response.usage from events like
	// response.completed. This captures provider-reported token usage from
	// the Responses API SSE payload (SDD S05/D04).
	if a.usage.inputTokens == 0 && a.usage.outputTokens == 0 {
		var respEvent struct {
			Response struct {
				Usage *providerUsageEnvelope `json:"usage"`
			} `json:"response"`
		}
		if err := json.Unmarshal([]byte(payload), &respEvent); err == nil {
			if respEvent.Response.Usage != nil {
				a.recordUsage(respEvent.Response.Usage)
			}
		}
	}
}

func (a *providerChatAssembler) consumeDelta(delta providerChatDeltaEnvelope) {
	a.content.WriteString(delta.Content)
	a.reasoning.WriteString(delta.ReasoningContent)
	a.reasoning.WriteString(delta.Reasoning)
	for _, call := range delta.ToolCalls {
		if name := strings.TrimSpace(call.Function.Name); name != "" {
			a.toolCallNames = append(a.toolCallNames, name)
		}
	}
}

func (a *providerChatAssembler) observation() *providerAssembledObservation {
	if !a.streaming && !a.nonStreamingParsed {
		a.nonStreamingParsed = true
		var resp struct {
			Choices []struct {
				Message providerChatDeltaEnvelope `json:"message"`
			} `json:"choices"`
			Usage *providerUsageEnvelope `json:"usage"`
		}
		if err := json.Unmarshal(a.pending, &resp); err == nil {
			for _, choice := range resp.Choices {
				a.consumeDelta(choice.Message)
			}
			a.recordUsage(resp.Usage)
		}
	}
	return &providerAssembledObservation{
		Content:       a.content.String(),
		Reasoning:     a.reasoning.String(),
		ToolCallNames: a.toolCallNames,
		BodyBytes:     a.bodyBytes,
	}
}

// hopByHopResponseHeaders are transport-level headers owned by each hop; they
// are not copied from the provider response to the caller response.
var hopByHopResponseHeaders = map[string]struct{}{
	"Connection":          {},
	"Keep-Alive":          {},
	"Proxy-Authenticate":  {},
	"Proxy-Authorization": {},
	"Te":                  {},
	"Trailer":             {},
	"Transfer-Encoding":   {},
	"Upgrade":             {},
}

func copyProviderResponseHeaders(dst http.Header, headers map[string]string) {
	for k, v := range headers {
		if _, hop := hopByHopResponseHeaders[http.CanonicalHeaderKey(k)]; hop {
			continue
		}
		dst.Set(k, v)
	}
}

type streamToolTextFilter struct {
	pending string
}

func (f *streamToolTextFilter) Append(delta string) string {
	f.pending += delta
	if idx := firstStreamToolTextCandidateIndex(f.pending); idx >= 0 {
		out := strings.TrimRightFunc(f.pending[:idx], unicode.IsSpace)
		f.pending = f.pending[idx:]
		return out
	}
	flushLen := streamToolTextSafeFlushLen(f.pending)
	out := strings.TrimRightFunc(f.pending[:flushLen], unicode.IsSpace)
	f.pending = f.pending[len(out):]
	return out
}

func (f *streamToolTextFilter) Flush() string {
	out := f.pending
	f.pending = ""
	return out
}

func (f *streamToolTextFilter) FlushNonCandidate() string {
	if firstStreamToolTextCandidateIndex(f.pending) >= 0 {
		f.pending = ""
		return ""
	}
	// Even without a full candidate, chunk-boundary protection can leave a
	// partial candidate suffix (e.g. "<tool_ca", trailing "{") in pending.
	// Return only the safe non-candidate prefix and drop the partial suffix so
	// a native tool_calls completion never flushes a raw candidate fragment.
	flushLen := streamToolTextSafeFlushLen(f.pending)
	out := strings.TrimRightFunc(f.pending[:flushLen], unicode.IsSpace)
	f.pending = ""
	return out
}

func firstStreamToolTextCandidateIndex(s string) int {
	xml := strings.Index(strings.ToLower(s), "<tool_call")
	mustache := strings.Index(s, textMustacheToolCallOpenHint)
	switch {
	case xml < 0:
		return mustache
	case mustache < 0:
		return xml
	case xml < mustache:
		return xml
	default:
		return mustache
	}
}

func streamToolTextSafeFlushLen(s string) int {
	if s == "" {
		return 0
	}
	const xmlHint = "<tool_call"
	lower := strings.ToLower(s)
	start := len(s) - len(xmlHint) + 1
	if start < 0 {
		start = 0
	}
	for i := start; i < len(s); i++ {
		if strings.HasPrefix(xmlHint, lower[i:]) ||
			strings.HasPrefix(textMustacheToolCallOpenHint, s[i:]) {
			return i
		}
	}
	return len(s)
}

func unstreamedCleanedText(cleaned, emitted string) string {
	if cleaned == "" {
		return ""
	}
	if emitted == "" {
		return cleaned
	}
	if strings.HasPrefix(cleaned, emitted) {
		return cleaned[len(emitted):]
	}
	cleaned = strings.TrimSpace(cleaned)
	emitted = strings.TrimSpace(emitted)
	if cleaned == "" || cleaned == emitted {
		return ""
	}
	if emitted != "" && strings.HasPrefix(cleaned, emitted) {
		return strings.TrimSpace(strings.TrimPrefix(cleaned, emitted))
	}
	return ""
}

func writeContentDeltaSSE(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string, content string) {
	if content == "" {
		return
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Content: content},
		}},
	})
}

func writeToolCallsDeltaSSE(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string, toolCalls []any) {
	if len(toolCalls) == 0 {
		return
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{ToolCalls: toolCallsForStreamDelta(toolCalls)},
		}},
	})
}

// streamBufferedChatCompletion collects the full run output, validates tool
// calls against the request schema before emitting anything, and — when
// validation fails on a run that has not yet written a user-visible chunk —
// resubmits the same request as a bounded exact replay. Only a validated (or
// validation-disabled) result reaches the SSE writer; a malformed tool call
// that survives the retry budget is surfaced as an SSE error instead of a
// successful tool_calls chunk.
func (s *Server) streamBufferedChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, submitReq edgeservice.SubmitRunRequest, handle edgeservice.RunResult, flusher http.Flusher, outputPolicy strictOutputPolicy, validation toolValidationContract, retrySubmit func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error)) {
	attempt := 1
	metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModeNormalized)
	for {
		result, err := collectChatCompletionOutput(r.Context(), req, handle, outputPolicy)
		if err != nil {
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), err)
			handle.Close()
			emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
			writeSSEError(w, flusher, err.Error())
			return
		}
		verr := result.toolValidationErr
		if verr == nil {
			verr = validateToolCallResponse(validation, result.toolCalls, result.toolCallOrigin)
		}
		if verr != nil {
			failedRunID := handle.Dispatch().RunID
			if attempt < maxToolValidationAttempts {
				handle.Close()
				attempt++
				retryReq := submitReq
				retryReq.Metadata = toolValidationAttemptMetadata(submitReq.Metadata, attempt, failedRunID, verr.Error())
				s.logger.Warn("openai chat completion stream tool validation retry",
					zap.String("run_id", failedRunID),
					zap.Int("attempt", attempt),
					zap.String("reason", verr.Error()),
				)
				next, submitErr := retrySubmit(r.Context(), retryReq)
				if submitErr != nil {
					emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
					writeSSEErrorWithType(w, flusher, "tool_validation_retry_error", submitErr.Error())
					return
				}
				// Extract RunResult from either a direct RunResult or a
				// ProviderPoolDispatchResult (pool retry). For provider-pool
				// retries, only normalized results with a non-nil Run are
				// accepted; tunnel or nil Run results are rejected to prevent
				// nil handle panic downstream.
				if rr, ok := next.(edgeservice.RunResult); ok {
					handle = rr
				} else if ppd, ok := next.(*edgeservice.ProviderPoolDispatchResult); ok {
					if ppd.Path != edgeservice.ProviderPoolPathNormalized || ppd.Run == nil {
						if ppd.Tunnel != nil {
							ppd.Tunnel.Close()
						}
						writeSSEErrorWithType(w, flusher, "tool_validation_retry_error", "provider-pool retry selected tunnel path")
						return
					}
					handle = ppd.Run
				} else {
					handle.Close()
					writeSSEErrorWithType(w, flusher, "run_error", "unexpected retry result type")
					return
				}
				continue
			}
			handle.Close()
			s.logger.Warn("openai chat completion stream tool validation failed",
				zap.String("run_id", failedRunID),
				zap.Int("attempt", attempt),
				zap.String("reason", verr.Error()),
			)
			emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
			writeSSEErrorWithType(w, flusher, "tool_validation_error", verr.Error())
			return
		}
		handle.Close()
		s.writeBufferedStreamOutput(w, flusher, req, handle.Dispatch(), result, outputPolicy, openAICompatTraceStreamEnabled(submitReq.Metadata))
		emitUsageMetrics(metricLabels, usageStatusSuccess, usageObservationFromOpenAIUsage(result.usage, result.reasoningLen))
		return
	}
}

// writeBufferedStreamOutput emits a validated buffered result as SSE chunks:
// the role chunk, an optional content chunk, an optional tool_calls chunk, and
// the terminal finish chunk.
func (s *Server) writeBufferedStreamOutput(w http.ResponseWriter, flusher http.Flusher, req chatCompletionRequest, dispatch edgeservice.RunDispatch, result chatCompletionOutput, outputPolicy strictOutputPolicy, traceStream bool) {
	created := time.Now().Unix()
	id := "chatcmpl-" + dispatch.RunID
	model := responseModel(req.Model, dispatch.Target)
	content := result.message.Content
	toolCallNames := extractToolCallNames(result.toolCalls)
	s.logger.Info("openai chat completion stream closed",
		zap.String("run_id", dispatch.RunID),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("normalized", result.normalized),
		zap.String("response_mode", result.responseMode),
		zap.Int("content_len", len(content)),
		zap.Int("reasoning_len", result.reasoningLen),
		zap.String("assembled_content", content),
		zap.String("assembled_reasoning", result.message.ReasoningContent),
		zap.Strings("assembled_tool_calls", toolCallNames),
		zap.Int("assembled_tool_call_count", len(toolCallNames)),
	)
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})
	if traceStream {
		logOpenAICompatStreamOutput(s.logger, dispatch.RunID, 1, "content", content)
	}
	writeContentDeltaSSE(w, flusher, id, created, model, content)
	if traceStream && len(result.toolCalls) > 0 {
		s.logger.Info("openai chat completion stream output chunk",
			zap.String("run_id", dispatch.RunID),
			zap.Int("seq", 2),
			zap.String("type", "tool_calls"),
			zap.Int("tool_call_count", len(result.toolCalls)),
		)
	}
	writeToolCallsDeltaSSE(w, flusher, id, created, model, result.toolCalls)
	if traceStream {
		s.logger.Info("openai chat completion stream output done",
			zap.String("run_id", dispatch.RunID),
			zap.Int("seq", 3),
			zap.String("finish_reason", result.finishReason),
			zap.Int("content_len", len(content)),
			zap.Int("reasoning_len", result.reasoningLen),
		)
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index:        0,
			Delta:        chatDelta{},
			FinishReason: result.finishReason,
		}},
	})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}

func writeSSE(w http.ResponseWriter, flusher http.Flusher, v any) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	fmt.Fprintf(w, "data: %s\n\n", b)
	flusher.Flush()
}

func writeSSEError(w http.ResponseWriter, flusher http.Flusher, message string) {
	writeSSEErrorWithType(w, flusher, "run_error", message)
}

func writeSSEErrorWithType(w http.ResponseWriter, flusher http.Flusher, errType, message string) {
	writeSSE(w, flusher, errorResponse{Error: errorBody{Type: errType, Message: message}})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}

func openAICompatTraceStreamEnabled(metadata map[string]string) bool {
	if traceBool(os.Getenv(streamTraceEnvKey)) {
		return true
	}
	if metadata == nil {
		return false
	}
	return traceBool(metadata[streamTraceMetadataKey])
}

func traceBool(v string) bool {
	switch strings.ToLower(strings.TrimSpace(v)) {
	case "1", "true", "yes", "on":
		return true
	default:
		return false
	}
}

func logOpenAICompatStreamOutput(logger *zap.Logger, runID string, seq int, typ string, delta string, extra ...zap.Field) {
	fields := []zap.Field{
		zap.String("run_id", runID),
		zap.Int("seq", seq),
		zap.String("type", typ),
		zap.Int("delta_len", len(delta)),
		zap.Bool("delta_has_open_think", hasOpenThinkTag(delta)),
		zap.Bool("delta_has_close_think", hasCloseThinkTag(delta)),
	}
	fields = append(fields, extra...)
	logger.Info("openai chat completion stream output chunk", fields...)
}

func hasOpenThinkTag(s string) bool {
	return strings.Contains(strings.ToLower(s), "<think")
}

func hasCloseThinkTag(s string) bool {
	return strings.Contains(strings.ToLower(s), "</think")
}