// 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"
	"log"
	"net/http"
	"strings"

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

	"github.com/google/uuid"
)

func (h *handler) writeResponse(w http.ResponseWriter, completion *sdk.ChatCompletion, id string, created int64, model string, toolChoiceNone, legacyMode, singleToolOnly bool) {
	toolCalls := ensureToolCallIDs(completion.ToolCalls)
	if toolChoiceNone {
		toolCalls = nil
	}
	if singleToolOnly && len(toolCalls) > 1 {
		toolCalls = toolCalls[:1]
	}

	choices := openAIResponseChoices(completion)
	if len(toolCalls) > 0 && len(choices) > 0 {
		// Apply tool_calls (or legacy function_call shape) to the first
		// choice; multi-choice responses don't realistically combine with
		// tool calling for the models on this backend.
		var finish string
		if legacyMode {
			first := toolCalls[0]
			choices[0].Message.FunctionCall = &openAIFunctionCall{
				Name:      first.Function.Name,
				Arguments: first.Function.Arguments,
			}
			finish = "function_call"
		} else {
			choices[0].Message.ToolCalls = toOutboundToolCalls(toolCalls)
			finish = "tool_calls"
		}
		if choices[0].FinishReason == nil || *choices[0].FinishReason == "" {
			choices[0].FinishReason = &finish
		} else if legacyMode && *choices[0].FinishReason == "tool_calls" {
			choices[0].FinishReason = &finish
		}
	}
	if toolChoiceNone {
		for i := range choices {
			if choices[i].FinishReason != nil && (*choices[i].FinishReason == "tool_calls" || *choices[i].FinishReason == "function_call") {
				stop := "stop"
				choices[i].FinishReason = &stop
			}
		}
	}

	resp := openAIResponse{
		ID:      id,
		Object:  "chat.completion",
		Created: created,
		Model:   model,
		Choices: choices,
		Usage:   *completionUsage(completion),
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

// toOutboundToolCalls converts SDK tool calls into the response shape with
// `index` and a defaulted `type`. Some consumers expect `index`
// on every tool_calls entry even in non-stream responses (OpenAI spec
// itself doesn't strictly require it, but we emit it for compatibility).
func toOutboundToolCalls(in []sdk.ToolCall) []openAIToolCall {
	if len(in) == 0 {
		return nil
	}
	out := make([]openAIToolCall, len(in))
	for i, tc := range in {
		idx := i
		t := tc.Type
		if t == "" {
			t = "function"
		}
		out[i] = openAIToolCall{
			Index:    &idx,
			ID:       tc.ID,
			Type:     t,
			Function: tc.Function,
		}
	}
	return out
}

// validateToolChoice enforces the OpenAI tool_choice contract at the proxy
// layer so malformed client requests fail fast with 400 instead of being
// passed through to the backend (where they either 500 or silently produce
// unexpected behavior).
//
// Allowed shapes:
//   - string: "auto" / "none" / "required"
//   - object: {"type":"function","function":{"name":<name>}} where <name>
//     refers to a tool actually present in tools[].
func validateToolChoice(raw json.RawMessage, tools []sdk.Tool) error {
	var asStr string
	if json.Unmarshal(raw, &asStr) == nil {
		switch asStr {
		case "auto", "none", "required":
			return nil
		default:
			return fmt.Errorf("tool_choice string must be one of: auto, none, required (got %q)", asStr)
		}
	}
	var obj struct {
		Type     string `json:"type"`
		Function struct {
			Name string `json:"name"`
		} `json:"function"`
	}
	if err := json.Unmarshal(raw, &obj); err != nil {
		return fmt.Errorf("tool_choice must be a string or an object: %w", err)
	}
	if obj.Type != "function" {
		return fmt.Errorf("tool_choice object type must be 'function' (got %q)", obj.Type)
	}
	if obj.Function.Name == "" {
		return fmt.Errorf("tool_choice object requires function.name")
	}
	for _, t := range tools {
		if t.Function.Name == obj.Function.Name {
			return nil
		}
	}
	return fmt.Errorf("tool_choice references function %q not present in tools", obj.Function.Name)
}

// isToolChoiceNone returns true when the serialized tool_choice value is the
// JSON string "none".
func isToolChoiceNone(raw json.RawMessage) bool {
	if len(raw) == 0 {
		return false
	}
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return s == "none"
	}
	return false
}

// ensureToolCallIDs returns a copy of tool calls with any missing ID fields
// filled in with generated "call_<uuid>" values.
func ensureToolCallIDs(calls []sdk.ToolCall) []sdk.ToolCall {
	if len(calls) == 0 {
		return nil
	}
	out := make([]sdk.ToolCall, len(calls))
	copy(out, calls)
	for i := range out {
		if out[i].ID == "" {
			out[i].ID = "call_" + uuid.New().String()
		}
	}
	return out
}

func openAIResponseChoices(completion *sdk.ChatCompletion) []openAIChoice {
	if len(completion.Choices) > 0 {
		choices := make([]openAIChoice, 0, len(completion.Choices))
		for _, choice := range completion.Choices {
			var finishReason *string
			if choice.FinishReason != "" {
				finishReasonValue := choice.FinishReason
				finishReason = &finishReasonValue
			}
			choices = append(choices, openAIChoice{
				Index: choice.Index,
				Message: &openAIMessage{
					Role:             "assistant",
					Content:          choice.Content,
					ReasoningContent: normalizedChoiceReasoningContent(choice),
				},
				FinishReason: finishReason,
				Logprobs:     openAILogprobs(choice.Logprobs),
			})
		}
		return choices
	}

	var finishReason *string
	if completion.FinishReason != "" {
		finishReason = &completion.FinishReason
	}
	return []openAIChoice{{
		Index:        0,
		Message:      &openAIMessage{Role: "assistant", Content: completion.Content, ReasoningContent: normalizedReasoningContent(completion)},
		FinishReason: finishReason,
		Logprobs:     openAILogprobs(completion.Logprobs),
	}}
}

func normalizedReasoningContent(completion *sdk.ChatCompletion) string {
	if completion.ReasoningContent != "" {
		return completion.ReasoningContent
	}
	return completion.Reasoning
}

func normalizedChoiceReasoningContent(choice sdk.ChatCompletionChoice) string {
	if choice.ReasoningContent != "" {
		return choice.ReasoningContent
	}
	return choice.Reasoning
}

func includeStreamUsage(opts *sdk.StreamOptions) bool {
	// ModelStudio default: always include usage in streaming responses
	// unless the client explicitly sets include_usage=false.
	if opts == nil {
		return true
	}
	return opts.IncludeUsage
}

func openAILogprobs(raw json.RawMessage) interface{} {
	if len(raw) == 0 || string(raw) == "null" {
		return nil
	}
	var value interface{}
	if err := json.Unmarshal(raw, &value); err != nil {
		return nil
	}
	return value
}

func normalizedDeltaReasoningContent(delta sdk.StreamDelta) string {
	if delta.ReasoningContent != "" {
		return delta.ReasoningContent
	}
	return delta.Reasoning
}

func (h *handler) writeStream(w http.ResponseWriter, completion *sdk.ChatCompletion, id string, created int64, model string, includeUsage, toolChoiceNone, legacyMode, singleToolOnly bool) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeOpenAIError(w, http.StatusInternalServerError, "streaming not supported", "", "")
		return
	}

	streamStarted := false

	writeChunk := func(chunk openAIChunk) {
		if !streamStarted {
			w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
			w.Header().Set("Cache-Control", "no-cache")
			w.Header().Set("Connection", "keep-alive")
			w.Header().Set("X-Accel-Buffering", "no")
			streamStarted = true
		}
		data, _ := json.Marshal(chunk)
		fmt.Fprintf(w, "data: %s\n\n", data)
		flusher.Flush()
	}

	roleSent := map[int]bool{}
	includeRoleForDelta := func(index int) bool {
		return !roleSent[index]
	}
	markRoleSent := func(index int) {
		roleSent[index] = true
	}
	writeRoleChunk := func(index int) {
		if !includeRoleForDelta(index) {
			return
		}
		markRoleSent(index)
		writeChunk(openAIChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []openAIChoice{
				{Index: index, Delta: map[string]any{"role": "assistant", "content": ""}, FinishReason: nil},
			},
		})
	}

	streamDelta := func(role, content, reasoning string) map[string]any {
		delta := map[string]any{}
		if role != "" {
			delta["role"] = role
		}
		if content != "" {
			delta["content"] = content
		}
		if reasoning != "" {
			delta["reasoning_content"] = reasoning
		}
		return delta
	}

	// Legacy v1 streaming response shape: model emits one (and only one)
	// function_call. We track whether the name has been emitted so
	// subsequent fragments only carry `function_call.arguments` deltas
	// per the v1 SSE contract.
	legacyNameEmitted := false
	emitLegacyFunctionCallDelta := func(d sdk.ToolCallDelta, includeRole bool) {
		fc := map[string]any{}
		if d.Name != "" && !legacyNameEmitted {
			fc["name"] = d.Name
			legacyNameEmitted = true
		}
		if d.Arguments != "" {
			fc["arguments"] = d.Arguments
		}
		if len(fc) == 0 {
			return
		}
		delta := map[string]any{"function_call": fc}
		if includeRole {
			delta["role"] = "assistant"
			markRoleSent(d.ChoiceIndex)
		}
		chunk := map[string]any{
			"id":      id,
			"object":  "chat.completion.chunk",
			"created": created,
			"model":   model,
			"choices": []map[string]any{{
				"index":         d.ChoiceIndex,
				"delta":         delta,
				"finish_reason": nil,
			}},
		}
		data, _ := json.Marshal(chunk)
		fmt.Fprintf(w, "data: %s\n\n", data)
		flusher.Flush()
	}

	// Track generated IDs for tool-call deltas by index so that
	// subsequent fragments for the same index reuse the same ID.
	toolCallIDsByIndex := map[[2]int]string{}

	// emitToolCallDelta writes a streaming chunk for a tool-call fragment.
	emitToolCallDelta := func(d sdk.ToolCallDelta, includeRole bool) {
		if d.Name == "" && d.Arguments == "" {
			return
		}
		tc := map[string]any{"index": d.Index}
		key := [2]int{d.ChoiceIndex, d.Index}
		tcID := d.ID
		isFirstFragment := false
		if tcID == "" {
			if existing, ok := toolCallIDsByIndex[key]; ok {
				tcID = existing
			} else {
				tcID = "call_" + uuid.New().String()
				toolCallIDsByIndex[key] = tcID
				isFirstFragment = true
			}
		} else {
			if _, existed := toolCallIDsByIndex[key]; !existed {
				isFirstFragment = true
			}
			toolCallIDsByIndex[key] = tcID
		}
		tc["id"] = tcID
		// type is required on every fragment per OpenAI spec. Some backends
		// only send it on the first fragment per index; we fill
		// it for subsequent ones with the OpenAI-standard "function".
		tcType := d.Type
		if tcType == "" {
			tcType = "function"
		}
		tc["type"] = tcType
		fn := map[string]any{}
		if d.Name != "" {
			fn["name"] = d.Name
		}
		// Always emit `arguments` on the first fragment per index, even if
		// the model called a parameterless function (so the accumulator
		// surfaces it as `""`). Without this, downstream history would
		// preserve a malformed `function:{name:...}` with no arguments
		// field at all, which the inbound validator (and ModelStudio) reject
		// when replayed.
		if d.Arguments != "" || isFirstFragment {
			fn["arguments"] = d.Arguments
		}
		if len(fn) > 0 {
			tc["function"] = fn
		}
		delta := map[string]any{"tool_calls": []map[string]any{tc}}
		if includeRole {
			delta["role"] = "assistant"
			markRoleSent(d.ChoiceIndex)
		}
		chunk := map[string]any{
			"id":      id,
			"object":  "chat.completion.chunk",
			"created": created,
			"model":   model,
			"choices": []map[string]any{{
				"index":         d.ChoiceIndex,
				"delta":         delta,
				"finish_reason": nil,
			}},
		}
		data, _ := json.Marshal(chunk)
		fmt.Fprintf(w, "data: %s\n\n", data)
		flusher.Flush()
	}

	// Multiplex content (Deltas or Chunks, mutually exclusive in the SDK)
	// with tool-call deltas (always on a separate channel). With
	// StructuredStream=true (req.Stream=true on this proxy) the SDK feeds
	// completion.Deltas; without it, completion.Chunks. Tool deltas always
	// arrive on completion.ToolCallChunks.
	deltas := completion.Deltas
	chunks := completion.Chunks
	if deltas != nil {
		// When the structured channel is in use, the SDK does not feed
		// completion.Chunks, so don't read from it (would block forever).
		chunks = nil
	}
	toolChunks := completion.ToolCallChunks
	sawToolCall := false
	for deltas != nil || chunks != nil || toolChunks != nil {
		select {
		case delta, ok := <-deltas:
			if !ok {
				deltas = nil
				continue
			}
			logprobsText := strings.TrimSpace(string(delta.Logprobs))
			if delta.Content == "" && delta.ReasoningContent == "" && delta.Reasoning == "" && (len(delta.Logprobs) == 0 || logprobsText == "" || logprobsText == "null") {
				continue
			}
			writeRoleChunk(delta.Index)
			writeChunk(openAIChunk{
				ID:      id,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   model,
				Choices: []openAIChoice{
					{Index: delta.Index, Delta: streamDelta("", delta.Content, normalizedDeltaReasoningContent(delta)), FinishReason: nil, Logprobs: openAILogprobs(delta.Logprobs)},
				},
			})
		case text, ok := <-chunks:
			if !ok {
				chunks = nil
				continue
			}
			if text == "" {
				continue
			}
			writeRoleChunk(0)
			writeChunk(openAIChunk{
				ID:      id,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   model,
				Choices: []openAIChoice{
					{Index: 0, Delta: streamDelta("", text, ""), FinishReason: nil},
				},
			})
		case d, ok := <-toolChunks:
			if !ok {
				toolChunks = nil
				continue
			}
			if toolChoiceNone {
				continue // suppress tool calls when tool_choice is "none"
			}
			// parallel_tool_calls=false: keep only the first tool-call index.
			if singleToolOnly && d.Index > 0 {
				continue
			}
			sawToolCall = true
			includeRole := includeRoleForDelta(d.ChoiceIndex)
			if legacyMode {
				emitLegacyFunctionCallDelta(d, includeRole)
			} else {
				emitToolCallDelta(d, includeRole)
			}
		}
	}

	// Check for streaming error
	if completion.Error != nil {
		log.Printf("[WARN] streaming error id=%s task_id=%s error=%v", id, completion.TaskID, completion.Error)
		if !streamStarted {
			writeSDKError(w, completion.Error, completion.TaskID)
			return
		}
	}

	// Final chunk: finish_reason. OpenAI-compatible usage is sent as a separate
	// chunk with choices=[] when stream_options.include_usage is true.
	writeRoleChunk(0)
	finishChoices := completion.FinishChoices
	if len(finishChoices) == 0 {
		finish := completion.FinishReason
		if finish == "" {
			finish = "stop"
		}
		finishChoices = []sdk.ChatCompletionChoice{{Index: 0, FinishReason: finish, Logprobs: completion.FinishLogprobs}}
	}
	choices := make([]openAIChoice, 0, len(finishChoices))
	for _, choice := range finishChoices {
		finish := choice.FinishReason
		if toolChoiceNone && (finish == "" || finish == "tool_calls" || finish == "function_call") {
			finish = "stop"
		} else if legacyMode && finish == "tool_calls" {
			finish = "function_call"
		}
		if finish == "" {
			if legacyMode {
				finish = "function_call"
			} else if sawToolCall {
				finish = "tool_calls"
			} else {
				finish = "stop"
			}
		}
		choices = append(choices, openAIChoice{
			Index:        choice.Index,
			Delta:        map[string]interface{}{},
			FinishReason: &finish,
			Logprobs:     openAILogprobs(choice.Logprobs),
		})
	}
	for _, choice := range choices {
		writeChunk(openAIChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []openAIChoice{choice},
		})
	}

	if includeUsage {
		usageChunk := openAIChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []openAIChoice{},
			Usage:   completionUsage(completion),
		}
		writeChunk(usageChunk)
	}

	fmt.Fprintf(w, "data: [DONE]\n\n")
	flusher.Flush()
}

func completionUsage(completion *sdk.ChatCompletion) *openAIUsage {
	return &openAIUsage{
		PromptTokens:            completion.Usage.PromptTokens,
		CompletionTokens:        completion.Usage.CompletionTokens,
		TotalTokens:             completion.Usage.TotalTokens,
		PromptTokensDetails:     completion.Usage.PromptTokensDetails,
		CompletionTokensDetails: completion.Usage.CompletionTokensDetails,
	}
}

// writeCompletionsResponse writes non-streaming response for Completions API
func (h *handler) writeCompletionsResponse(w http.ResponseWriter, completion *sdk.ChatCompletion, id string, created int64, model string) {
	stopStr := "stop"
	resp := openAIResponse{
		ID:      id,
		Object:  "text_completion",
		Created: created,
		Model:   model,
		Choices: []openAIChoice{
			{
				Index:        0,
				Text:         completion.Content,
				FinishReason: &stopStr,
			},
		},
		Usage: *completionUsage(completion),
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

// writeCompletionsStream writes streaming response for Completions API
func (h *handler) writeCompletionsStream(w http.ResponseWriter, completion *sdk.ChatCompletion, id string, created int64, model string) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeOpenAIError(w, http.StatusInternalServerError, "streaming not supported", "", "")
		return
	}

	streamStarted := false

	writeChunk := func(chunk openAIChunk) {
		if !streamStarted {
			w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
			w.Header().Set("Cache-Control", "no-cache")
			w.Header().Set("Connection", "keep-alive")
			w.Header().Set("X-Accel-Buffering", "no")
			streamStarted = true
		}
		data, _ := json.Marshal(chunk)
		fmt.Fprintf(w, "data: %s\n\n", data)
		flusher.Flush()
	}

	// Content chunks from SDK (no role chunk for completions API)
	for text := range completion.Chunks {
		if text == "" {
			continue
		}
		contentChunk := openAIChunk{
			ID:      id,
			Object:  "text_completion.chunk",
			Created: created,
			Model:   model,
			Choices: []openAIChoice{
				{Index: 0, Text: text, FinishReason: nil},
			},
		}
		writeChunk(contentChunk)
	}

	// Check for streaming error
	if completion.Error != nil {
		log.Printf("[WARN] streaming error id=%s task_id=%s error=%v", id, completion.TaskID, completion.Error)
		if !streamStarted {
			writeSDKError(w, completion.Error, completion.TaskID)
			return
		}
	}

	// Final chunk: finish_reason + usage
	stopStr := "stop"
	finalChunk := openAIChunk{
		ID:      id,
		Object:  "text_completion.chunk",
		Created: created,
		Model:   model,
		Choices: []openAIChoice{
			{Index: 0, Text: "", FinishReason: &stopStr},
		},
		Usage: completionUsage(completion),
	}
	writeChunk(finalChunk)

	fmt.Fprintf(w, "data: [DONE]\n\n")
	flusher.Flush()
}
