// 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 sdk

import (
	"context"
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestBuildParametersIncludesChatTemplateKwargs(t *testing.T) {
	sessionKey := []byte("0123456789abcdef0123456789abcdef")
	req := &ChatCompletionRequest{
		ThinkingTokenBudget: Int(20),
		ResponseFormat:      json.RawMessage(`{"type":"json_object"}`),
		Logprobs:            Bool(true),
		TopLogprobs:         Int(2),
		ParallelToolCalls:   Bool(true),
		ChatTemplateKwargs: map[string]interface{}{
			"enable_thinking": false,
		},
	}

	params, err := buildParameters(req, sessionKey)
	require.NoError(t, err)

	kwargs := params["chat_template_kwargs"].(map[string]interface{})
	if got := kwargs["enable_thinking"]; got != false {
		t.Fatalf("chat_template_kwargs.enable_thinking not serialized, got %#v", got)
	}
	responseFormat := params["response_format"].(map[string]interface{})
	if got := responseFormat["type"]; got != "json_object" {
		t.Fatalf("response_format not serialized, got %#v", responseFormat)
	}
	if got := params["thinking_token_budget"]; got != float64(20) {
		t.Fatalf("thinking_token_budget not serialized, got %#v", got)
	}
	if got := params["logprobs"]; got != true {
		t.Fatalf("logprobs not serialized, got %#v", got)
	}
	if got := params["top_logprobs"]; got != float64(2) {
		t.Fatalf("top_logprobs not serialized, got %#v", got)
	}
	if got := params["parallel_tool_calls"]; got != true {
		t.Fatalf("parallel_tool_calls not serialized, got %#v", got)
	}
}

func TestBuildParametersEncryptsToolsAndToolChoice(t *testing.T) {
	sessionKey := []byte("0123456789abcdef0123456789abcdef")
	toolChoice := json.RawMessage(`{"type":"function","function":{"name":"get_weather"}}`)
	req := &ChatCompletionRequest{
		Tools: []Tool{{
			Type: "function",
			Function: FunctionSchema{
				Name:        "get_weather",
				Description: "private schema",
				Parameters: map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"city": map[string]interface{}{"type": "string"},
					},
				},
			},
		}},
		ToolChoice: toolChoice,
	}

	params, err := buildParameters(req, sessionKey)
	require.NoError(t, err)

	encryptedTools, ok := params["tools"].(string)
	require.True(t, ok)
	assert.NotContains(t, encryptedTools, "private schema")
	decryptedTools, err := decryptContent(encryptedTools, sessionKey)
	require.NoError(t, err)
	assert.Contains(t, decryptedTools, "private schema")

	encryptedChoice, ok := params["tool_choice"].(string)
	require.True(t, ok)
	assert.NotContains(t, encryptedChoice, "get_weather")
	decryptedChoice, err := decryptContent(encryptedChoice, sessionKey)
	require.NoError(t, err)
	assert.JSONEq(t, string(toolChoice), decryptedChoice)
}

func TestBuildParametersEncryptsStopFields(t *testing.T) {
	sessionKey := []byte("0123456789abcdef0123456789abcdef")
	req := &ChatCompletionRequest{
		Stop:         []string{"private-stop"},
		StopTokenIDs: []int{109266},
	}

	params, err := buildParameters(req, sessionKey)
	require.NoError(t, err)

	encryptedStop, ok := params["stop"].(string)
	require.True(t, ok)
	assert.NotContains(t, encryptedStop, "private-stop")
	decryptedStop, err := decryptContent(encryptedStop, sessionKey)
	require.NoError(t, err)
	assert.JSONEq(t, `["private-stop"]`, decryptedStop)

	encryptedStopTokenIDs, ok := params["stop_token_ids"].(string)
	require.True(t, ok)
	decryptedStopTokenIDs, err := decryptContent(encryptedStopTokenIDs, sessionKey)
	require.NoError(t, err)
	assert.JSONEq(t, `[109266]`, decryptedStopTokenIDs)
}

func TestEncryptOutboundMessagesEncryptsToolCallNameAndArguments(t *testing.T) {
	sessionKey := []byte("0123456789abcdef0123456789abcdef")
	toolCallIndex := 2
	messages := []Message{{
		Role: "assistant",
		ToolCalls: []ToolCall{{
			Index: &toolCallIndex,
			ID:    "call_1",
			Type:  "function",
			Function: FunctionCall{
				Name:      "get_weather",
				Arguments: `{"city":"hangzhou"}`,
			},
		}},
	}}

	encrypted, err := encryptOutboundMessages(context.Background(), messages, sessionKey)
	require.NoError(t, err)
	require.NotNil(t, encrypted[0].ToolCalls[0].Index)
	assert.Equal(t, 2, *encrypted[0].ToolCalls[0].Index)
	assert.Equal(t, "call_1", encrypted[0].ToolCalls[0].ID)
	assert.NotContains(t, encrypted[0].ToolCalls[0].Function.Name, "get_weather")
	assert.NotContains(t, encrypted[0].ToolCalls[0].Function.Arguments, "hangzhou")

	decrypted, err := decryptToolCalls(encrypted[0].ToolCalls, sessionKey)
	require.NoError(t, err)
	require.NotNil(t, decrypted[0].Index)
	assert.Equal(t, 2, *decrypted[0].Index)
	assert.Equal(t, "get_weather", decrypted[0].Function.Name)
	assert.JSONEq(t, `{"city":"hangzhou"}`, decrypted[0].Function.Arguments)
}

// TestEncryptOutboundMessagesEncryptsMultimodalContent verifies multimodal text and image data stay encrypted on the wire.
func TestEncryptOutboundMessagesEncryptsMultimodalContent(t *testing.T) {
	sessionKey := []byte("0123456789abcdef0123456789abcdef")
	imageDataURL := "data:image/png;base64,iVBORw0KGgo="
	messages := []Message{{
		Role: "user",
		ContentParts: []ContentPart{
			{Type: "text", Text: "describe private image"},
			{Type: "image_url", ImageURL: &ImageURL{URL: imageDataURL, Detail: "auto"}},
		},
	}}

	encrypted, err := encryptOutboundMessages(context.Background(), messages, sessionKey)
	require.NoError(t, err)
	require.Len(t, encrypted, 1)
	require.Len(t, encrypted[0].ContentParts, 2)

	assert.Empty(t, encrypted[0].Content)
	assert.NotContains(t, encrypted[0].ContentParts[0].Text, "describe private image")
	assert.NotContains(t, encrypted[0].ContentParts[1].ImageURL.URL, "data:image/png")

	decryptedText, err := decryptContent(encrypted[0].ContentParts[0].Text, sessionKey)
	require.NoError(t, err)
	assert.Equal(t, "describe private image", decryptedText)

	decryptedImage, err := decryptContent(encrypted[0].ContentParts[1].ImageURL.URL, sessionKey)
	require.NoError(t, err)
	assert.Equal(t, imageDataURL, decryptedImage)
	assert.Equal(t, "auto", encrypted[0].ContentParts[1].ImageURL.Detail)
}

func TestInferenceOutputPreservesStreamToolCallIndex(t *testing.T) {
	var output InferenceOutput
	err := json.Unmarshal([]byte(`{
		"choices": [{
			"index": 0,
			"delta": {
				"tool_calls": [{
					"index": 3,
					"id": "call_3",
					"type": "function",
					"function": {"name": "cipher-name", "arguments": "cipher-args"}
				}]
			}
		}]
	}`), &output)
	require.NoError(t, err)
	require.Len(t, output.Choices, 1)
	require.Len(t, output.Choices[0].Delta.ToolCalls, 1)
	require.NotNil(t, output.Choices[0].Delta.ToolCalls[0].Index)
	assert.Equal(t, 3, *output.Choices[0].Delta.ToolCalls[0].Index)
}

func TestExtractResponseContentIncludesReasoningAliases(t *testing.T) {
	output := InferenceOutput{
		Choices: []InferenceChoice{
			{
				Message: &ResponseMessage{
					Role:             "assistant",
					Content:          "cipher-content",
					ReasoningContent: "cipher-reasoning-content",
					Reasoning:        "cipher-reasoning",
				},
			},
		},
	}
	outputJSON, err := json.Marshal(output)
	if err != nil {
		t.Fatalf("marshal output: %v", err)
	}

	resp := &DashScopeMessage{
		Payload: DashScopePayload{
			Output: outputJSON,
			Usage:  json.RawMessage(`{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"completion_tokens_details":{"reasoning_tokens":1,"text_tokens":1}}`),
		},
	}

	parsed, err := parseInferenceResponse(resp)
	if err != nil {
		t.Fatalf("extract response message: %v", err)
	}
	if len(parsed.Choices) != 1 {
		t.Fatalf("choices mismatch: %#v", parsed.Choices)
	}
	first := parsed.Choices[0]
	if first.Content != "cipher-content" {
		t.Fatalf("content mismatch: %q", first.Content)
	}
	if first.ReasoningContent != "cipher-reasoning-content" {
		t.Fatalf("reasoning_content mismatch: %q", first.ReasoningContent)
	}
	if first.Reasoning != "cipher-reasoning" {
		t.Fatalf("reasoning mismatch: %q", first.Reasoning)
	}
	if len(first.Logprobs) != 0 {
		t.Fatalf("logprobs mismatch: %s", string(first.Logprobs))
	}
	if parsed.Usage.PromptTokens != 1 || parsed.Usage.CompletionTokens != 2 || parsed.Usage.TotalTokens != 3 {
		t.Fatalf("usage mismatch: %#v", parsed.Usage)
	}
	if parsed.Usage.CompletionTokensDetails["reasoning_tokens"] != float64(1) {
		t.Fatalf("completion_tokens_details mismatch: %#v", parsed.Usage.CompletionTokensDetails)
	}
}

func TestDecryptLogprobs(t *testing.T) {
	sessionKey := []byte("0123456789abcdef0123456789abcdef")
	rawLogprobs := `{"content":[{"token":"hello","logprob":-0.1}]}`
	encrypted, err := encryptContent(rawLogprobs, sessionKey)
	if err != nil {
		t.Fatalf("encrypt logprobs: %v", err)
	}

	decrypted, err := decryptLogprobs(json.RawMessage(`"`+encrypted+`"`), sessionKey)
	if err != nil {
		t.Fatalf("decrypt logprobs: %v", err)
	}
	if string(decrypted) != rawLogprobs {
		t.Fatalf("logprobs mismatch: %s", string(decrypted))
	}

	decryptedNull, err := decryptLogprobs(json.RawMessage(`null`), sessionKey)
	if err != nil {
		t.Fatalf("decrypt null logprobs: %v", err)
	}
	if string(decryptedNull) != "null" {
		t.Fatalf("null logprobs mismatch: %s", string(decryptedNull))
	}
}

func TestResponseErrorReadsEncryptedStatusMessage(t *testing.T) {
	sessionKey := make([]byte, 32)
	encrypted, err := encryptContent("top_k must be 0 (disable), or at least 1, got -2.", sessionKey)
	require.NoError(t, err)

	resp := &DashScopeMessage{
		Header: DashScopeHeader{
			ErrorMessage: encryptedBackendErrorStatusPrefix + encrypted,
		},
	}

	got := responseError(resp, sessionKey)
	var apiErr *APIError
	require.ErrorAs(t, got, &apiErr)
	assert.Equal(t, 0, apiErr.StatusCode)
	assert.Equal(t, "top_k must be 0 (disable), or at least 1, got -2.", apiErr.Message)
}

func TestResponseErrorEncryptedStatusDecryptFailureIsGeneric(t *testing.T) {
	resp := &DashScopeMessage{
		Header: DashScopeHeader{
			StatusCode:    400,
			StatusMessage: encryptedBackendErrorStatusPrefix + "not-valid-ciphertext",
		},
	}

	got := responseError(resp, make([]byte, 32))
	var apiErr *APIError
	require.ErrorAs(t, got, &apiErr)
	assert.Equal(t, 400, apiErr.StatusCode)
	assert.Equal(t, "An internal error has occured during inference, please try again later or contact service support.", apiErr.Message)
}

func TestResponseErrorFallsBackToHeaderMessage(t *testing.T) {
	resp := &DashScopeMessage{
		Header: DashScopeHeader{
			StatusCode:    500,
			StatusMessage: "An internal error has occured",
		},
	}
	got := responseError(resp, nil)
	var apiErr *APIError
	require.ErrorAs(t, got, &apiErr)
	assert.Equal(t, 500, apiErr.StatusCode)
	assert.Equal(t, "An internal error has occured", apiErr.Message)
}

func TestResponseErrorInfersStatusCodeFromErrorName(t *testing.T) {
	resp := &DashScopeMessage{
		Header: DashScopeHeader{
			ErrorName:    "InvalidParameter",
			ErrorMessage: "Required parameter(s) missing or invalid, please check the request parameters.",
		},
	}
	got := responseError(resp, nil)
	var apiErr *APIError
	require.ErrorAs(t, got, &apiErr)
	assert.Equal(t, 400, apiErr.StatusCode)
	assert.Equal(t, "Required parameter(s) missing or invalid, please check the request parameters.", apiErr.Message)
}

func TestResponseErrorInfersStatusCodeFromErrorCode(t *testing.T) {
	resp := &DashScopeMessage{
		Header: DashScopeHeader{
			ErrorCode:    "InvalidParameter",
			ErrorMessage: "Required parameter(s) missing or invalid, please check the request parameters.",
		},
	}
	got := responseError(resp, nil)
	var apiErr *APIError
	require.ErrorAs(t, got, &apiErr)
	assert.Equal(t, 400, apiErr.StatusCode)
	assert.Equal(t, "Required parameter(s) missing or invalid, please check the request parameters.", apiErr.Message)
}

func TestResponseErrorInfersStatusCodeFromStatusName(t *testing.T) {
	resp := &DashScopeMessage{
		Header: DashScopeHeader{
			StatusName:    "UnprocessableEntity",
			StatusMessage: "Key negotiation required.",
		},
	}
	got := responseError(resp, nil)
	var apiErr *APIError
	require.ErrorAs(t, got, &apiErr)
	assert.Equal(t, 422, apiErr.StatusCode)
	assert.Equal(t, "Key negotiation required.", apiErr.Message)
}

func TestResponseErrorWithoutStatusOrCodeLeavesStatusCodeUnset(t *testing.T) {
	resp := &DashScopeMessage{
		Header: DashScopeHeader{
			ErrorMessage: "Required parameter(s) missing or invalid, please check the request parameters.",
		},
	}
	got := responseError(resp, nil)
	var apiErr *APIError
	require.ErrorAs(t, got, &apiErr)
	assert.Equal(t, 0, apiErr.StatusCode)
	assert.Equal(t, "Required parameter(s) missing or invalid, please check the request parameters.", apiErr.Message)
}

func TestIsErrorResponseDetectsStatusCode(t *testing.T) {
	assert.True(t, isErrorResponse(&DashScopeMessage{Header: DashScopeHeader{StatusCode: 400}}))
	assert.True(t, isErrorResponse(&DashScopeMessage{Header: DashScopeHeader{Event: "task-failed"}}))
	assert.True(t, isErrorResponse(&DashScopeMessage{Header: DashScopeHeader{StatusName: "InvalidParameter"}}))
	assert.False(t, isErrorResponse(&DashScopeMessage{Header: DashScopeHeader{StatusCode: 200, StatusName: "Success"}}))
	assert.False(t, isErrorResponse(&DashScopeMessage{}))
}
