// 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"
	"strconv"

	"github.com/dashscope/dashscope-confidential-maas/sdk"
)

// OpenAI Chat Completions API request
type openAIChatRequest struct {
	Model               string                 `json:"model"`
	Messages            []openAIMessage        `json:"messages"`
	Stream              bool                   `json:"stream"`
	MaxTokens           *int                   `json:"max_tokens,omitempty"`
	MaxCompletionTokens *int                   `json:"max_completion_tokens,omitempty"`
	MaxLength           *int                   `json:"max_length,omitempty"`
	Temperature         *float64               `json:"temperature,omitempty"`
	TopP                *float64               `json:"top_p,omitempty"`
	FrequencyPenalty    *float64               `json:"frequency_penalty,omitempty"`
	PresencePenalty     *float64               `json:"presence_penalty,omitempty"`
	RepetitionPenalty   *float64               `json:"repetition_penalty,omitempty"`
	TopK                *int                   `json:"top_k,omitempty"`
	Logprobs            *bool                  `json:"logprobs,omitempty"`
	TopLogprobs         *int                   `json:"top_logprobs,omitempty"`
	Stop                stopValue              `json:"stop,omitempty"`
	Seed                *int                   `json:"seed,omitempty"`
	N                   *int                   `json:"n,omitempty"`
	StreamOptions       *sdk.StreamOptions     `json:"stream_options,omitempty"`
	ResponseFormat      json.RawMessage        `json:"response_format,omitempty"`
	ChatTemplateKwargs  map[string]interface{} `json:"chat_template_kwargs,omitempty"`
	EnableThinking      *bool                  `json:"enable_thinking,omitempty"`
	ThinkingBudget      *int                   `json:"thinking_budget,omitempty"`
	Tools               json.RawMessage        `json:"tools,omitempty"`
	ToolChoice          json.RawMessage        `json:"tool_choice,omitempty"`
	// ParallelToolCalls=false asks for at most one tool call; the proxy
	// enforces that response shape for compatibility.
	ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
	// Legacy OpenAI function calling fields (deprecated but still sent by
	// older clients). Converted to tools/tool_choice in applyChatParameters.
	Functions    json.RawMessage `json:"functions,omitempty"`
	FunctionCall json.RawMessage `json:"function_call,omitempty"`
}

// OpenAI Completions API request (--backend vllm / openai)
type openAICompletionsRequest struct {
	Model               string      `json:"model"`
	Prompt              promptValue `json:"prompt"` // can be string or array
	Stream              bool        `json:"stream"`
	MaxTokens           int         `json:"max_tokens,omitempty"`
	MaxCompletionTokens int         `json:"max_completion_tokens,omitempty"`
	Temperature         float64     `json:"temperature,omitempty"`
	TopP                float64     `json:"top_p,omitempty"`
}

// promptValue handles both string and array prompt formats
type promptValue struct {
	String  string
	Array   []string
	IsArray bool
}

type stopValue struct {
	Values   []string
	TokenIDs []int
	Set      bool
}

func (s *stopValue) UnmarshalJSON(data []byte) error {
	if string(data) == "null" {
		return nil
	}
	var one string
	if err := json.Unmarshal(data, &one); err == nil {
		s.Values = []string{one}
		s.Set = true
		return nil
	}
	var oneToken int
	if err := json.Unmarshal(data, &oneToken); err == nil {
		s.TokenIDs = []int{oneToken}
		s.Set = true
		return nil
	}
	var many []json.RawMessage
	if err := json.Unmarshal(data, &many); err == nil {
		for _, raw := range many {
			var text string
			if err := json.Unmarshal(raw, &text); err == nil {
				s.Values = append(s.Values, text)
				continue
			}
			var tokenID int
			if err := json.Unmarshal(raw, &tokenID); err == nil {
				s.TokenIDs = append(s.TokenIDs, tokenID)
				continue
			}
			return fmt.Errorf("stop array elements must be string or integer token id")
		}
		s.Set = true
		return nil
	}
	return fmt.Errorf("stop must be string, integer token id, or array of strings/token ids")
}

func (p *promptValue) UnmarshalJSON(data []byte) error {
	// Try string first
	if err := json.Unmarshal(data, &p.String); err == nil {
		p.IsArray = false
		return nil
	}
	// Try array
	if err := json.Unmarshal(data, &p.Array); err == nil {
		p.IsArray = true
		return nil
	}
	return fmt.Errorf("prompt must be string or array of strings")
}

func (p *promptValue) First() string {
	if p.IsArray && len(p.Array) > 0 {
		return p.Array[0]
	}
	return p.String
}

type openAIMessage struct {
	Role             string           `json:"role"`
	Content          string           `json:"content"` // simple string format
	ReasoningContent string           `json:"reasoning_content,omitempty"`
	ToolCalls        []openAIToolCall `json:"tool_calls,omitempty"`
	ToolCallID       string           `json:"tool_call_id,omitempty"`
	// FunctionCall is the legacy OpenAI v1 response shape, populated when
	// the client sent the request using `functions` instead of `tools`.
	FunctionCall *openAIFunctionCall `json:"function_call,omitempty"`
	// For multimodal format, content is an array of content parts
	ContentParts []openAIContentPart `json:"-"`
}

// openAIToolCall mirrors OpenAI's tool_call shape. Index is set on the
// non-stream response to keep response shape stable (OpenAI spec
// itself does not require it on non-stream tool_calls); on inbound parsing it stays nil and serializes
// out via omitempty.
type openAIToolCall struct {
	Index    *int             `json:"index,omitempty"`
	ID       string           `json:"id"`
	Type     string           `json:"type"`
	Function sdk.FunctionCall `json:"function"`
}

func (tc *openAIToolCall) UnmarshalJSON(data []byte) error {
	type Alias struct {
		Index    json.RawMessage  `json:"index,omitempty"`
		ID       string           `json:"id"`
		Type     string           `json:"type"`
		Function sdk.FunctionCall `json:"function"`
	}
	var aux Alias
	if err := json.Unmarshal(data, &aux); err != nil {
		return err
	}
	tc.ID = aux.ID
	tc.Type = aux.Type
	tc.Function = aux.Function
	if len(aux.Index) == 0 || string(aux.Index) == "null" {
		return nil
	}
	var idx int
	if err := json.Unmarshal(aux.Index, &idx); err == nil {
		tc.Index = &idx
		return nil
	}
	var idxStr string
	if err := json.Unmarshal(aux.Index, &idxStr); err != nil {
		return fmt.Errorf("index must be integer")
	}
	parsed, err := strconv.Atoi(idxStr)
	if err != nil {
		return fmt.Errorf("index must be integer")
	}
	tc.Index = &parsed
	return nil
}

// openAIFunctionCall is the legacy v1 response shape, equivalent to a
// single tool_call but flattened.
type openAIFunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

type openAIContentPart struct {
	Type     string        `json:"type"`
	Text     string        `json:"text,omitempty"`
	ImageURL *sdk.ImageURL `json:"image_url,omitempty"`
}

// Custom UnmarshalJSON to support both string and array content formats
// while preserving tool_calls / tool_call_id pass-through.
func (m *openAIMessage) UnmarshalJSON(data []byte) error {
	// Use an alias to avoid recursive UnmarshalJSON
	type Alias struct {
		Role             string              `json:"role"`
		Content          json.RawMessage     `json:"content"`
		ReasoningContent string              `json:"reasoning_content"`
		ToolCalls        []openAIToolCall    `json:"tool_calls,omitempty"`
		ToolCallID       string              `json:"tool_call_id,omitempty"`
		FunctionCall     *openAIFunctionCall `json:"function_call,omitempty"`
	}
	var aux Alias
	if err := json.Unmarshal(data, &aux); err != nil {
		return err
	}
	m.Role = aux.Role
	m.ReasoningContent = aux.ReasoningContent
	m.ToolCalls = aux.ToolCalls
	m.ToolCallID = aux.ToolCallID
	m.FunctionCall = aux.FunctionCall

	// content is optional when the message carries tool_calls only.
	if len(aux.Content) == 0 || string(aux.Content) == "null" {
		return nil
	}

	// Try to parse content as string first
	var contentStr string
	if err := json.Unmarshal(aux.Content, &contentStr); err == nil {
		m.Content = contentStr
		return nil
	}

	// Try to parse content as array of content parts
	var contentParts []openAIContentPart
	if err := json.Unmarshal(aux.Content, &contentParts); err == nil {
		m.ContentParts = contentParts
		// Extract text from first text part
		for _, part := range contentParts {
			if part.Type == "text" {
				m.Content = part.Text
				break
			}
		}
		return nil
	}

	return fmt.Errorf("content must be string or array of content parts")
}

// non-streaming response
type openAIResponse struct {
	ID      string         `json:"id"`
	Object  string         `json:"object"`
	Created int64          `json:"created"`
	Model   string         `json:"model"`
	Choices []openAIChoice `json:"choices"`
	Usage   openAIUsage    `json:"usage"`
}

type openAIChoice struct {
	Index        int            `json:"index"`
	Message      *openAIMessage `json:"message,omitempty"` // for chat completions
	Delta        interface{}    `json:"delta,omitempty"`   // for chat completions streaming
	Text         string         `json:"text,omitempty"`    // for completions API
	FinishReason *string        `json:"finish_reason"`
	Logprobs     interface{}    `json:"logprobs"`
}

type openAIUsage struct {
	PromptTokens            int                    `json:"prompt_tokens"`
	CompletionTokens        int                    `json:"completion_tokens"`
	TotalTokens             int                    `json:"total_tokens"`
	PromptTokensDetails     map[string]interface{} `json:"prompt_tokens_details,omitempty"`
	CompletionTokensDetails map[string]interface{} `json:"completion_tokens_details,omitempty"`
}

type openAIErrorResponse struct {
	Error     openAIError `json:"error"`
	ID        string      `json:"id"`
	RequestID string      `json:"request_id"`
}

type openAIError struct {
	Message string      `json:"message"`
	Type    string      `json:"type"`
	Param   interface{} `json:"param"`
	Code    string      `json:"code"`
}

// streaming chunk
type openAIChunk struct {
	ID      string         `json:"id"`
	Object  string         `json:"object"`
	Created int64          `json:"created"`
	Model   string         `json:"model"`
	Choices []openAIChoice `json:"choices"`
	Usage   *openAIUsage   `json:"usage,omitempty"`
}
