iop/apps/edge/internal/openai/chat_handler.go

package openai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"go.uber.org/zap"

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

func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
		return
	}
	defer r.Body.Close()

	// The raw body is kept for the provider tunnel passthrough path, which
	// forwards the caller's own payload (model rewritten to the served target).
	rawBody, err := io.ReadAll(r.Body)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "failed to read request body")
		return
	}

	// The route decision runs before strict normalization: only the
	// routing-relevant envelope fields are decoded leniently so provider
	// passthrough can forward Codex/provider-specific payloads verbatim.
	// Strict field validation stays on the non-provider path.
	env, err := decodeChatCompletionEnvelope(rawBody)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	dispatch, ok := s.resolveRouteDispatch(env.Model)
	if !ok {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "model is required")
		return
	}

	var req chatCompletionRequest
	if dispatch.ProviderPool {
		// Provider-pool path: lenient decode so caller unknown/Codex/
		// provider-specific fields survive into the tunnel body.
		if err := decodeChatCompletionRequestLenient(json.NewDecoder(bytes.NewReader(rawBody)), &req); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
			return
		}
	} else {
		// Non-provider path keeps the strict allowlist-based decode.
		if err := decodeChatCompletionRequest(json.NewDecoder(bytes.NewReader(rawBody)), &req); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
			return
		}
	}
	runMeta, workspace, err := parseOpenAIMetadata(req.Metadata)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	// Overwrite (not merge-if-absent): the authenticated caller identity must
	// win over any caller-supplied metadata.iop_principal_* spoof attempt.
	for k, v := range principalMetadata(r.Context()) {
		runMeta[k] = v
	}
	if err := validateWorkspaceForRoute(dispatch, workspace); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	providerTunnelRoute := routeUsesProviderTunnel(dispatch)
	basePrompt := promptFromMessages(req.Messages)
	if strings.TrimSpace(basePrompt) == "" {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "messages are required")
		return
	}
	outputPolicy := s.resolveOutputPolicy(basePrompt)
	if catalogEntry := s.findProviderPoolEntry(req.Model); catalogEntry != nil {
		applyModelCatalogGenerationPolicyToChat(&req, *catalogEntry, outputPolicy.Strict)
	}

	// The response path is decided by the resolved route, never by caller
	// metadata: provider routes relay pure passthrough over the raw tunnel;
	// every other route uses the normalized RunEvent path. Caller metadata is
	// arbitrary context and does not select route or response shape (SDD S01).
	// Pre-compute estimate/contextClass early so provider-pool path can use them.
	estimate := s.estimateChatInputTokens(basePrompt, runMeta, req.Tools, req.ToolChoice)
	contextClass := classifyContext(estimate, s.longContextThreshold())

	// Provider-pool: use one-shot SubmitProviderPool which dispatches either
	// tunnel or normalized based on the selected provider's executionPath.
	if dispatch.ProviderPool {
		// Build input and messages for provider-pool dispatch.
		messages := req.Messages
		prompt := promptFromMessages(messages)
		input := req.runInput(prompt, messages, outputPolicy.Strict, routeSupportsNativeToolCalls(dispatch))
		validation := toolValidationContractFromRequest(req)
		metadata := chatRunMetadata(runMeta, req, outputPolicy)
		metadata["estimated_input_tokens"] = strconv.Itoa(estimate)
		metadata["context_class"] = contextClass

		// strict-output output policy only applies to normalized dispatch,
		// not to raw tunnel passthrough (SDD D02).
		s.handleChatCompletionsProviderPool(r, req, dispatch, workspace, basePrompt, prompt, rawBody, input, metadata, estimate, contextClass, outputPolicy, validation, w)
		return
	}

	if providerTunnelRoute {
		s.tunnelChatCompletionPassthrough(w, r, req, dispatch, runMeta, rawBody, estimate, contextClass)
		return
	}

	// Non-provider-pool normalized path: build messages, prompt, estimate.
	messages := req.Messages
	if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" {
		messages = prependSystemMessage(messages, instruction)
	}
	prompt := promptFromMessages(messages)
	input := req.runInput(prompt, messages, outputPolicy.Strict, routeSupportsNativeToolCalls(dispatch))
	// Re-compute estimate with instruction-augmented prompt for normalized path.
	estimate = s.estimateChatInputTokens(prompt, runMeta, req.Tools, req.ToolChoice)
	contextClass = classifyContext(estimate, s.longContextThreshold())

	validation := toolValidationContractFromRequest(req)
	metadata := chatRunMetadata(runMeta, req, outputPolicy)
	metadata["estimated_input_tokens"] = strconv.Itoa(estimate)
	metadata["context_class"] = contextClass
	if validation.enabled {
		metadata = toolValidationAttemptMetadata(metadata, 1, "", "")
	}
	submitReq := chatSubmitRunRequest(dispatch, req, workspace, prompt, input, metadata)
	submitReq.EstimatedInputTokens = estimate
	submitReq.ContextClass = contextClass
	metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModeNormalized)
	handle, err := s.service.SubmitRun(r.Context(), submitReq)
	if err != nil {
		emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
		writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
		return
	}
	s.logger.Info("openai chat completion 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.Int("estimated_input_tokens", handle.Dispatch().EstimatedInputTokens),
		zap.String("context_class", handle.Dispatch().ContextClass),
		zap.String("queue_reason", handle.Dispatch().QueueReason),
	)

	if req.Stream {
		retryFn := func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error) {
			return s.service.SubmitRun(ctx, req)
		}
		s.streamChatCompletion(w, r, req, submitReq, handle, outputPolicy, validation, retryFn)
		return
	}
	retryFn := func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error) {
		return s.service.SubmitRun(ctx, req)
	}
	s.completeChatCompletion(w, r, req, submitReq, handle, outputPolicy, validation, retryFn)
}

// handleChatCompletionsProviderPool dispatches a provider-pool chat completion
// through the one-shot SubmitProviderPool surface. After a single queue admission,
// the selected candidate's executionPath determines whether the response is
// delivered via tunnel passthrough or the normalized RunEvent path.
//
// Provider auth is validated via providerTunnelAuthHeaders which is called for
// tunnel-path only. For normalized paths, the auth check is skipped (SDD D01).
func (s *Server) handleChatCompletionsProviderPool(
	r *http.Request,
	req chatCompletionRequest,
	dispatch routeDispatch,
	workspace, basePrompt, prompt string,
	rawBody []byte,
	input map[string]any,
	metadata map[string]string,
	estimate int,
	contextClass string,
	outputPolicy strictOutputPolicy,
	validation toolValidationContract,
	w http.ResponseWriter,
) {
	poolReq := edgeservice.ProviderPoolDispatchRequest{
		Run: edgeservice.SubmitRunRequest{
			NodeRef:              dispatch.NodeRef,
			ModelGroupKey:        strings.TrimSpace(req.Model),
			SessionID:            dispatch.SessionID,
			Workspace:            workspace,
			Prompt:               prompt,
			Input:                input,
			TimeoutSec:           dispatch.TimeoutSec,
			MaxQueue:             dispatch.MaxQueue,
			QueueTimeoutMS:       dispatch.QueueTimeoutMS,
			Metadata:             metadata,
			EstimatedInputTokens: estimate,
			ContextClass:         contextClass,
			ProviderPool:         true,
		},
		Tunnel: edgeservice.SubmitProviderTunnelRequest{
			ModelGroupKey:        strings.TrimSpace(req.Model),
			SessionID:            dispatch.SessionID,
			Method:               http.MethodPost,
			Path:                 "/v1/chat/completions",
			Stream:               req.Stream,
			TimeoutSec:           dispatch.TimeoutSec,
			MaxQueue:             dispatch.MaxQueue,
			QueueTimeoutMS:       dispatch.QueueTimeoutMS,
			Metadata:             metadata,
			EstimatedInputTokens: estimate,
			ContextClass:         contextClass,
			ProviderPool:         true,
		},
	}

	// Pre-dispatch provider auth header injection. Runs inside SubmitProviderPool
	// BEFORE buildProviderTunnelRequest and the Node Send step, so the auth
	// header lands on the actual wire request. On failure the slot is released
	// and no tunnel request is sent (SDD S03).
	poolReq.PrepareTunnel = func(tunnelReq edgeservice.SubmitProviderTunnelRequest) (edgeservice.SubmitProviderTunnelRequest, error) {
		headers, err := s.providerTunnelAuthHeaders(r)
		if err != nil {
			return tunnelReq, err
		}
		tunnelReq.Headers = headers
		return tunnelReq, nil
	}

	// strict-output output policy only applies to normalized dispatch,
	// not to raw tunnel passthrough (SDD D02).
	poolReq.Tunnel.BuildBody = func(target string) ([]byte, error) {
		return rewriteChatCompletionModel(rawBody, target, req)
	}

	result, err := s.service.SubmitProviderPool(r.Context(), poolReq)
	if err != nil {
		metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModePassthrough)
		emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
		// Provider auth failure is a client request error (400), not a
		// backend dispatch error. The auth check now runs inside
		// SubmitProviderPool (PrepareTunnel) before any tunnel request is sent.
		if errors.Is(err, errProviderAuthRequired) {
			writeError(w, http.StatusBadRequest, "invalid_request_error", "provider auth token is required")
			return
		}
		writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
		return
	}

	s.logger.Info("openai chat completion provider-pool dispatch",
		zap.String("run_id", result.DispatchInfo.RunID),
		zap.String("node_id", result.DispatchInfo.NodeID),
		zap.String("provider_id", result.DispatchInfo.ProviderID),
		zap.String("provider_type", result.DispatchInfo.ProviderType),
		zap.String("execution_path", result.DispatchInfo.ExecutionPath),
		zap.String("model_group", result.DispatchInfo.ModelGroupKey),
		zap.String("adapter", result.DispatchInfo.Adapter),
		zap.String("target", result.DispatchInfo.Target),
		zap.String("path", string(result.Path)),
		zap.Int("estimated_input_tokens", result.DispatchInfo.EstimatedInputTokens),
		zap.String("context_class", result.DispatchInfo.ContextClass),
		zap.String("queue_reason", result.DispatchInfo.QueueReason),
	)

	// Prepare a retry submit function that preserves provider-pool context.
	// For normalized path, retry must call SubmitProviderPool (not SubmitRun)
	// to preserve ModelGroupKey, provider-pool metadata, and input.
	poolRetrySubmit := func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error) {
		return s.service.SubmitProviderPool(ctx, edgeservice.ProviderPoolDispatchRequest{
			Run:           req,
			Tunnel:        poolReq.Tunnel,
			PrepareTunnel: poolReq.PrepareTunnel,
		})
	}

	switch result.Path {
	case edgeservice.ProviderPoolPathTunnel:
		// Tunnel path: provider auth was already validated and injected via
		// PrepareTunnel before dispatch; on failure SubmitProviderPool returns
		// an error and no tunnel handle exists. Provider bytes are relayed as
		// pure passthrough; caller metadata never selects a sideband surface.
		metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModePassthrough)
		s.writeProviderTunnelResponse(w, r, result.Tunnel, req.Stream, req.Model, metricLabels)

	case edgeservice.ProviderPoolPathNormalized:
		// Normalized path: no auth required, collect from RunEvent stream.
		handle := result.Run
		if handle == nil {
			writeError(w, http.StatusInternalServerError, "run_error", "provider-pool selection returned normalized path but no run result")
			return
		}

		if req.Stream {
			s.streamChatCompletion(w, r, req, poolReq.Run, handle, outputPolicy, validation, poolRetrySubmit)
		} else {
			s.completeChatCompletion(w, r, req, poolReq.Run, handle, outputPolicy, validation, poolRetrySubmit)
		}
	}
}

func decodeChatCompletionRequest(dec *json.Decoder, req *chatCompletionRequest) error {
	var raw map[string]json.RawMessage
	if err := dec.Decode(&raw); err != nil {
		return fmt.Errorf("invalid JSON request")
	}
	for key := range raw {
		switch key {
		case "model", "messages", "stream", "metadata", "max_tokens", "max_completion_tokens", "temperature", "top_p", "presence_penalty", "frequency_penalty", "seed", "stop", "response_format", "tools", "tool_choice", "parallel_tool_calls", "stream_options", "store", "think", "reasoning_effort", "thinking_token_budget", "include_reasoning":
		default:
			return fmt.Errorf("%s is not supported for /v1/chat/completions", key)
		}
	}
	normalized, err := json.Marshal(raw)
	if err != nil {
		return fmt.Errorf("invalid JSON request")
	}
	if err := json.Unmarshal(normalized, req); err != nil {
		return fmt.Errorf("invalid /v1/chat/completions request format")
	}
	if req.MaxTokens != nil && *req.MaxTokens <= 0 {
		return fmt.Errorf("max_tokens must be greater than zero")
	}
	if req.MaxCompletionTokens != nil && *req.MaxCompletionTokens <= 0 {
		return fmt.Errorf("max_completion_tokens must be greater than zero")
	}
	if req.Temperature != nil && (*req.Temperature < 0 || *req.Temperature > 2) {
		return fmt.Errorf("temperature must be between 0 and 2")
	}
	if req.TopP != nil && (*req.TopP < 0 || *req.TopP > 1) {
		return fmt.Errorf("top_p must be between 0 and 1")
	}
	if err := validateThinkControl(req); err != nil {
		return err
	}
	return nil
}

// chatCompletionEnvelope holds the routing-relevant fields decoded leniently
// from a /v1/chat/completions request body before strict normalization. Unknown
// fields are ignored so the provider tunnel passthrough can forward Codex/
// provider-specific payloads verbatim. Strict field validation stays on the
// normalized non-provider path.
type chatCompletionEnvelope struct {
	Model    string          `json:"model"`
	Metadata json.RawMessage `json:"metadata,omitempty"`
	Stream   bool            `json:"stream"`
}

// decodeChatCompletionEnvelope leniently extracts only the routing-relevant
// fields (model, metadata, stream) from a /v1/chat/completions request body.
// It uses json.Decoder to read only the first JSON value so caller-provided
// trailing payload (provider-specific extensions beyond the root object) does
// not invalidate the envelope; strict field validation stays on the normalized
// non-provider path.
func decodeChatCompletionEnvelope(rawBody []byte) (chatCompletionEnvelope, error) {
	var env chatCompletionEnvelope
	dec := json.NewDecoder(bytes.NewReader(rawBody))
	if err := dec.Decode(&env); err != nil {
		return chatCompletionEnvelope{}, fmt.Errorf("invalid JSON request")
	}
	return env, nil
}

// decodeChatCompletionRequestLenient decodes a /v1/chat/completions request
// body while preserving unknown top-level fields. Known-field syntax and
// validation are still enforced so malformed standard fields cannot reach the
// provider tunnel. This is the provider-pool ingress path.
func decodeChatCompletionRequestLenient(dec *json.Decoder, req *chatCompletionRequest) error {
	var raw map[string]json.RawMessage
	if err := dec.Decode(&raw); err != nil {
		return fmt.Errorf("invalid JSON request")
	}
	for key := range raw {
		switch key {
		case "model", "messages", "stream", "metadata", "max_tokens", "max_completion_tokens", "temperature", "top_p", "presence_penalty", "frequency_penalty", "seed", "stop", "response_format", "tools", "tool_choice", "parallel_tool_calls", "stream_options", "store", "think", "reasoning_effort", "thinking_token_budget", "include_reasoning":
		default:
			// Unknown fields are tolerated for provider-pool passthrough.
		}
	}
	normalized, err := json.Marshal(raw)
	if err != nil {
		return fmt.Errorf("invalid JSON request")
	}
	if err := json.Unmarshal(normalized, req); err != nil {
		return fmt.Errorf("invalid /v1/chat/completions request format")
	}
	// Reuse the same validation rules on known fields for consistency.
	if req.MaxTokens != nil && *req.MaxTokens <= 0 {
		return fmt.Errorf("max_tokens must be greater than zero")
	}
	if req.MaxCompletionTokens != nil && *req.MaxCompletionTokens <= 0 {
		return fmt.Errorf("max_completion_tokens must be greater than zero")
	}
	if req.Temperature != nil && (*req.Temperature < 0 || *req.Temperature > 2) {
		return fmt.Errorf("temperature must be between 0 and 2")
	}
	if req.TopP != nil && (*req.TopP < 0 || *req.TopP > 1) {
		return fmt.Errorf("top_p must be between 0 and 1")
	}
	if err := validateThinkControl(req); err != nil {
		return err
	}
	return nil
}

func validateThinkControl(req *chatCompletionRequest) error {
	if req.ThinkingTokenBudget != nil && *req.ThinkingTokenBudget < 0 {
		return fmt.Errorf("thinking_token_budget must be non-negative")
	}
	if req.ReasoningEffort != nil {
		eff := *req.ReasoningEffort
		if eff == "" {
			return fmt.Errorf("reasoning_effort cannot be empty when present")
		}
		if eff != "none" && eff != "low" && eff != "medium" && eff != "high" {
			return fmt.Errorf("reasoning_effort must be one of none, low, medium, or high")
		}
	}
	if req.Think != nil && !*req.Think {
		if req.ReasoningEffort != nil && *req.ReasoningEffort != "none" {
			return fmt.Errorf("think=false conflicts with reasoning_effort=%s", *req.ReasoningEffort)
		}
	}
	if req.ThinkingTokenBudget != nil {
		if req.Think != nil && !*req.Think {
			return fmt.Errorf("thinking_token_budget cannot be set when think=false")
		}
		if req.ReasoningEffort != nil && *req.ReasoningEffort == "none" {
			return fmt.Errorf("thinking_token_budget cannot be set when reasoning_effort=none")
		}
	}
	return nil
}

func applyModelCatalogGenerationPolicyToChat(req *chatCompletionRequest, entry config.ModelCatalogEntry, strictOutput bool) {
	if req == nil {
		return
	}
	applyOutputTokenPolicy(&req.MaxTokens, req.MaxCompletionTokens, entry.DefaultMaxTokens, entry.MinMaxTokens)
	if entry.DefaultThinkingTokenBudget > 0 && chatRequestAllowsDefaultThinkingBudget(*req) {
		if req.ThinkingTokenBudget == nil {
			req.ThinkingTokenBudget = intPtr(entry.DefaultThinkingTokenBudget)
		}
		if strictOutput && req.Think == nil {
			req.Think = boolPtr(true)
		}
	}
}

func applyModelCatalogGenerationPolicyToResponses(req *responsesRequest, entry config.ModelCatalogEntry) {
	if req == nil {
		return
	}
	applyOutputTokenPolicy(&req.MaxOutputTokens, nil, entry.DefaultMaxTokens, entry.MinMaxTokens)
}

func applyOutputTokenPolicy(primary **int, fallback *int, defaultTokens, minTokens int) {
	if primary == nil {
		return
	}
	current := 0
	if *primary != nil {
		current = **primary
	} else if fallback != nil {
		current = *fallback
	}

	next := current
	if next == 0 && defaultTokens > 0 {
		next = defaultTokens
	}
	if minTokens > 0 && (next == 0 || next < minTokens) {
		next = minTokens
	}
	if next > 0 && next != current {
		*primary = intPtr(next)
	}
}

func chatRequestAllowsDefaultThinkingBudget(req chatCompletionRequest) bool {
	if req.Think != nil && !*req.Think {
		return false
	}
	if req.ReasoningEffort != nil && *req.ReasoningEffort == "none" {
		return false
	}
	return true
}

func intPtr(v int) *int {
	return &v
}

func boolPtr(v bool) *bool {
	return &v
}

func chatRunMetadata(runMeta map[string]string, req chatCompletionRequest, outputPolicy strictOutputPolicy) map[string]string {
	if runMeta == nil {
		runMeta = make(map[string]string)
	}
	runMeta["openai_model"] = req.Model
	runMeta["openai_stream"] = fmt.Sprintf("%t", req.Stream)
	runMeta["strict_output"] = fmt.Sprintf("%t", outputPolicy.Strict)
	return runMeta
}

const (
	// responseModePassthrough and responseModeNormalized are internal execution
	// labels for the response_mode usage metric. They are derived from the
	// handler execution path, never from caller metadata: provider tunnel routes
	// report passthrough, normalized RunEvent routes report normalized. Callers
	// cannot select a response mode through OpenAI metadata.
	responseModePassthrough = "passthrough"
	responseModeNormalized  = "normalized"
)

// routeUsesProviderTunnel reports whether the resolved dispatch targets an
// OpenAI-compatible provider that serves raw tunnel passthrough. Provider-pool
// catalog routes and openai_compat/vllm type routes qualify; CLI and other
// legacy adapters keep the normalized RunEvent path.
func routeUsesProviderTunnel(d routeDispatch) bool {
	if d.ProviderPool {
		return true
	}
	switch strings.TrimSpace(d.Adapter) {
	case "openai_compat", "vllm":
		return true
	default:
		return false
	}
}

// rewriteChatCompletionModel replaces only the model field of the caller's
// original request JSON so the provider receives its served model name. The
// rest of the caller payload is forwarded without IOP rewriting.
func rewriteChatCompletionModel(rawBody []byte, target string, req chatCompletionRequest) ([]byte, error) {
	if strings.TrimSpace(target) == "" && req.MaxTokens == nil && req.ThinkingTokenBudget == nil && req.Think == nil {
		return rawBody, nil
	}
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(rawBody, &raw); err != nil {
		return nil, fmt.Errorf("invalid JSON request")
	}
	if strings.TrimSpace(target) != "" {
		modelJSON, err := json.Marshal(target)
		if err != nil {
			return nil, err
		}
		raw["model"] = modelJSON
	}
	if req.MaxTokens != nil {
		maxTokensJSON, err := json.Marshal(*req.MaxTokens)
		if err != nil {
			return nil, err
		}
		raw["max_tokens"] = maxTokensJSON
		delete(raw, "max_completion_tokens")
	}
	if req.ThinkingTokenBudget != nil {
		budgetJSON, err := json.Marshal(*req.ThinkingTokenBudget)
		if err != nil {
			return nil, err
		}
		raw["thinking_token_budget"] = budgetJSON
	}
	if req.Think != nil {
		thinkJSON, err := json.Marshal(*req.Think)
		if err != nil {
			return nil, err
		}
		raw["think"] = thinkJSON
	}
	return json.Marshal(raw)
}

func routeSupportsNativeToolCalls(dispatch routeDispatch) bool {
	if dispatch.ProviderPool {
		return true
	}
	switch strings.TrimSpace(dispatch.Adapter) {
	case "openai_compat", "vllm", "ollama":
		return true
	default:
		return false
	}
}

func shouldSynthesizeTextToolCalls(dispatch edgeservice.RunDispatch, textToolFallback bool) bool {
	return textToolFallback || strings.TrimSpace(dispatch.Adapter) == "cli"
}

func chatSubmitRunRequest(dispatch routeDispatch, req chatCompletionRequest, workspace, prompt string, input map[string]any, metadata map[string]string) edgeservice.SubmitRunRequest {
	return edgeservice.SubmitRunRequest{
		NodeRef:        dispatch.NodeRef,
		ModelGroupKey:  strings.TrimSpace(req.Model),
		Adapter:        dispatch.Adapter,
		Target:         dispatch.Target,
		SessionID:      dispatch.SessionID,
		Workspace:      workspace,
		Prompt:         prompt,
		Input:          input,
		TimeoutSec:     dispatch.TimeoutSec,
		MaxQueue:       dispatch.MaxQueue,
		QueueTimeoutMS: dispatch.QueueTimeoutMS,
		Metadata:       metadata,
		ProviderPool:   dispatch.ProviderPool,
	}
}

func (s *Server) completeChatCompletion(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)) {
	attempt := 1
	for {
		result, err := collectChatCompletionOutput(r.Context(), req, handle, outputPolicy)
		if err != nil {
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), err)
			handle.Close()
			emitUsageMetrics(
				s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModeNormalized),
				usageStatusForError(err), usageObservation{},
			)
			writeError(w, httpStatusForRunError(err), "run_error", err.Error())
			return
		}
		valErr := result.toolValidationErr
		if valErr == nil {
			valErr = validateToolCallResponse(validation, result.toolCalls, result.toolCallOrigin)
		}
		if valErr != nil {
			failedRunID := handle.Dispatch().RunID
			if attempt < maxToolValidationAttempts {
				handle.Close()
				attempt++
				retryReq := submitReq
				retryReq.Metadata = toolValidationAttemptMetadata(submitReq.Metadata, attempt, failedRunID, valErr.Error())
				s.logger.Warn("openai chat completion tool validation retry",
					zap.String("run_id", failedRunID),
					zap.Int("attempt", attempt),
					zap.String("reason", valErr.Error()),
				)
				next, submitErr := retrySubmit(r.Context(), retryReq)
				if submitErr != nil {
					emitUsageMetrics(
						s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModeNormalized),
						usageStatusError, usageObservation{},
					)
					writeError(w, http.StatusBadGateway, "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()
						}
						writeError(w, http.StatusBadGateway, "tool_validation_retry_error", "provider-pool retry selected tunnel path")
						return
					}
					handle = ppd.Run
				} else {
					handle.Close()
					writeError(w, http.StatusInternalServerError, "run_error", "unexpected retry result type")
					return
				}
				continue
			}
			handle.Close()
			s.logger.Warn("openai chat completion tool validation failed",
				zap.String("run_id", failedRunID),
				zap.Int("attempt", attempt),
				zap.String("reason", valErr.Error()),
			)
			emitUsageMetrics(
				s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModeNormalized),
				usageStatusError, usageObservation{},
			)
			writeError(w, http.StatusBadGateway, "tool_validation_error", valErr.Error())
			return
		}
		handle.Close()
		emitUsageMetrics(
			s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, result.responseMode),
			usageStatusSuccess, usageObservationFromOpenAIUsage(result.usage, result.reasoningLen),
		)
		toolCallNames := extractToolCallNames(result.toolCalls)
		s.logger.Info("openai chat completion output",
			zap.String("run_id", handle.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", result.contentLen),
			zap.Int("reasoning_len", result.reasoningLen),
			zap.String("assembled_content", result.message.Content),
			zap.String("assembled_reasoning", result.message.ReasoningContent),
			zap.Strings("assembled_tool_calls", toolCallNames),
			zap.Int("assembled_tool_call_count", len(toolCallNames)),
		)
		created := time.Now().Unix()
		writeJSON(w, http.StatusOK, chatCompletionResponse{
			ID:      "chatcmpl-" + handle.Dispatch().RunID,
			Object:  "chat.completion",
			Created: created,
			Model:   responseModel(req.Model, handle.Dispatch().Target),
			Choices: []chatCompletionChoice{{
				Index:        0,
				Message:      result.message,
				FinishReason: result.finishReason,
			}},
			Usage: result.usage,
		})
		return
	}
}

type chatCompletionOutput struct {
	message           chatMessage
	finishReason      string
	usage             *openAIUsage
	normalized        bool
	responseMode      string
	contentLen        int
	reasoningLen      int
	toolCalls         []any
	toolCallOrigin    string
	toolValidationErr error
}

func collectChatCompletionOutput(ctx context.Context, req chatCompletionRequest, handle edgeservice.RunResult, outputPolicy strictOutputPolicy) (chatCompletionOutput, error) {
	text, reasoning, finishReason, nativeToolCalls, usage, _, err := collectRunResult(ctx, handle.Stream(), handle.WaitTimeout())
	if err != nil {
		return chatCompletionOutput{}, err
	}
	text, reasoning, normalized := normalizeCompletionOutput(outputPolicy, text, reasoning, req.explicitlyIncludesReasoning())
	message := chatMessage{Role: "assistant", Content: text, ReasoningContent: reasoning}
	includeReasoning := req.includeReasoning()
	if !includeReasoning {
		message.ReasoningContent = ""
	}
	var toolCalls []any
	toolCallOrigin := ""
	var toolValidationErr error

	if len(nativeToolCalls) > 0 {
		if len(req.Tools) > 0 && strings.TrimSpace(message.Content) != "" {
			filter := &streamToolTextFilter{}
			message.Content = filter.Append(message.Content) + filter.FlushNonCandidate()
		}
		nativeToolCalls = normalizeNativeToolCallsForResponse(nativeToolCalls, req.Tools)
		message.ToolCalls = nativeToolCalls
		toolCalls = nativeToolCalls
		toolCallOrigin = toolCallOriginNative
		finishReason = "tool_calls"
	} else if len(req.Tools) > 0 {
		res := synthesizeToolCallsFromTextResult(text, req.Tools, handle.Dispatch().RunID)
		if res.validationErr != nil {
			toolValidationErr = res.validationErr
		} else if len(res.toolCalls) > 0 {
			message.Content = res.cleaned
			message.ReasoningContent = ""
			message.ToolCalls = res.toolCalls
			toolCalls = res.toolCalls
			toolCallOrigin = toolCallOriginText
			finishReason = "tool_calls"
		}
	}
	if len(message.ToolCalls) == 0 && strings.TrimSpace(message.Content) == "" && strings.TrimSpace(reasoning) != "" {
		if includeReasoning {
			message.Content = reasoningOnlyFallbackContent(reasoning, finishReason)
		} else {
			message.Content = hiddenReasoningFallbackContent(finishReason)
		}
	}

	return chatCompletionOutput{
		message:           message,
		finishReason:      finishReason,
		usage:             usage,
		normalized:        normalized,
		responseMode:      responseModeNormalized,
		contentLen:        len(message.Content),
		reasoningLen:      len(reasoning),
		toolCalls:         toolCalls,
		toolCallOrigin:    toolCallOrigin,
		toolValidationErr: toolValidationErr,
	}, nil
}

func reasoningOnlyFallbackContent(reasoning, finishReason string) string {
	content := strings.TrimSpace(reasoning)
	if content == "" {
		return ""
	}
	reason := strings.TrimSpace(finishReason)
	if reason == "" || reason == "stop" {
		return content
	}
	return content + "\n\n[IOP notice: model finished before producing final assistant content; finish_reason=" + reason + "]"
}

func hiddenReasoningFallbackContent(finishReason string) string {
	reason := strings.TrimSpace(finishReason)
	if reason == "" {
		reason = "unknown"
	}
	return "[IOP notice: model produced reasoning but no final assistant content; reasoning was hidden by include_reasoning=false; finish_reason=" + reason + "]"
}

func (s *Server) resolveAdapter() string {
	if s.cfg.Adapter != "" {
		return s.cfg.Adapter
	}
	return "ollama"
}

func (s *Server) resolveTarget(model string) string {
	if s.cfg.Target != "" {
		return s.cfg.Target
	}
	return strings.TrimSpace(model)
}

// routeDispatch holds fully-resolved dispatch parameters for a single request.
type routeDispatch struct {
	NodeRef           string
	Adapter           string
	Target            string
	SessionID         string
	TimeoutSec        int
	MaxQueue          int
	QueueTimeoutMS    int
	WorkspaceRequired bool
	// ProviderPool is true when the request model matched a provider-pool
	// catalog entry. Adapter and Target are empty; the service layer resolves
	// them per-candidate and rewrites Target after admission.
	ProviderPool bool
}

// resolveRoute returns the first catalog entry whose Model matches model.
// Entries with an empty Target are skipped.
func (s *Server) resolveRoute(model string) *config.OpenAIRouteEntry {
	model = strings.TrimSpace(model)
	if model == "" {
		return nil
	}
	for i := range s.cfg.ModelRoutes {
		r := &s.cfg.ModelRoutes[i]
		if strings.TrimSpace(r.Model) == model && r.Target != "" {
			return r
		}
	}
	return nil
}

// findProviderPoolEntry returns the catalog entry matching model, or nil.
func (s *Server) findProviderPoolEntry(model string) *config.ModelCatalogEntry {
	model = strings.TrimSpace(model)
	if model == "" {
		return nil
	}
	modelCatalog := s.modelCatalogSnapshot()
	for i := range modelCatalog {
		if modelCatalog[i].ID == model {
			entry := modelCatalog[i]
			return &entry
		}
	}
	return nil
}

// resolveRouteDispatch returns fully-resolved dispatch params for model.
// Priority: provider-pool catalog → legacy model_routes → single-target fallback.
// Returns (dispatch, true) on success; (zero, false) when no target can be resolved.
func (s *Server) resolveRouteDispatch(model string) (routeDispatch, bool) {
	// Provider-pool catalog takes highest priority.
	if s.findProviderPoolEntry(model) != nil {
		return routeDispatch{
			SessionID:    s.resolveSessionID(),
			TimeoutSec:   s.resolveTimeoutSec(),
			ProviderPool: true,
		}, true
	}

	if route := s.resolveRoute(model); route != nil {
		adapter := route.Adapter
		if adapter == "" {
			adapter = s.resolveAdapter()
		}
		nodeRef := route.NodeRef
		if nodeRef == "" {
			nodeRef = s.cfg.NodeRef
		}
		sessionID := route.SessionID
		if sessionID == "" {
			sessionID = s.resolveSessionID()
		}
		timeoutSec := route.TimeoutSec
		if timeoutSec <= 0 {
			timeoutSec = s.resolveTimeoutSec()
		}
		return routeDispatch{
			NodeRef:           nodeRef,
			Adapter:           adapter,
			Target:            route.Target,
			SessionID:         sessionID,
			TimeoutSec:        timeoutSec,
			MaxQueue:          route.MaxQueue,
			QueueTimeoutMS:    route.QueueTimeoutMS,
			WorkspaceRequired: route.WorkspaceRequired,
		}, true
	}
	target := s.resolveTarget(model)
	if target == "" {
		return routeDispatch{}, false
	}
	return routeDispatch{
		NodeRef:    s.cfg.NodeRef,
		Adapter:    s.resolveAdapter(),
		Target:     target,
		SessionID:  s.resolveSessionID(),
		TimeoutSec: s.resolveTimeoutSec(),
	}, true
}

func (s *Server) resolveSessionID() string {
	if s.cfg.SessionID != "" {
		return s.cfg.SessionID
	}
	return edgeservice.DefaultSessionID
}

func (s *Server) resolveTimeoutSec() int {
	if s.cfg.TimeoutSec > 0 {
		return s.cfg.TimeoutSec
	}
	return edgeservice.DefaultTimeoutSec
}

func (s *Server) resolveStrictOutput() bool {
	return s.cfg.StrictOutput
}

func (s *Server) resolveStrictStreamBuffer() bool {
	return s.cfg.StrictStreamBuffer
}

func (s *Server) resolveOutputPolicy(prompt string) strictOutputPolicy {
	policy := strictOutputPolicy{
		Strict:       s.resolveStrictOutput(),
		StreamBuffer: s.resolveStrictStreamBuffer(),
	}
	if !policy.Strict {
		return policy
	}
	policy.XMLCompletionTool, policy.XMLResultTag = inferXMLCompletionContract(prompt)
	policy.ContractInstruction = policy.XMLCompletionTool != ""
	return policy
}

func promptFromMessages(messages []chatMessage) string {
	var b strings.Builder
	for _, msg := range messages {
		content := strings.TrimSpace(msg.Content)
		if content == "" {
			continue
		}
		role := strings.TrimSpace(msg.Role)
		if role == "" {
			role = "user"
		}
		if b.Len() > 0 {
			b.WriteString("\n")
		}
		b.WriteString(role)
		b.WriteString(": ")
		b.WriteString(content)
	}
	return b.String()
}

func responseModel(requestModel, target string) string {
	if requestModel != "" {
		return requestModel
	}
	return target
}

func httpStatusForRunError(err error) int {
	if errors.Is(err, context.Canceled) {
		return http.StatusRequestTimeout
	}
	return http.StatusBadGateway
}

func validateWorkspaceForRoute(d routeDispatch, workspace string) error {
	if !d.WorkspaceRequired {
		return nil
	}
	if strings.TrimSpace(workspace) == "" {
		return fmt.Errorf("workspace is required for this model route")
	}
	if !filepath.IsAbs(workspace) {
		return fmt.Errorf("workspace must be an absolute path")
	}
	return nil
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, status int, code, message string) {
	writeJSON(w, status, errorResponse{Error: errorBody{Type: code, Message: message}})
}

// estimateChatInputTokens computes a conservative token-count approximation for
// a chat completion request.  The prompt argument is the already-computed
// payload (including any strict-output contract instruction) so the estimate
// does not double-count the contract text that is also sent to the node.
func (s *Server) estimateChatInputTokens(prompt string, runMeta map[string]string, tools []any, toolChoice any) int {
	textParts := []string{prompt}
	// Metadata keys+values.
	for k, v := range runMeta {
		textParts = append(textParts, k, v)
	}
	// Tool schemas.
	for _, t := range tools {
		textParts = append(textParts, toJSON(t))
	}
	if toolChoice != nil {
		textParts = append(textParts, toJSON(toolChoice))
	}
	input := strings.Join(textParts, "\n")
	return estimateInputTokens(input, runMeta, tools, toolChoice)
}

// toJSON is a lightweight serialiser for any json.RawMessage-compatible value.
func toJSON(v any) string {
	b, err := json.Marshal(v)
	if err != nil {
		return ""
	}
	return string(b)
}