package openai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"time"

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

type edgeserviceRunner interface {
	SubmitProviderPool(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error)
}

// singleRequestProviderStage is the private, fail-closed provider codec used by
// the fixed single-request stage drivers.
type singleRequestProviderStage struct{ service edgeserviceRunner }

func newSingleRequestProviderStage(service edgeserviceRunner) *singleRequestProviderStage {
	return &singleRequestProviderStage{service: service}
}

type singleRequestProviderStageRequest struct {
	StageBinding     edgeservice.SingleRequestStageBinding
	Limits           edgeservice.SingleRequestLimits
	Messages         []chatMessage
	ResponseFormat   *singleRequestProviderResponseFormat
	NodeRef          string
	SessionID        string
	UsageAttribution string
	Quality          *singleRequestQualityGate
}

type singleRequestProviderStageResponse struct {
	Output   string
	Dispatch edgeservice.RunDispatch
}

type singleRequestProviderBodyBuilder func(string) ([]byte, error)

var (
	errProviderStageMissingBinding = errors.New("provider stage: missing stage binding")
	errProviderStageMissingInput   = errors.New("provider stage: missing messages")
	errProviderStageGeneric        = errors.New("provider stage: failed")
	errProviderStageMalformed      = errors.New("provider stage: malformed response")
	errProviderStageOutputLimit    = errors.New("provider stage: output limit")
	errProviderStageContextLimit   = errors.New("provider stage: context limit")
)

func (s *singleRequestProviderStage) submit(ctx context.Context, req singleRequestProviderStageRequest) (*singleRequestProviderStageResponse, error) {
	quality := singleRequestQualityGateOrNew(req.Quality)
	if s == nil || s.service == nil || req.StageBinding.Dispatch == nil || req.Limits.StageTimeoutMS < 1 || req.Limits.MaxOutputBytes < 1 {
		return nil, quality.validation(errProviderStageMissingBinding)
	}
	if len(req.Messages) == 0 {
		return nil, quality.validation(errProviderStageMissingInput)
	}
	dispatch := req.StageBinding.Dispatch
	stageCtx, cancel := providerStageContext(ctx, req.Limits.StageTimeoutMS)
	defer cancel()

	requirements := singleRequestProviderRequirements(req.StageBinding.Options, nil, req.ResponseFormat)
	poolReq := edgeservice.ProviderPoolDispatchRequest{
		Run: edgeservice.SubmitRunRequest{NodeRef: req.NodeRef, ModelGroupKey: dispatch.ModelGroupKey, ProviderID: dispatch.ProviderID, UsageAttribution: req.UsageAttribution, SessionID: req.SessionID, TimeoutSec: dispatch.TimeoutSec, MaxQueue: dispatch.MaxQueue, QueueTimeoutMS: dispatch.QueueTimeoutMS, ProviderPool: true},
		Tunnel: edgeservice.SubmitProviderTunnelRequest{
			CredentialBinding: dispatch.CredentialBindingSnapshot(), NodeRef: req.NodeRef, ModelGroupKey: dispatch.ModelGroupKey, ProviderID: dispatch.ProviderID, UsageAttribution: req.UsageAttribution,
			Adapter: "openai_compat", Target: dispatch.UpstreamModel, SessionID: req.SessionID, Method: http.MethodPost, Path: "/v1/chat/completions", Stream: false,
			TimeoutSec: dispatch.TimeoutSec, MaxQueue: dispatch.MaxQueue, QueueTimeoutMS: dispatch.QueueTimeoutMS, ProviderPool: true,
			BuildBody: func(target string) ([]byte, error) {
				return buildSingleRequestChatBody(req.Messages, req.StageBinding.Options, req.ResponseFormat, target)
			},
		},
		AcceptCandidate: singleRequestProviderCandidatePredicate(dispatch, requirements),
	}
	poolReq.PrepareProtocolTunnel = singleRequestProviderTunnelPreparer(requirements, func(target string) ([]byte, error) {
		return buildSingleRequestChatBody(req.Messages, req.StageBinding.Options, req.ResponseFormat, target)
	})
	result, err := s.service.SubmitProviderPool(stageCtx, poolReq)
	if err != nil || result == nil {
		return nil, quality.providerFailure(stageCtx, err, errProviderStageGeneric)
	}
	if result.Tunnel == nil {
		return nil, quality.providerFailure(stageCtx, errProviderStageGeneric, errProviderStageGeneric)
	}
	defer result.Tunnel.Close()
	if !providerStageDispatchMatches(result, dispatch) {
		return nil, quality.validation(errProviderStageGeneric)
	}

	body, err := collectProviderStageFrames(stageCtx, result.Tunnel.Stream().Frames, req.Limits.MaxOutputBytes)
	if err != nil {
		return nil, quality.providerFailure(stageCtx, err, errProviderStageGeneric)
	}
	body, err = normalizeSingleRequestProviderResponse(body, result.DispatchInfo)
	if err != nil {
		return nil, quality.providerFailure(stageCtx, err, errProviderStageGeneric)
	}
	response, err := decodeSingleRequestChatResponse(body, result.DispatchInfo)
	if err != nil {
		return nil, quality.providerFailure(stageCtx, err, errProviderStageGeneric)
	}
	return response, nil
}

func providerStageContext(parent context.Context, timeoutMS int) (context.Context, context.CancelFunc) {
	deadline := time.Now().Add(time.Duration(timeoutMS) * time.Millisecond)
	if callerDeadline, ok := parent.Deadline(); ok && callerDeadline.Before(deadline) {
		deadline = callerDeadline
	}
	return context.WithDeadline(parent, deadline)
}

func providerStageDispatchMatches(result *edgeservice.ProviderPoolDispatchResult, frozen *edgeservice.SingleRequestStageDispatchBinding) bool {
	d := result.DispatchInfo
	operation := config.ProtocolOperation(d.ProfileOperation)
	driver := config.ProtocolDriver(d.ProfileDriver)
	// A default managed resource selector intentionally freezes no provider ID:
	// the provider pool owns the concrete provider choice while profile, target,
	// credential, and execution-path fences remain fixed. Explicit selectors
	// continue to require the exact admitted provider.
	providerMatches := frozen.ProviderID == "" || d.ProviderID == frozen.ProviderID
	return result.Path == edgeservice.ProviderPoolPathTunnel &&
		(operation == config.OperationChatCompletions || operation == config.OperationResponses || operation == "") &&
		(driver == config.ProtocolDriverOpenAIChat || driver == config.ProtocolDriverOpenAIResponses) &&
		d.ModelGroupKey == frozen.ModelGroupKey && providerMatches &&
		d.Target == frozen.UpstreamModel && d.ProfileID == frozen.ProfileID &&
		d.CredentialSlotRef == frozen.CredentialSlotRef && d.CredentialRevision == frozen.CredentialRevision &&
		d.ExecutionPath == string(edgeservice.ProviderPoolPathTunnel)
}

func collectProviderStageFrames(ctx context.Context, frames <-chan *iop.ProviderTunnelFrame, maximum int) ([]byte, error) {
	if frames == nil {
		return nil, errProviderStageGeneric
	}
	var body bytes.Buffer
	state := 0 // 0=start, 1=body, 2=terminal
	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case frame, ok := <-frames:
			if !ok {
				if state != 2 {
					return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
				}
				return body.Bytes(), nil
			}
			if frame == nil {
				return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
			}
			switch frame.GetKind() {
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
				if state != 0 {
					return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
				}
				if frame.GetStatusCode() == http.StatusRequestEntityTooLarge {
					return nil, errors.Join(errProviderStageGeneric, errProviderStageContextLimit)
				}
				if frame.GetStatusCode() < 200 || frame.GetStatusCode() >= 300 {
					return nil, errProviderStageGeneric
				}
				state = 1
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
				if state != 1 {
					return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
				}
				if body.Len()+len(frame.GetBody()) > maximum {
					return nil, errors.Join(errProviderStageGeneric, errProviderStageOutputLimit)
				}
				body.Write(frame.GetBody())
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
				if state != 1 {
					return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
				}
				state = 2
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
				return nil, errProviderStageGeneric
			default:
				return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
			}
		}
	}
}

type singleRequestProviderResponseFormat struct {
	Type       string                                  `json:"type"`
	JSONSchema singleRequestProviderResponseJSONSchema `json:"json_schema"`
}

type singleRequestProviderResponseJSONSchema struct {
	Name   string                            `json:"name"`
	Strict bool                              `json:"strict"`
	Schema singleRequestProviderOutputSchema `json:"schema"`
}

type singleRequestProviderOutputSchema struct {
	Type                 string                                         `json:"type"`
	Properties           map[string]singleRequestProviderOutputProperty `json:"properties"`
	Required             []string                                       `json:"required"`
	AdditionalProperties bool                                           `json:"additionalProperties"`
}

type singleRequestProviderOutputProperty struct {
	Type        string                               `json:"type"`
	Description string                               `json:"description,omitempty"`
	Enum        []string                             `json:"enum,omitempty"`
	MinLength   int                                  `json:"minLength,omitempty"`
	Items       *singleRequestProviderOutputProperty `json:"items,omitempty"`
	MinItems    int                                  `json:"minItems,omitempty"`
	MaxItems    int                                  `json:"maxItems,omitempty"`
}

// buildSingleRequestChatBody owns all request authority. Stage options are
// frozen admission facts; callers cannot select messages, tools, credentials,
// or the stage-owned structured-output contract.
func buildSingleRequestChatBody(messages []chatMessage, options map[string]any, responseFormat *singleRequestProviderResponseFormat, targetModel string) ([]byte, error) {
	body := map[string]any{"model": targetModel, "messages": messages, "stream": false}
	for key, value := range options {
		switch key {
		case "model", "messages", "tools", "tool_choice", "stream", "credential", "credential_binding", "response_format":
			continue
		}
		body[key] = value
	}
	if responseFormat != nil {
		body["response_format"] = responseFormat
	}
	return json.Marshal(body)
}

type singleRequestChatResponse struct {
	ID      string                    `json:"id"`
	Object  string                    `json:"object"`
	Created int64                     `json:"created"`
	Model   string                    `json:"model"`
	Choices []singleRequestChatChoice `json:"choices"`
	Usage   *singleRequestChatUsage   `json:"usage,omitempty"`
}

// singleRequestChatUsage admits only standard bounded Chat Completions
// bookkeeping. The private stage does not retain or authorize on these values;
// provider usage observation remains owned by the existing tunnel path.
type singleRequestChatUsage struct {
	PromptTokens            int             `json:"prompt_tokens"`
	CompletionTokens        int             `json:"completion_tokens"`
	TotalTokens             int             `json:"total_tokens"`
	PromptTokensDetails     json.RawMessage `json:"prompt_tokens_details,omitempty"`
	CompletionTokensDetails json.RawMessage `json:"completion_tokens_details,omitempty"`
}

type singleRequestChatUsageAlias singleRequestChatUsage

func (u *singleRequestChatUsage) UnmarshalJSON(data []byte) error {
	if err := validateSingleRequestObjectFields(data, "prompt_tokens", "completion_tokens", "total_tokens", "prompt_tokens_details", "completion_tokens_details"); err != nil {
		return err
	}
	var a singleRequestChatUsageAlias
	if err := json.Unmarshal(data, &a); err != nil {
		return err
	}
	if a.PromptTokens < 0 || a.CompletionTokens < 0 || a.TotalTokens < 0 {
		return errors.New("provider stage: invalid usage")
	}
	*u = singleRequestChatUsage(a)
	return nil
}

type singleRequestChatChoice struct {
	Index        int                      `json:"index"`
	FinishReason string                   `json:"finish_reason"`
	Message      singleRequestChatMessage `json:"message"`
}

type singleRequestChatMessage struct {
	Role             string                          `json:"role"`
	Content          *string                         `json:"content"`
	ToolCalls        []any                           `json:"tool_calls,omitempty"`
	ReasoningContent *string                         `json:"reasoning_content,omitempty"`
	ExtraContent     singleRequestGeminiExtraContent `json:"extra_content"`
}

// singleRequestGeminiExtraContent is the only provider extension admitted by
// the private Gemini stage codecs. A text response drops it; a Review tool
// continuation may replay it through the request-local tool call only.
type singleRequestGeminiExtraContent struct {
	Google  singleRequestGeminiExtraContentGoogle `json:"google"`
	present bool
}

type singleRequestGeminiExtraContentGoogle struct {
	ThoughtSignature string `json:"thought_signature"`
}

func (v *singleRequestGeminiExtraContent) UnmarshalJSON(data []byte) error {
	if bytes.Equal(bytes.TrimSpace(data), []byte("null")) {
		return errors.New("provider stage: invalid Gemini extra content")
	}
	if err := validateSingleRequestObjectFields(data, "google"); err != nil {
		return err
	}
	type alias singleRequestGeminiExtraContent
	var decoded alias
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	if strings.TrimSpace(decoded.Google.ThoughtSignature) == "" {
		return errors.New("provider stage: invalid Gemini thought signature")
	}
	decoded.present = true
	*v = singleRequestGeminiExtraContent(decoded)
	return nil
}

func (v *singleRequestGeminiExtraContentGoogle) UnmarshalJSON(data []byte) error {
	if bytes.Equal(bytes.TrimSpace(data), []byte("null")) {
		return errors.New("provider stage: invalid Gemini extra content")
	}
	if err := validateSingleRequestObjectFields(data, "thought_signature"); err != nil {
		return err
	}
	type alias singleRequestGeminiExtraContentGoogle
	var decoded alias
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	*v = singleRequestGeminiExtraContentGoogle(decoded)
	return nil
}

func validateSingleRequestObjectFields(data []byte, allowed ...string) error {
	decoder := json.NewDecoder(bytes.NewReader(data))
	tok, err := decoder.Token()
	if err != nil {
		return err
	}
	delim, ok := tok.(json.Delim)
	if !ok || delim != '{' {
		return errors.New("json: expected object")
	}
	allowedMap := make(map[string]bool, len(allowed))
	for _, a := range allowed {
		allowedMap[a] = true
	}
	seen := make(map[string]bool, len(allowed))
	for decoder.More() {
		keyTok, err := decoder.Token()
		if err != nil {
			return err
		}
		key, ok := keyTok.(string)
		if !ok {
			return errors.New("json: object key must be string")
		}
		if !allowedMap[key] || seen[key] {
			return errors.New("json: unknown or non-canonical field: " + key)
		}
		seen[key] = true
		var val json.RawMessage
		if err := decoder.Decode(&val); err != nil {
			return err
		}
	}
	return nil
}

type singleRequestChatResponseAlias singleRequestChatResponse

func (r *singleRequestChatResponse) UnmarshalJSON(data []byte) error {
	if err := validateSingleRequestObjectFields(data, "id", "object", "created", "model", "choices", "usage"); err != nil {
		return err
	}
	var a singleRequestChatResponseAlias
	if err := json.Unmarshal(data, &a); err != nil {
		return err
	}
	*r = singleRequestChatResponse(a)
	return nil
}

type singleRequestChatChoiceAlias singleRequestChatChoice

func (c *singleRequestChatChoice) UnmarshalJSON(data []byte) error {
	if err := validateSingleRequestObjectFields(data, "index", "finish_reason", "message"); err != nil {
		return err
	}
	var a singleRequestChatChoiceAlias
	if err := json.Unmarshal(data, &a); err != nil {
		return err
	}
	*c = singleRequestChatChoice(a)
	return nil
}

type singleRequestChatMessageAlias singleRequestChatMessage

func (m *singleRequestChatMessage) UnmarshalJSON(data []byte) error {
	if err := validateSingleRequestObjectFields(data, "role", "content", "tool_calls", "reasoning_content", "extra_content"); err != nil {
		return err
	}
	var a singleRequestChatMessageAlias
	if err := json.Unmarshal(data, &a); err != nil {
		return err
	}
	*m = singleRequestChatMessage(a)
	return nil
}

func validateSingleRequestJSON(data []byte) error {
	decoder := json.NewDecoder(bytes.NewReader(data))
	tok, err := decoder.Token()
	if err != nil {
		return err
	}
	if err := validateJSONValue(decoder, tok); err != nil {
		return err
	}
	var extra json.RawMessage
	if err := decoder.Decode(&extra); err != io.EOF {
		return errors.New("json: trailing content")
	}
	return nil
}

func validateJSONValue(decoder *json.Decoder, tok json.Token) error {
	delim, ok := tok.(json.Delim)
	if !ok {
		return nil
	}
	switch delim {
	case '{':
		seen := make(map[string]bool)
		for decoder.More() {
			keyTok, err := decoder.Token()
			if err != nil {
				return err
			}
			key, ok := keyTok.(string)
			if !ok {
				return errors.New("json: object key must be string")
			}
			if seen[key] {
				return errors.New("json: duplicate key")
			}
			seen[key] = true

			valTok, err := decoder.Token()
			if err != nil {
				return err
			}
			if err := validateJSONValue(decoder, valTok); err != nil {
				return err
			}
		}
		closingTok, err := decoder.Token()
		if err != nil {
			return err
		}
		if closingDelim, ok := closingTok.(json.Delim); !ok || closingDelim != '}' {
			return errors.New("json: expected closing brace")
		}
		return nil
	case '[':
		for decoder.More() {
			elemTok, err := decoder.Token()
			if err != nil {
				return err
			}
			if err := validateJSONValue(decoder, elemTok); err != nil {
				return err
			}
		}
		closingTok, err := decoder.Token()
		if err != nil {
			return err
		}
		if closingDelim, ok := closingTok.(json.Delim); !ok || closingDelim != ']' {
			return errors.New("json: expected closing bracket")
		}
		return nil
	default:
		return errors.New("json: unexpected delimiter")
	}
}

func decodeSingleRequestChatResponse(body []byte, dispatch edgeservice.RunDispatch) (*singleRequestProviderStageResponse, error) {
	if err := validateSingleRequestJSON(body); err != nil {
		return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
	}
	var decoded singleRequestChatResponse
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&decoded); err != nil || decoder.More() || len(decoded.Choices) != 1 {
		return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
	}
	choice := decoded.Choices[0]
	if choice.FinishReason == "length" {
		return nil, errors.Join(errProviderStageGeneric, errProviderStageOutputLimit)
	}
	if choice.FinishReason == "context_length" || choice.FinishReason == "context_length_exceeded" {
		return nil, errors.Join(errProviderStageGeneric, errProviderStageContextLimit)
	}
	if choice.Index != 0 || choice.Message.Role != "assistant" || choice.FinishReason != "stop" || len(choice.Message.ToolCalls) != 0 || choice.Message.Content == nil {
		return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
	}
	var extra any
	if decoder.Decode(&extra) != io.EOF {
		return nil, errors.Join(errProviderStageGeneric, errProviderStageMalformed)
	}
	return &singleRequestProviderStageResponse{Output: *choice.Message.Content, Dispatch: dispatch}, nil
}
