package openai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"go.uber.org/zap"
	"io"
	edgeservice "iop/apps/edge/internal/service"
	"net/http"
	"strconv"
	"strings"
)

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 := resolveCallerIdentity(r, dispatch, req.Metadata)
	if 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)
	// The mutable model-catalog generation policy is applied here, at the single
	// ingress point, before req is frozen into the dispatch context. Every stage
	// downstream reads the post-policy request.
	if catalogEntry := s.findProviderPoolEntry(req.Model); catalogEntry != nil {
		applyModelCatalogGenerationPolicyToChat(&req, *catalogEntry, outputPolicy.Strict, chatRequestHasProviderNativeThinking(rawBody))
	}

	// 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())

	requestCtx := openAIRequestContext{
		r:              r,
		route:          dispatch,
		rawBody:        rawBody,
		callerMetadata: runMeta,
		workspace:      workspace,
		estimate:       estimate,
		contextClass:   contextClass,
		endpoint:       usageEndpointChatCompletions,
	}

	// Provider-pool: use one-shot SubmitProviderPool which dispatches either
	// tunnel or normalized based on the selected provider's executionPath.
	dc := s.newChatDispatchContext(requestCtx, req, basePrompt, outputPolicy)
	if dispatch.ProviderPool {
		s.handleChatCompletionsProviderPool(w, dc)
		return
	}

	if providerTunnelRoute {
		s.tunnelChatCompletionPassthrough(w, dc)
		return
	}

	metricLabels := dc.usageLabels(s, responseModeNormalized)
	handle, err := s.service.SubmitRun(r.Context(), dc.submitReq)
	if err != nil {
		emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
		writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
		return
	}
	s.logChatDispatch("openai chat completion dispatch", handle.Dispatch())

	dc = dc.withRetrySubmit(func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error) {
		return s.service.SubmitRun(ctx, req)
	})
	if req.Stream {
		s.streamChatCompletion(w, dc, handle)
		return
	}
	s.completeChatCompletion(w, dc, handle)
}

// newChatDispatchContext freezes the ingress-resolved values into the immutable
// dispatch context for the resolved route. The provider-pool route leaves
// adapter/target unset: the service layer resolves them per candidate at
// admission and rewrites the target to the winning provider's served model.
func (s *Server) newChatDispatchContext(requestCtx openAIRequestContext, req chatCompletionRequest, basePrompt string, outputPolicy strictOutputPolicy) *chatDispatchContext {
	dc := &chatDispatchContext{
		openAIRequestContext: requestCtx,
		req:                  req,
		basePrompt:           basePrompt,
		outputPolicy:         outputPolicy,
		validation:           toolValidationContractFromRequest(req),
	}

	messages := req.Messages
	if !requestCtx.route.ProviderPool {
		// Strict output is a normalized-dispatch contract only; the pool route
		// decides the path after admission, so the instruction is not baked in
		// here (SDD D02).
		if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" {
			messages = prependSystemMessage(messages, instruction)
		}
	}
	dc.prompt = promptFromMessages(messages)
	dc.input = req.runInput(dc.prompt, messages, outputPolicy.Strict, routeSupportsNativeToolCalls(requestCtx.route))

	if !requestCtx.route.ProviderPool {
		// Re-estimate against the instruction-augmented prompt for the
		// normalized path.
		dc.estimate = s.estimateChatInputTokens(dc.prompt, requestCtx.callerMetadata, req.Tools, req.ToolChoice)
		dc.contextClass = classifyContext(dc.estimate, s.longContextThreshold())
	}

	dc.runMetadata = chatRunMetadata(cloneMetadata(requestCtx.callerMetadata), req, outputPolicy)
	dc.runMetadata["estimated_input_tokens"] = strconv.Itoa(dc.estimate)
	dc.runMetadata["context_class"] = dc.contextClass

	if requestCtx.route.ProviderPool {
		dc.submitReq = edgeservice.SubmitRunRequest{
			NodeRef:              requestCtx.route.NodeRef,
			ModelGroupKey:        strings.TrimSpace(req.Model),
			SessionID:            requestCtx.route.SessionID,
			Workspace:            requestCtx.workspace,
			Prompt:               dc.prompt,
			Input:                dc.input,
			TimeoutSec:           requestCtx.route.TimeoutSec,
			MaxQueue:             requestCtx.route.MaxQueue,
			QueueTimeoutMS:       requestCtx.route.QueueTimeoutMS,
			Metadata:             dc.runMetadata,
			EstimatedInputTokens: dc.estimate,
			ContextClass:         dc.contextClass,
			ProviderPool:         true,
		}
		return dc
	}

	if dc.validation.enabled {
		dc.runMetadata = toolValidationAttemptMetadata(dc.runMetadata, 1, "", "")
	}
	dc.submitReq = chatSubmitRunRequest(requestCtx.route, req, requestCtx.workspace, dc.prompt, dc.input, dc.runMetadata)
	dc.submitReq.EstimatedInputTokens = dc.estimate
	dc.submitReq.ContextClass = dc.contextClass
	return dc
}

func (s *Server) logChatDispatch(msg string, disp edgeservice.RunDispatch, extra ...zap.Field) {
	fields := []zap.Field{
		zap.String("run_id", disp.RunID),
		zap.String("node_id", disp.NodeID),
		zap.String("provider_id", disp.ProviderID),
		zap.String("provider_type", disp.ProviderType),
		zap.String("execution_path", disp.ExecutionPath),
		zap.String("model_group", disp.ModelGroupKey),
		zap.String("adapter", disp.Adapter),
		zap.String("target", disp.Target),
		zap.Int("estimated_input_tokens", disp.EstimatedInputTokens),
		zap.String("context_class", disp.ContextClass),
		zap.String("queue_reason", disp.QueueReason),
	}
	s.logger.Info(msg, append(fields, extra...)...)
}

// 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(w http.ResponseWriter, dc *chatDispatchContext) {
	r := dc.r
	req := dc.req
	poolReq := edgeservice.ProviderPoolDispatchRequest{
		Run: dc.submitReq,
		Tunnel: edgeservice.SubmitProviderTunnelRequest{
			ModelGroupKey:        strings.TrimSpace(req.Model),
			SessionID:            dc.route.SessionID,
			Method:               http.MethodPost,
			Path:                 "/v1/chat/completions",
			Stream:               req.Stream,
			TimeoutSec:           dc.route.TimeoutSec,
			MaxQueue:             dc.route.MaxQueue,
			QueueTimeoutMS:       dc.route.QueueTimeoutMS,
			Metadata:             dc.runMetadata,
			EstimatedInputTokens: dc.estimate,
			ContextClass:         dc.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). The caller's raw body is the
	// passthrough source; only the model field is rewritten to the served target.
	poolReq.Tunnel.BuildBody = func(target string) ([]byte, error) {
		return rewriteChatCompletionModel(dc.rawBody, target, req)
	}

	// Normalized branch: chat completions use lenient decode so unknown
	// fields survive; PrepareRun passes through the already-built Run
	// without strict validation (SDD S02/S03). The tunnel branch is
	// unaffected.
	poolReq.PrepareRun = func(runReq edgeservice.SubmitRunRequest) (edgeservice.SubmitRunRequest, error) {
		return runReq, nil
	}

	result, err := s.service.SubmitProviderPool(r.Context(), poolReq)
	if err != nil {
		emitUsageMetrics(dc.usageLabels(s, responseModePassthrough), 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.logChatDispatch("openai chat completion provider-pool dispatch", result.DispatchInfo,
		zap.String("path", string(result.Path)),
	)

	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.
		s.writeProviderTunnelResponse(w, r, result.Tunnel, req.Stream, req.Model, dc.usageLabels(s, responseModePassthrough))

	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
		}

		// Retry must re-enter SubmitProviderPool (not SubmitRun) so a bounded
		// tool-validation replay keeps the ModelGroupKey, provider-pool
		// metadata, and input of the original dispatch.
		poolDC := dc.withRetrySubmit(func(ctx context.Context, retryReq edgeservice.SubmitRunRequest) (any, error) {
			return s.service.SubmitProviderPool(ctx, edgeservice.ProviderPoolDispatchRequest{
				Run:           retryReq,
				Tunnel:        poolReq.Tunnel,
				PrepareTunnel: poolReq.PrepareTunnel,
				PrepareRun:    poolReq.PrepareRun,
			})
		})
		if req.Stream {
			s.streamChatCompletion(w, poolDC, handle)
		} else {
			s.completeChatCompletion(w, poolDC, handle)
		}
	}
}
