// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"fmt"
	"strings"
)

type requestFieldError struct {
	Field   string
	Message string
	Code    string
}

func (e *requestFieldError) Error() string {
	return e.Message
}

type chatFieldPolicy int

const (
	forwardChatField chatFieldPolicy = iota
	ignoreChatField
	rejectChatField
)

var chatRequestFieldPolicies = map[string]chatFieldPolicy{
	// Parsed and forwarded to the SDK, or used by response shaping.
	"model":                 forwardChatField,
	"messages":              forwardChatField,
	"stream":                forwardChatField,
	"max_tokens":            forwardChatField,
	"max_completion_tokens": forwardChatField,
	"max_length":            forwardChatField,
	"temperature":           forwardChatField,
	"top_p":                 forwardChatField,
	"frequency_penalty":     forwardChatField,
	"presence_penalty":      forwardChatField,
	"repetition_penalty":    forwardChatField,
	"top_k":                 forwardChatField,
	"stop":                  forwardChatField,
	"seed":                  forwardChatField,
	"n":                     forwardChatField,
	"stream_options":        forwardChatField,
	"response_format":       forwardChatField,
	"logprobs":              forwardChatField,
	"top_logprobs":          forwardChatField,
	"tools":                 forwardChatField,
	"tool_choice":           forwardChatField,
	"parallel_tool_calls":   forwardChatField,
	"functions":             forwardChatField,
	"function_call":         forwardChatField,
	"enable_thinking":       forwardChatField,
	"thinking_budget":       forwardChatField,
	"chat_template_kwargs":  forwardChatField,

	// Accepted for compatibility but not forwarded.
	"dashscope_custom":           ignoreChatField,
	"debug":                      ignoreChatField,
	"head_response_per_tokens":   ignoreChatField,
	"include_chunk_usage":        ignoreChatField,
	"incremental_output":         ignoreChatField,
	"max_input_tokens":           ignoreChatField,
	"output_in_full":             ignoreChatField,
	"output_tokens_per_response": ignoreChatField,
	"tool_stream":                ignoreChatField,
	"prompt":                     ignoreChatField,
	"result_format":              ignoreChatField,
	"task_group":                 ignoreChatField,
	"task":                       ignoreChatField,
	"use_raw_prompt":             ignoreChatField,
	"function":                   ignoreChatField,
	"search_options":             ignoreChatField,
	"enable_search":              ignoreChatField,
	"stop_words":                 ignoreChatField,
	"enableEncrypt":              ignoreChatField,
	"user":                       ignoreChatField,
	"request_id":                 ignoreChatField,
	"cache_control":              ignoreChatField,
	"store":                      ignoreChatField,

	// Recognized but not implemented by this proxy.
	"audio":                   rejectChatField,
	"modalities":              rejectChatField,
	"preserve_thinking":       rejectChatField,
	"reasoning_effort":        rejectChatField,
	"enable_code_interpreter": rejectChatField,
	"metadata":                rejectChatField,
	"web_search_options":      rejectChatField,
	"service_tier":            rejectChatField,
	"prediction":              rejectChatField,
	"logit_bias":              rejectChatField,
}

func validateChatRequestFields(body []byte) ([]byte, error) {
	normalized, err := normalizeChatRequestFields(body)
	if err != nil {
		return nil, err
	}
	for field, raw := range normalized {
		if err := validateChatRequestField(field, raw); err != nil {
			return nil, err
		}
	}
	return json.Marshal(normalized)
}

func normalizeChatRequestFields(body []byte) (map[string]json.RawMessage, error) {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(body, &fields); err != nil {
		return nil, err
	}
	normalized := make(map[string]json.RawMessage, len(fields))
	for field, raw := range fields {
		trimmed := strings.TrimSpace(field)
		if trimmed == "" {
			return nil, fmt.Errorf("request field name must not be empty")
		}
		if _, exists := normalized[trimmed]; exists {
			return nil, fmt.Errorf("duplicate request field after trimming whitespace: %q", trimmed)
		}
		normalized[trimmed] = raw
	}
	return normalized, nil
}

func validateChatRequestField(field string, raw json.RawMessage) error {
	policy, ok := chatRequestFieldPolicies[field]
	if !ok {
		return &requestFieldError{Field: field, Message: fmt.Sprintf("unknown chat completion field %q", field)}
	}
	switch policy {
	case forwardChatField:
		return nil
	case ignoreChatField:
		return validateIgnoredChatField(field, raw)
	case rejectChatField:
		return &requestFieldError{Field: field, Message: fmt.Sprintf("unsupported chat completion field %q", field)}
	default:
		return &requestFieldError{Field: field, Message: fmt.Sprintf("unknown chat completion field %q", field)}
	}
}

func validateIgnoredChatField(field string, raw json.RawMessage) error {
	if len(raw) == 0 || string(raw) == "null" || field != "cache_control" {
		return nil
	}
	return validateCacheControlType(raw, "cache_control")
}

func validateInboundToolCalls(body []byte) error {
	var raw struct {
		Messages []map[string]json.RawMessage `json:"messages"`
	}
	if err := json.Unmarshal(body, &raw); err != nil {
		return nil // let the main parser surface the error instead
	}
	for i, msg := range raw.Messages {
		tcRaw, ok := msg["tool_calls"]
		if !ok || len(tcRaw) == 0 || string(tcRaw) == "null" {
			continue
		}
		var tcs []map[string]json.RawMessage
		if err := json.Unmarshal(tcRaw, &tcs); err != nil {
			return fmt.Errorf("messages[%d].tool_calls must be an array", i)
		}
		for j, tc := range tcs {
			fnRaw, ok := tc["function"]
			if !ok || len(fnRaw) == 0 || string(fnRaw) == "null" {
				return fmt.Errorf("messages[%d].tool_calls[%d].function is required", i, j)
			}
			var fn map[string]json.RawMessage
			if err := json.Unmarshal(fnRaw, &fn); err != nil {
				return fmt.Errorf("messages[%d].tool_calls[%d].function must be an object", i, j)
			}
			if _, has := fn["name"]; !has {
				return fmt.Errorf("messages[%d].tool_calls[%d].function.name is required", i, j)
			}
			if _, has := fn["arguments"]; !has {
				return fmt.Errorf("<400> InternalError.Algo.InvalidParameter: If tool_calls are present in the message, function.arguments must be defined. (messages[%d].tool_calls[%d])", i, j)
			}
		}
	}
	return nil
}

func validateInboundMessageContent(body []byte) error {
	var raw struct {
		Messages []map[string]json.RawMessage `json:"messages"`
	}
	if err := json.Unmarshal(body, &raw); err != nil {
		return nil // let the main parser surface the error instead
	}
	for i, msg := range raw.Messages {
		// Validate cache_control at message level
		if ccRaw, ok := msg["cache_control"]; ok {
			if err := validateCacheControlType(ccRaw, fmt.Sprintf("messages[%d].cache_control", i)); err != nil {
				return err
			}
		}

		contentRaw, ok := msg["content"]
		if !ok || len(contentRaw) == 0 || string(contentRaw) == "null" {
			continue
		}
		var contentStr string
		if err := json.Unmarshal(contentRaw, &contentStr); err == nil {
			continue
		}
		var contentParts []openAIContentPart
		if err := json.Unmarshal(contentRaw, &contentParts); err == nil {
			// Validate cache_control inside each content part
			if err := validateContentPartsCacheControl(contentRaw, i); err != nil {
				return err
			}
			continue
		}
		if len(contentRaw) > 0 && contentRaw[0] == '[' {
			return fmt.Errorf("if content is list. item must be dict (messages.%d.content)", i)
		}
		return fmt.Errorf("messages.%d.content input should be a valid string or list", i)
	}
	return nil
}

// validateCacheControlType checks that a cache_control JSON object has
// type == "ephemeral" (the only value allowed by the ModelStudio contract).
func validateCacheControlType(raw json.RawMessage, path string) error {
	if len(raw) == 0 || string(raw) == "null" {
		return nil
	}
	var cc struct {
		Type string `json:"type"`
	}
	if err := json.Unmarshal(raw, &cc); err != nil || cc.Type != "ephemeral" {
		return &requestFieldError{
			Field:   "cache_control.type",
			Message: fmt.Sprintf("<400> InternalError.Algo.InvalidParameter: %s.type must be 'ephemeral'", path),
		}
	}
	return nil
}

// validateContentPartsCacheControl checks that any cache_control.type
// inside content array elements is exactly "ephemeral" (the only value
// allowed by the ModelStudio contract).
func validateContentPartsCacheControl(contentRaw json.RawMessage, msgIdx int) error {
	var parts []struct {
		CacheControl *struct {
			Type string `json:"type"`
		} `json:"cache_control,omitempty"`
	}
	if err := json.Unmarshal(contentRaw, &parts); err != nil {
		return nil
	}
	for j, part := range parts {
		if part.CacheControl == nil {
			continue
		}
		if part.CacheControl.Type != "ephemeral" {
			return &requestFieldError{
				Field:   "cache_control.type",
				Message: fmt.Sprintf("<400> InternalError.Algo.InvalidParameter: messages[%d].content[%d].cache_control.type must be 'ephemeral'", msgIdx, j),
			}
		}
	}
	return nil
}

func validateChatRequestSemantics(req *openAIChatRequest) error {
	if strings.TrimSpace(req.Model) == "" {
		return &requestFieldError{Field: "model", Message: "required model is missing"}
	}
	for i, msg := range req.Messages {
		if !isValidChatRole(msg.Role) {
			return &requestFieldError{Field: "messages.role", Message: fmt.Sprintf("messages[%d].role must be one of ['system', 'assistant', 'user', 'tool', 'function']", i)}
		}
	}
	if req.N != nil && *req.N < 0 {
		return &requestFieldError{Field: "n", Message: "<400> InternalError.Algo.InvalidParameter: Range of n should be [1, 4]", Code: "integer_below_min_value"}
	}
	if req.N != nil && *req.N > 4 {
		return &requestFieldError{Field: "n", Message: "<400> InternalError.Algo.InvalidParameter: Range of n should be [1, 4]", Code: "integer_below_min_value"}
	}
	if req.MaxLength != nil && *req.MaxLength <= 0 {
		return &requestFieldError{Field: "max_length", Message: "max_length must be a positive integer"}
	}
	if req.ThinkingBudget != nil && *req.ThinkingBudget <= 0 {
		return &requestFieldError{Field: "thinking_budget", Message: "<400> InternalError.Algo.InvalidParameter: thinking_budget must be a positive integer"}
	}
	if req.Temperature != nil && (*req.Temperature < 0 || *req.Temperature >= 2) {
		return &requestFieldError{Field: "temperature", Message: "<400> InternalError.Algo.InvalidParameter: Temperature should be in [0.0, 2.0)"}
	}
	if req.TopP != nil && (*req.TopP < 0 || *req.TopP > 1) {
		return &requestFieldError{Field: "top_p", Message: "'top_p' must be in (0.0, 1.0]"}
	}
	if req.StreamOptions != nil && !req.Stream {
		return &requestFieldError{Field: "stream_options", Message: "stream_options can only be set when stream is true"}
	}
	if enableThinking, ok := effectiveEnableThinking(req); ok && enableThinking && req.N != nil && *req.N > 1 {
		return &requestFieldError{Field: "n", Message: "<400> InternalError.Algo.InvalidParameter: The n parameter must be 1 when enable_thinking is true"}
	}
	if enableThinking, ok := effectiveEnableThinking(req); ok && enableThinking && thinkingModeRejectsToolChoice(req.ToolChoice) {
		return &requestFieldError{Field: "tool_choice", Message: "<400> InternalError.Algo.InvalidParameter: The tool_choice parameter does not support being set to required or object in thinking mode"}
	}
	if len(req.ResponseFormat) > 0 && strings.TrimSpace(string(req.ResponseFormat)) != "null" {
		if err := validateResponseFormat(req.ResponseFormat, req.Messages); err != nil {
			return err
		}
	}
	return nil
}

func thinkingModeRejectsToolChoice(raw json.RawMessage) bool {
	if len(raw) == 0 || strings.TrimSpace(string(raw)) == "null" {
		return false
	}
	var value string
	if err := json.Unmarshal(raw, &value); err == nil {
		return value == "required"
	}
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return false
	}
	_, hasType := obj["type"]
	_, hasFunction := obj["function"]
	return hasType || hasFunction
}

func isValidChatRole(role string) bool {
	switch role {
	case "system", "assistant", "user", "tool", "function":
		return true
	default:
		return false
	}
}

func effectiveEnableThinking(req *openAIChatRequest) (bool, bool) {
	if req.EnableThinking != nil {
		return *req.EnableThinking, true
	}
	if len(req.ChatTemplateKwargs) == 0 {
		return false, false
	}
	value, ok := req.ChatTemplateKwargs["enable_thinking"]
	if !ok {
		return false, false
	}
	enableThinking, ok := value.(bool)
	return enableThinking, ok
}

func validateResponseFormat(raw json.RawMessage, messages []openAIMessage) error {
	var value map[string]json.RawMessage
	if err := json.Unmarshal(raw, &value); err != nil {
		return &requestFieldError{Field: "response_format", Message: "response_format must be a valid dictionary"}
	}
	typeRaw, ok := value["type"]
	if !ok {
		return &requestFieldError{Field: "response_format", Message: "response_format.type is required"}
	}
	var typ string
	if err := json.Unmarshal(typeRaw, &typ); err != nil {
		return &requestFieldError{Field: "response_format", Message: "response_format.type must be a string"}
	}
	if typ != "json_schema" {
		if _, exists := value["json_schema"]; exists {
			return &requestFieldError{Field: "response_format", Message: fmt.Sprintf("<400> InternalError.Algo.InvalidParameter: Unknown parameter: 'response_format.json_schema'. 'response_format.json_schema' cannot be provided when 'response_format.type' is '%s'.", typ)}
		}
	}
	switch typ {
	case "text":
		return nil
	case "json_object":
		if !messagesContainJSON(messages) {
			return &requestFieldError{Field: "messages", Message: "'messages' must contain the word 'json' in some form, to use 'response_format' of type 'json_object'"}
		}
		return nil
	case "json_schema":
		schema, exists := value["json_schema"]
		if !exists {
			return &requestFieldError{Field: "response_format", Message: "response_format.json_schema is required when response_format.type is json_schema"}
		}
		var schemaObj map[string]interface{}
		if err := json.Unmarshal(schema, &schemaObj); err != nil || schemaObj == nil {
			return &requestFieldError{Field: "response_format", Message: "response_format.json_schema must be a valid dictionary"}
		}
		return nil
	default:
		return &requestFieldError{Field: "response_format", Message: fmt.Sprintf("<400> InternalError.Algo.InvalidParameter: 'response_format.type' Invalid value: '%s'. Supported values are: 'json_object' and 'text'.", typ)}
	}
}

func messagesContainJSON(messages []openAIMessage) bool {
	for _, msg := range messages {
		if strings.Contains(strings.ToLower(msg.Content), "json") {
			return true
		}
		for _, part := range msg.ContentParts {
			if strings.Contains(strings.ToLower(part.Text), "json") {
				return true
			}
		}
	}
	return false
}
