// 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 (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"

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

	"github.com/google/uuid"
)

const maxRequestBodyBytes int64 = 64 * 1024 * 1024

var errRequestBodyTooLarge = errors.New("request body is too large")

type handler struct {
	pool  *pool
	model string
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/v1/chat/completions", "/chat/completions":
		h.chatCompletions(w, r)
	case "/v1/completions", "/completions":
		h.completions(w, r)
	case "/v1/models", "/models":
		h.listModels(w, r)
	default:
		http.NotFound(w, r)
	}
}

// listModels returns a minimal OpenAI-compatible model list so benchmark
// tools that call GET /v1/models don't fail.
func (h *handler) listModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeOpenAIError(w, http.StatusMethodNotAllowed, "method not allowed", "", "")
		return
	}
	resp := map[string]interface{}{
		"object": "list",
		"data": []map[string]interface{}{
			{
				"id":       h.model,
				"object":   "model",
				"created":  time.Now().Unix(),
				"owned_by": "cmaas",
			},
		},
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

func readRequestBody(w http.ResponseWriter, r *http.Request) ([]byte, error) {
	return readRequestBodyWithLimit(w, r, maxRequestBodyBytes)
}

func readRequestBodyWithLimit(w http.ResponseWriter, r *http.Request, limit int64) ([]byte, error) {
	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, limit))
	if err != nil {
		var maxErr *http.MaxBytesError
		if errors.As(err, &maxErr) {
			return nil, errRequestBodyTooLarge
		}
		return nil, err
	}
	return body, nil
}

func writeRequestBodyError(w http.ResponseWriter, err error) {
	if errors.Is(err, errRequestBodyTooLarge) {
		writeOpenAIError(w, http.StatusRequestEntityTooLarge, "request body is too large", "request_body_too_large", "")
		return
	}
	writeOpenAIError(w, http.StatusBadRequest, "failed to read request body", "", "")
}

func writeAcquireError(w http.ResponseWriter, err error, requestID string) {
	message := "failed to acquire backend worker"
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		message = "request canceled before an available worker was assigned"
	}
	writeOpenAIError(w, http.StatusServiceUnavailable, message, "service_unavailable", requestID)
}

func (h *handler) chatCompletions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeOpenAIError(w, http.StatusMethodNotAllowed, "method not allowed", "", "")
		return
	}

	body, err := readRequestBody(w, r)
	if err != nil {
		writeRequestBodyError(w, err)
		return
	}
	defer r.Body.Close()

	var req openAIChatRequest
	body, err = validateChatRequestFields(body)
	if err != nil {
		writeOpenAIError(w, http.StatusBadRequest, "invalid request: "+err.Error(), "", "")
		return
	}
	if err := validateInboundToolCalls(body); err != nil {
		writeOpenAIError(w, http.StatusBadRequest, err.Error(), "", "")
		return
	}
	if err := validateInboundMessageContent(body); err != nil {
		writeOpenAIError(w, http.StatusBadRequest, err.Error(), "", "")
		return
	}
	if err := json.Unmarshal(body, &req); err != nil {
		writeOpenAIError(w, http.StatusBadRequest, jsonUnmarshalMessage(err), "", "")
		return
	}
	if err := validateChatRequestSemantics(&req); err != nil {
		var fieldErr *requestFieldError
		if errors.As(err, &fieldErr) {
			writeOpenAIError(w, http.StatusBadRequest, fieldErr.Message, fieldErr.Code, "")
			return
		}
		writeOpenAIError(w, http.StatusBadRequest, "invalid request: "+err.Error(), "", "")
		return
	}

	msgs := convertOpenAIMessagesToSDK(req.Messages)

	requestID := uuid.New().String()
	respID := "chatcmpl-" + requestID
	sdkReq := &sdk.ChatCompletionRequest{
		RequestID:        requestID,
		Messages:         msgs,
		Stream:           req.Stream,
		StructuredStream: req.Stream,
	}
	if err := applyChatParameters(&req, sdkReq); err != nil {
		writeOpenAIError(w, http.StatusBadRequest, "invalid request: "+err.Error(), "", "")
		return
	}

	log.Printf("[INFO] request started id=%s endpoint=/v1/chat/completions stream=%t backend_model=%q", respID, req.Stream, h.model)

	ctx := r.Context()
	client, err := h.pool.acquire(ctx)
	if err != nil {
		writeAcquireError(w, err, requestID)
		return
	}
	defer h.pool.release(client)
	completion, err := client.Chat(ctx, sdkReq)
	if err != nil {
		log.Printf("[WARN] request failed id=%s task_id=%s error=%v", respID, requestID, err)
		writeSDKError(w, err, requestID)
		return
	}
	log.Printf("[INFO] cmaas request mapped id=%s task_id=%s stream=%t", respID, completion.TaskID, req.Stream)

	created := time.Now().Unix()

	// Determine whether tool_choice was "none" so the response writer can
	// strip any tool_calls the backend may have returned anyway.
	toolChoiceNone := isToolChoiceNone(sdkReq.ToolChoice)
	// Track whether the client used the legacy `functions` field so we
	// can mirror the legacy `function_call` shape in the response.
	legacyMode := hasJSONValue(req.Functions)
	// When parallel_tool_calls=false, defensively cap the response to one
	// tool call (some backends do not enforce the hint reliably).
	singleToolOnly := req.ParallelToolCalls != nil && !*req.ParallelToolCalls

	if req.Stream {
		h.writeStream(w, completion, respID, created, h.model, includeStreamUsage(req.StreamOptions), toolChoiceNone, legacyMode, singleToolOnly)
	} else {
		h.writeResponse(w, completion, respID, created, h.model, toolChoiceNone, legacyMode, singleToolOnly)
	}
}

// convertOpenAIMessagesToSDK maps OpenAI chat messages to the SDK wire shape.
// OpenAI-compatible providers accept requests without a user message; the
// backend expects one, so append an empty user turn in that case.
func convertOpenAIMessagesToSDK(messages []openAIMessage) []sdk.Message {
	// Drop the optional index from inbound history tool_calls; the backend only
	// needs id/type/function when replaying completed assistant calls.
	msgs := make([]sdk.Message, 0, len(messages)+1)
	hasUserMessage := false
	for _, m := range messages {
		if m.Role == "user" {
			hasUserMessage = true
		}
		var sdkTCs []sdk.ToolCall
		if len(m.ToolCalls) > 0 {
			sdkTCs = make([]sdk.ToolCall, len(m.ToolCalls))
			for j, tc := range m.ToolCalls {
				sdkTCs[j] = sdk.ToolCall{ID: tc.ID, Type: tc.Type, Function: tc.Function}
			}
		}
		msg := sdk.Message{
			Role:       m.Role,
			Content:    m.Content,
			ToolCalls:  sdkTCs,
			ToolCallID: m.ToolCallID,
		}
		if len(m.ContentParts) > 0 {
			msg.Content = ""
			msg.ContentParts = make([]sdk.ContentPart, len(m.ContentParts))
			for i, part := range m.ContentParts {
				msg.ContentParts[i] = sdk.ContentPart{
					Type:     part.Type,
					Text:     part.Text,
					ImageURL: part.ImageURL,
				}
			}
		}
		msgs = append(msgs, msg)
	}
	if !hasUserMessage {
		msgs = append(msgs, sdk.Message{Role: "user", Content: ""})
	}
	return msgs
}

// hasJSONValue returns true when raw is non-empty and not the JSON literal null.
func hasJSONValue(raw json.RawMessage) bool {
	t := strings.TrimSpace(string(raw))
	return t != "" && t != "null"
}

func applyChatParameters(req *openAIChatRequest, sdkReq *sdk.ChatCompletionRequest) error {
	if req.MaxCompletionTokens != nil {
		sdkReq.MaxTokens = req.MaxCompletionTokens
	} else if req.MaxTokens != nil {
		sdkReq.MaxTokens = req.MaxTokens
	} else if req.MaxLength != nil {
		sdkReq.MaxTokens = req.MaxLength
	}
	sdkReq.Temperature = req.Temperature
	if req.TopP != nil && *req.TopP > 0 {
		sdkReq.TopP = req.TopP
	}
	sdkReq.FrequencyPenalty = req.FrequencyPenalty
	sdkReq.PresencePenalty = req.PresencePenalty
	sdkReq.RepetitionPenalty = req.RepetitionPenalty
	sdkReq.TopK = req.TopK
	sdkReq.Logprobs = req.Logprobs
	if req.Logprobs != nil && *req.Logprobs {
		sdkReq.TopLogprobs = req.TopLogprobs
	}
	if req.Stop.Set {
		sdkReq.Stop = req.Stop.Values
		sdkReq.StopTokenIDs = req.Stop.TokenIDs
	}
	sdkReq.Seed = req.Seed
	if req.N != nil && *req.N > 0 {
		sdkReq.N = req.N
	}
	sdkReq.ParallelToolCalls = req.ParallelToolCalls
	sdkReq.StreamOptions = req.StreamOptions
	sdkReq.ResponseFormat = req.ResponseFormat
	sdkReq.ThinkingTokenBudget = req.ThinkingBudget
	// Legacy functions → tools conversion (deprecated OpenAI v1 field).
	if len(req.Functions) > 0 && string(req.Functions) != "null" {
		if len(req.Tools) > 0 && string(req.Tools) != "null" {
			return fmt.Errorf("cannot specify both 'tools' and 'functions'")
		}
		var funcs []sdk.FunctionSchema
		if err := json.Unmarshal(req.Functions, &funcs); err != nil {
			return fmt.Errorf("invalid functions: %w", err)
		}
		for _, fn := range funcs {
			if fn.Name == "" {
				return fmt.Errorf("each function must have a 'name'")
			}
			sdkReq.Tools = append(sdkReq.Tools, sdk.Tool{
				Type:     "function",
				Function: sdk.FunctionSchema{Name: fn.Name, Description: fn.Description, Parameters: fn.Parameters},
			})
		}
		// Convert legacy function_call → tool_choice
		if len(req.FunctionCall) > 0 && string(req.FunctionCall) != "null" {
			var fcStr string
			if json.Unmarshal(req.FunctionCall, &fcStr) == nil {
				// "auto" / "none"
				tc, _ := json.Marshal(fcStr)
				sdkReq.ToolChoice = tc
			} else {
				// {"name": "X"} → {"type":"function","function":{"name":"X"}}
				var fcObj struct {
					Name string `json:"name"`
				}
				if err := json.Unmarshal(req.FunctionCall, &fcObj); err == nil && fcObj.Name != "" {
					tc, _ := json.Marshal(map[string]any{
						"type":     "function",
						"function": map[string]string{"name": fcObj.Name},
					})
					sdkReq.ToolChoice = tc
				}
			}
		}
	} else if len(req.Tools) > 0 && string(req.Tools) != "null" {
		var tools []sdk.Tool
		if err := json.Unmarshal(req.Tools, &tools); err != nil {
			return fmt.Errorf("invalid tools: %w", err)
		}
		for _, t := range tools {
			if t.Function.Name == "" {
				return fmt.Errorf("each tool function must have a 'name'")
			}
		}
		sdkReq.Tools = tools
	}
	// Only set ToolChoice from req when it wasn't already set by functions→tools conversion.
	if len(sdkReq.ToolChoice) == 0 {
		if len(req.ToolChoice) > 0 && string(req.ToolChoice) != "null" {
			if err := validateToolChoice(req.ToolChoice, sdkReq.Tools); err != nil {
				return err
			}
			sdkReq.ToolChoice = req.ToolChoice
		}
	}
	if len(req.ChatTemplateKwargs) > 0 {
		sdkReq.ChatTemplateKwargs = make(map[string]interface{}, len(req.ChatTemplateKwargs))
		for k, v := range req.ChatTemplateKwargs {
			sdkReq.ChatTemplateKwargs[k] = v
		}
	}
	if req.EnableThinking != nil {
		if sdkReq.ChatTemplateKwargs == nil {
			sdkReq.ChatTemplateKwargs = make(map[string]interface{}, 1)
		}
		if existing, exists := sdkReq.ChatTemplateKwargs["enable_thinking"]; exists {
			existingBool, ok := existing.(bool)
			if !ok {
				return fmt.Errorf("chat_template_kwargs.enable_thinking must be boolean")
			}
			if existingBool != *req.EnableThinking {
				return fmt.Errorf("enable_thinking conflicts with chat_template_kwargs.enable_thinking")
			}
			return nil
		}
		sdkReq.ChatTemplateKwargs["enable_thinking"] = *req.EnableThinking
	}
	return nil
}

// completions handles OpenAI Completions API requests (--backend vllm / openai)
func (h *handler) completions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeOpenAIError(w, http.StatusMethodNotAllowed, "method not allowed", "", "")
		return
	}

	body, err := readRequestBody(w, r)
	if err != nil {
		writeRequestBodyError(w, err)
		return
	}
	defer r.Body.Close()

	var req openAICompletionsRequest
	if err := json.Unmarshal(body, &req); err != nil {
		writeOpenAIError(w, http.StatusBadRequest, jsonUnmarshalMessage(err), "", "")
		return
	}

	// Convert prompt to messages format (user role)
	promptText := req.Prompt.First()
	msgs := []sdk.Message{{Role: "user", Content: promptText}}

	requestID := uuid.New().String()
	respID := "cmpl-" + requestID
	sdkReq := &sdk.ChatCompletionRequest{
		RequestID: requestID,
		Messages:  msgs,
		Stream:    req.Stream,
	}
	// max_completion_tokens takes precedence over max_tokens
	maxTokens := req.MaxCompletionTokens
	if maxTokens == 0 {
		maxTokens = req.MaxTokens
	}
	if maxTokens > 0 {
		sdkReq.MaxTokens = sdk.Int(maxTokens)
	}
	if req.Temperature != 0 {
		sdkReq.Temperature = sdk.Float(req.Temperature)
	}
	if req.TopP != 0 {
		sdkReq.TopP = sdk.Float(req.TopP)
	}

	log.Printf("[INFO] request started id=%s endpoint=/v1/completions stream=%t backend_model=%q", respID, req.Stream, h.model)

	ctx := r.Context()
	client, err := h.pool.acquire(ctx)
	if err != nil {
		writeAcquireError(w, err, requestID)
		return
	}
	defer h.pool.release(client)
	completion, err := client.Chat(ctx, sdkReq)
	if err != nil {
		log.Printf("[WARN] request failed id=%s task_id=%s error=%v", respID, requestID, err)
		writeSDKError(w, err, requestID)
		return
	}
	log.Printf("[INFO] cmaas request mapped id=%s task_id=%s stream=%t", respID, completion.TaskID, req.Stream)

	created := time.Now().Unix()

	if req.Stream {
		h.writeCompletionsStream(w, completion, respID, created, h.model)
	} else {
		h.writeCompletionsResponse(w, completion, respID, created, h.model)
	}
}
