// 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"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/patrickmn/go-cache"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestChatRejectsNilRequest(t *testing.T) {
	client := &Client{}

	_, err := client.Chat(context.Background(), nil)

	require.Error(t, err)
	assert.IsType(t, &ValidationError{}, err)
	assert.Equal(t, "request is required", err.Error())
}

// TestPubKeyCacheExpiry verifies cached public keys expire after their configured TTL.
func TestPubKeyCacheExpiry(t *testing.T) {
	client := &Client{
		pubKeyCacheTTL: 100 * time.Millisecond, // Short TTL for testing
		pubKeyCache:    cache.New(100*time.Millisecond, 50*time.Millisecond),
	}

	pubKey := "test-pubkey"
	result := &AttestationResult{
		VerifiedAt: time.Now(),
	}

	// Add to cache
	client.addPubKeyCache(pubKey, result)

	// Should find it immediately
	cached := client.findPubKeyCache(pubKey)
	require.NotNil(t, cached)

	// Wait for expiry + cleanup interval
	time.Sleep(200 * time.Millisecond)

	// Should not find it after expiry
	cached = client.findPubKeyCache(pubKey)
	assert.Nil(t, cached)

	// Cache should be empty after cleanup
	assert.Equal(t, 0, client.pubKeyCache.ItemCount())
}

// TestPubKeyCacheNoDuplicate verifies adding the same public key updates one cache entry.
func TestPubKeyCacheNoDuplicate(t *testing.T) {
	client := &Client{
		pubKeyCacheTTL: 1 * time.Hour,
		pubKeyCache:    cache.New(1*time.Hour, 2*time.Hour),
	}

	pubKey := "test-pubkey"
	result := &AttestationResult{
		VerifiedAt: time.Now(),
	}

	// Add same key twice
	client.addPubKeyCache(pubKey, result)
	client.addPubKeyCache(pubKey, result)

	// Should only have one entry (go-cache overwrites on duplicate Set)
	assert.Equal(t, 1, client.pubKeyCache.ItemCount())
}

// TestPubKeyCacheMultipleKeys verifies distinct server public keys are cached independently.
func TestPubKeyCacheMultipleKeys(t *testing.T) {
	client := &Client{
		pubKeyCacheTTL: 1 * time.Hour,
		pubKeyCache:    cache.New(1*time.Hour, 2*time.Hour),
	}

	// Add multiple keys (simulating load balancer with multiple backends)
	for i := 0; i < 5; i++ {
		pubKey := string(rune('A' + i)) // "A", "B", "C", "D", "E"
		result := &AttestationResult{
			VerifiedAt: time.Now(),
		}
		client.addPubKeyCache(pubKey, result)
	}

	// Should have 5 entries
	assert.Equal(t, 5, client.pubKeyCache.ItemCount())

	// Should find each key
	for i := 0; i < 5; i++ {
		pubKey := string(rune('A' + i))
		cached := client.findPubKeyCache(pubKey)
		assert.NotNil(t, cached)
	}
}

func TestBuildParametersOmitsRequestID(t *testing.T) {
	params, err := buildParameters(&ChatCompletionRequest{
		RequestID: "req-123",
		Stream:    true,
		MaxTokens: Int(8),
	}, []byte("01234567890123456789012345678901"))
	require.NoError(t, err)

	if _, ok := params["request_id"]; ok {
		t.Fatal("request_id must not be sent as a model parameter")
	}
	if _, ok := params["RequestID"]; ok {
		t.Fatal("RequestID must not be sent as a model parameter")
	}
	assert.Equal(t, float64(8), params["max_tokens"])
}

func TestExtractResponseContentIncludesReasoningContent(t *testing.T) {
	sessionKey := []byte("01234567890123456789012345678901")
	encryptedContent, err := encryptContent("answer", sessionKey)
	require.NoError(t, err)
	encryptedReasoning, err := encryptContent("thinking", sessionKey)
	require.NoError(t, err)
	encryptedSecondContent, err := encryptContent("second answer", sessionKey)
	require.NoError(t, err)

	output, err := json.Marshal(InferenceOutput{
		Choices: []InferenceChoice{{
			Index:        0,
			FinishReason: "length",
			Message: &ResponseMessage{
				Role:             "assistant",
				Content:          encryptedContent,
				ReasoningContent: encryptedReasoning,
			},
		}, {
			Index:        1,
			FinishReason: "stop",
			Message: &ResponseMessage{
				Role:    "assistant",
				Content: encryptedSecondContent,
			},
		}},
	})
	require.NoError(t, err)
	usage, err := json.Marshal(TokenUsage{PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3})
	require.NoError(t, err)

	parsed, err := parseInferenceResponse(&DashScopeMessage{
		Payload: DashScopePayload{
			Output: output,
			Usage:  usage,
		},
	})
	require.NoError(t, err)
	require.Len(t, parsed.Choices, 2)
	first := parsed.Choices[0]

	decryptedContent, err := decryptContent(first.Content, sessionKey)
	require.NoError(t, err)
	decryptedReasoningContent, err := decryptContent(first.ReasoningContent, sessionKey)
	require.NoError(t, err)

	assert.Equal(t, "answer", decryptedContent)
	assert.Equal(t, "thinking", decryptedReasoningContent)
	assert.Empty(t, first.Reasoning)
	assert.Equal(t, "length", first.FinishReason)
	assert.Equal(t, 0, parsed.Choices[0].Index)
	assert.Equal(t, "length", parsed.Choices[0].FinishReason)
	assert.Equal(t, 1, parsed.Choices[1].Index)
	assert.Equal(t, "stop", parsed.Choices[1].FinishReason)
	assert.Equal(t, TokenUsage{PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3}, parsed.Usage)
}

func TestExtractResponseContentWithReasoningField(t *testing.T) {
	sessionKey := []byte("01234567890123456789012345678901")
	encryptedContent, err := encryptContent("answer", sessionKey)
	require.NoError(t, err)
	encryptedReasoning, err := encryptContent("thinking", sessionKey)
	require.NoError(t, err)

	output := []byte(`{"choices":[{"message":{"role":"assistant","content":` +
		strconv.Quote(encryptedContent) + `,"reasoning":` + strconv.Quote(encryptedReasoning) + `}}]}`)

	parsed, err := parseInferenceResponse(&DashScopeMessage{
		Payload: DashScopePayload{Output: output},
	})
	require.NoError(t, err)
	require.Len(t, parsed.Choices, 1)
	first := parsed.Choices[0]

	decryptedContent, err := decryptContent(first.Content, sessionKey)
	require.NoError(t, err)
	decryptedReasoning, err := decryptContent(first.Reasoning, sessionKey)
	require.NoError(t, err)

	assert.Equal(t, "answer", decryptedContent)
	assert.Empty(t, first.ReasoningContent)
	assert.Equal(t, "thinking", decryptedReasoning)
}

func TestExtractResponseContentWithBothReasoningFields(t *testing.T) {
	sessionKey := []byte("01234567890123456789012345678901")
	encryptedContent, err := encryptContent("answer", sessionKey)
	require.NoError(t, err)
	encryptedReasoningContent, err := encryptContent("thinking-content", sessionKey)
	require.NoError(t, err)
	encryptedReasoning, err := encryptContent("thinking", sessionKey)
	require.NoError(t, err)

	output, err := json.Marshal(InferenceOutput{
		Choices: []InferenceChoice{{
			Message: &ResponseMessage{
				Role:             "assistant",
				Content:          encryptedContent,
				ReasoningContent: encryptedReasoningContent,
				Reasoning:        encryptedReasoning,
			},
		}},
	})
	require.NoError(t, err)

	parsed, err := parseInferenceResponse(&DashScopeMessage{
		Payload: DashScopePayload{Output: output},
	})
	require.NoError(t, err)
	require.Len(t, parsed.Choices, 1)
	first := parsed.Choices[0]

	decryptedContent, err := decryptContent(first.Content, sessionKey)
	require.NoError(t, err)
	decryptedReasoningContent, err := decryptContent(first.ReasoningContent, sessionKey)
	require.NoError(t, err)
	decryptedReasoning, err := decryptContent(first.Reasoning, sessionKey)
	require.NoError(t, err)

	assert.Equal(t, "answer", decryptedContent)
	assert.Equal(t, "thinking-content", decryptedReasoningContent)
	assert.Equal(t, "thinking", decryptedReasoning)
}

// TestChatMutexSerializesCalls verifies that one client allows only one active Chat lifecycle.
func TestChatMutexSerializesCalls(t *testing.T) {
	client := &Client{}

	client.chatMu.Lock()
	acquiredSecond := make(chan struct{})

	go func() {
		client.chatMu.Lock()
		defer client.chatMu.Unlock()
		close(acquiredSecond)
	}()

	select {
	case <-acquiredSecond:
		require.FailNow(t, "second chat acquired lock before first released it")
	case <-time.After(25 * time.Millisecond):
	}

	client.chatMu.Unlock()

	select {
	case <-acquiredSecond:
	case <-time.After(time.Second):
		require.FailNow(t, "second chat did not acquire lock after first released it")
	}
}

// TestCloseWaitsForActiveChat verifies Close does not close the session while a Chat call is active.
func TestCloseWaitsForActiveChat(t *testing.T) {
	session := newSession(nil, []byte("01234567890123456789012345678901"), nil, "", "task-1")
	client := &Client{currentSession: session}
	client.chatMu.Lock()
	closed := make(chan error)

	go func() {
		closed <- client.Close()
	}()

	select {
	case err := <-closed:
		require.NoError(t, err)
		require.FailNow(t, "Close returned before active Chat lock was released")
	case <-time.After(25 * time.Millisecond):
	}

	assert.NotNil(t, client.currentSession)
	client.chatMu.Unlock()

	select {
	case err := <-closed:
		require.NoError(t, err)
	case <-time.After(time.Second):
		require.FailNow(t, "Close did not return after active Chat lock was released")
	}
	assert.Nil(t, client.currentSession)
}

// TestChatStreamHoldsLockUntilChunksClose verifies streaming Chat holds the client lock until the stream ends.
func TestChatStreamHoldsLockUntilChunksClose(t *testing.T) {
	sessionKey := []byte("01234567890123456789012345678901")
	finishStream := make(chan struct{})
	serverErrors := make(chan error, 8)
	reportServerError := func(err error) bool {
		if err == nil {
			return false
		}
		serverErrors <- err
		return true
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upgrader := websocket.Upgrader{}
		conn, err := upgrader.Upgrade(w, r, nil)
		if reportServerError(err) {
			return
		}
		defer conn.Close()

		_, _, err = conn.ReadMessage()
		if reportServerError(err) {
			return
		}

		encryptedChunk, err := encryptContent("hello", sessionKey)
		if reportServerError(err) {
			return
		}
		output, err := json.Marshal(InferenceOutput{
			Choices: []InferenceChoice{{Delta: &ResponseMessage{Content: encryptedChunk}}},
		})
		if reportServerError(err) {
			return
		}
		err = conn.WriteJSON(DashScopeMessage{
			Header:  DashScopeHeader{Event: "result-generated"},
			Payload: DashScopePayload{Output: output},
		})
		if reportServerError(err) {
			return
		}

		<-finishStream

		output, err = json.Marshal(InferenceOutput{
			Choices: []InferenceChoice{{FinishReason: "stop"}},
		})
		if reportServerError(err) {
			return
		}
		err = conn.WriteJSON(DashScopeMessage{
			Header:  DashScopeHeader{Event: "result-generated"},
			Payload: DashScopePayload{Output: output},
		})
		_ = reportServerError(err)
	}))
	defer server.Close()

	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
	conn, _, err := websocket.DefaultDialer.DialContext(context.Background(), wsURL, nil)
	require.NoError(t, err)
	defer conn.Close()

	session := newSession(conn, sessionKey, nil, "", "task-1")
	client := &Client{
		currentSession: session,
		model:          "test-model",
		reuseSession:   true,
	}

	req := &ChatCompletionRequest{
		Messages: []Message{{Role: "user", Content: "hi"}},
		Stream:   true,
	}
	completion, err := client.Chat(context.Background(), req)
	require.NoError(t, err)
	require.NotNil(t, completion)
	assert.Nil(t, req.StreamOptions)

	select {
	case chunk := <-completion.Chunks:
		assert.Equal(t, "hello", chunk)
	case <-time.After(time.Second):
		require.FailNow(t, "stream did not return the first chunk")
	}

	acquiredSecond := make(chan struct{})
	go func() {
		client.chatMu.Lock()
		defer client.chatMu.Unlock()
		close(acquiredSecond)
	}()

	select {
	case <-acquiredSecond:
		require.FailNow(t, "chat lock was released before stream completed")
	case <-time.After(25 * time.Millisecond):
	}

	close(finishStream)
	for range completion.Chunks {
	}

	select {
	case <-acquiredSecond:
	case <-time.After(time.Second):
		require.FailNow(t, "chat lock was not released after stream completed")
	}
	select {
	case err := <-serverErrors:
		require.NoError(t, err)
	default:
	}
}

func TestChatStreamProcessesToolCallDeltaBeforeFinish(t *testing.T) {
	sessionKey := []byte("01234567890123456789012345678901")
	serverErrors := make(chan error, 8)
	reportServerError := func(err error) bool {
		if err == nil {
			return false
		}
		serverErrors <- err
		return true
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upgrader := websocket.Upgrader{}
		conn, err := upgrader.Upgrade(w, r, nil)
		if reportServerError(err) {
			return
		}
		defer conn.Close()

		_, _, err = conn.ReadMessage()
		if reportServerError(err) {
			return
		}

		name, err := encryptContent("get_current_weather", sessionKey)
		if reportServerError(err) {
			return
		}
		argsHead, err := encryptContent(`{"location": "Boston, MA"`, sessionKey)
		if reportServerError(err) {
			return
		}
		argsTail, err := encryptContent(`}`, sessionKey)
		if reportServerError(err) {
			return
		}

		toolCallIndex := 0
		output, err := json.Marshal(InferenceOutput{
			Choices: []InferenceChoice{{
				Index: 0,
				Delta: &ResponseMessage{ToolCalls: []ToolCall{{
					Index: &toolCallIndex,
					ID:    "call_1",
					Type:  "function",
					Function: FunctionCall{
						Name:      name,
						Arguments: argsHead,
					},
				}}},
			}},
		})
		if reportServerError(err) {
			return
		}
		err = conn.WriteJSON(DashScopeMessage{
			Header:  DashScopeHeader{Event: "result-generated"},
			Payload: DashScopePayload{Output: output},
		})
		if reportServerError(err) {
			return
		}

		output, err = json.Marshal(InferenceOutput{
			Choices: []InferenceChoice{{
				Index:        0,
				FinishReason: "tool_calls",
				Delta: &ResponseMessage{ToolCalls: []ToolCall{{
					Index: &toolCallIndex,
					Function: FunctionCall{
						Arguments: argsTail,
					},
				}}},
			}},
		})
		if reportServerError(err) {
			return
		}
		err = conn.WriteJSON(DashScopeMessage{
			Header: DashScopeHeader{Event: "result-generated"},
			Payload: DashScopePayload{
				Output: output,
				Usage:  json.RawMessage(`{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}`),
			},
		})
		_ = reportServerError(err)
	}))
	defer server.Close()

	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
	conn, _, err := websocket.DefaultDialer.DialContext(context.Background(), wsURL, nil)
	require.NoError(t, err)
	defer conn.Close()

	session := newSession(conn, sessionKey, nil, "", "task-1")
	client := &Client{
		currentSession: session,
		model:          "test-model",
		reuseSession:   true,
	}

	completion, err := client.Chat(context.Background(), &ChatCompletionRequest{
		Messages: []Message{{Role: "user", Content: "hi"}},
		Stream:   true,
	})
	require.NoError(t, err)
	require.NotNil(t, completion)

	var gotName string
	var gotArgs string
	for chunk := range completion.ToolCallChunks {
		if chunk.Name != "" {
			gotName = chunk.Name
		}
		gotArgs += chunk.Arguments
	}
	for range completion.Chunks {
	}

	require.NoError(t, completion.Error)
	assert.Equal(t, "get_current_weather", gotName)
	assert.Equal(t, `{"location": "Boston, MA"}`, gotArgs)
	assert.Equal(t, "tool_calls", completion.FinishReason)
	assert.Equal(t, 3, completion.Usage.TotalTokens)

	select {
	case err := <-serverErrors:
		require.NoError(t, err)
	default:
	}
}
