gochat/backend/internal/llm/llm_test.go

package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

// --- Interface compliance test ---

func TestOpenAIProvider_ImplementsProvider(t *testing.T) {
	var _ Provider = (*OpenAIProvider)(nil)
}

// --- Constructor tests ---

func TestNewOpenAIProvider_Defaults(t *testing.T) {
	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey: "test-key",
	})

	assert.Equal(t, "test-key", p.apiKey)
	assert.Equal(t, "https://api.openai.com/v1", p.baseURL)
	assert.Equal(t, "gpt-4", p.model)
	assert.Equal(t, "text-embedding-3-small", p.embedModel)
	assert.Equal(t, 3, p.maxRetries)
}

func TestNewOpenAIProvider_CustomConfig(t *testing.T) {
	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:     "custom-key",
		BaseURL:    "https://ark.cn-beijing.volces.com/api/v3",
		Model:      "doubao-pro-4k",
		EmbedModel: "doubao-embedding",
		MaxRetries: 5,
		Timeout:    120,
	})

	assert.Equal(t, "custom-key", p.apiKey)
	assert.Equal(t, "https://ark.cn-beijing.volces.com/api/v3", p.baseURL)
	assert.Equal(t, "doubao-pro-4k", p.model)
	assert.Equal(t, "doubao-embedding", p.embedModel)
	assert.Equal(t, 5, p.maxRetries)
}

func TestNewOpenAIProvider_TrailingSlashTrimmed(t *testing.T) {
	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:  "key",
		BaseURL: "https://api.openai.com/v1/",
	})
	assert.Equal(t, "https://api.openai.com/v1", p.baseURL)
}

// --- Request serialization tests ---

func TestChatRequest_MarshalJSON(t *testing.T) {
	req := ChatRequest{
		Model:       "gpt-4",
		Messages:    []ChatMessage{
			{Role: "system", Content: "You are a helpful assistant."},
			{Role: "user", Content: "Hello!"},
		},
		Temperature: 0.7,
		MaxTokens:   100,
	}

	body, err := json.Marshal(req)
	require.NoError(t, err)

	// Verify key fields are present
	assert.Contains(t, string(body), `"model":"gpt-4"`)
	assert.Contains(t, string(body), `"messages"`)
	assert.Contains(t, string(body), `"role":"system"`)
	assert.Contains(t, string(body), `"role":"user"`)
	assert.Contains(t, string(body), `"temperature":0.7`)
	assert.Contains(t, string(body), `"max_tokens":100`)
}

func TestChatRequest_MarshalJSON_WithTools(t *testing.T) {
	req := ChatRequest{
		Model:    "gpt-4",
		Messages: []ChatMessage{{Role: "user", Content: "Search for docs"}},
		Tools: []ToolDefinition{
			{
				Type: "function",
				Function: ToolFunction{
					Name:        "search_documents",
					Description: "Search knowledge base",
					Parameters: map[string]interface{}{
						"type": "object",
						"properties": map[string]interface{}{
							"query": map[string]interface{}{
								"type": "string",
							},
						},
					},
				},
			},
		},
	}

	body, err := json.Marshal(req)
	require.NoError(t, err)

	assert.Contains(t, string(body), `"tools"`)
	assert.Contains(t, string(body), `"search_documents"`)
}

func TestEmbeddingRequest_MarshalJSON(t *testing.T) {
	req := EmbeddingRequest{
		Model: "text-embedding-3-small",
		Input: []string{"Hello world", "Test embedding"},
	}

	body, err := json.Marshal(req)
	require.NoError(t, err)

	assert.Contains(t, string(body), `"model":"text-embedding-3-small"`)
	assert.Contains(t, string(body), `"Hello world"`)
	assert.Contains(t, string(body), `"Test embedding"`)
}

func TestChatRequest_OmitEmptyFields(t *testing.T) {
	req := ChatRequest{
		Model:    "gpt-4",
		Messages: []ChatMessage{{Role: "user", Content: "Hi"}},
	}

	body, err := json.Marshal(req)
	require.NoError(t, err)

	// temperature=0 should NOT be omitted since 0 is the zero value for float64
	// but "omitempty" will omit it — this is expected behavior
	decoded := ChatRequest{}
	err = json.Unmarshal(body, &decoded)
	require.NoError(t, err)
	assert.Equal(t, "gpt-4", decoded.Model)
	assert.Equal(t, 1, len(decoded.Messages))
}

// --- Response deserialization tests ---

func TestChatResponse_UnmarshalJSON(t *testing.T) {
	raw := `{
		"id": "chatcmpl-123",
		"object": "chat.completion",
		"created": 1677652288,
		"model": "gpt-4",
		"choices": [{
			"index": 0,
			"message": {"role": "assistant", "content": "Hello! How can I help you?"},
			"finish_reason": "stop"
		}],
		"usage": {
			"prompt_tokens": 9,
			"completion_tokens": 12,
			"total_tokens": 21
		}
	}`

	var resp ChatResponse
	err := json.Unmarshal([]byte(raw), &resp)
	require.NoError(t, err)

	assert.Equal(t, "chatcmpl-123", resp.ID)
	assert.Equal(t, 1, len(resp.Choices))
	assert.Equal(t, "assistant", resp.Choices[0].Message.Role)
	assert.Equal(t, "Hello! How can I help you?", resp.Choices[0].Message.Content)
	assert.Equal(t, "stop", resp.Choices[0].FinishReason)
	assert.Equal(t, 21, resp.Usage.TotalTokens)
}

func TestEmbeddingResponse_UnmarshalJSON(t *testing.T) {
	raw := `{
		"object": "list",
		"data": [{
			"object": "embedding",
			"index": 0,
			"embedding": [0.0023064255, -0.009327292, 0.015871]
		}],
		"model": "text-embedding-3-small",
		"usage": {
			"prompt_tokens": 5,
			"total_tokens": 5
		}
	}`

	var resp EmbeddingResponse
	err := json.Unmarshal([]byte(raw), &resp)
	require.NoError(t, err)

	assert.Equal(t, 1, len(resp.Data))
	assert.Equal(t, 0, resp.Data[0].Index)
	assert.Equal(t, 3, len(resp.Data[0].Embedding))
	assert.InDelta(t, 0.0023064255, resp.Data[0].Embedding[0], 0.0001)
}

// --- Mock server tests ---

func TestOpenAIProvider_ChatCompletion_MockServer(t *testing.T) {
	mockResp := ChatResponse{
		ID:     "chatcmpl-test",
		Object: "chat.completion",
		Model:  "gpt-4",
		Choices: []ChatChoice{
			{
				Index:        0,
				Message:      ChatMessage{Role: "assistant", Content: "Mocked response"},
				FinishReason: "stop",
			},
		},
		Usage: TokenUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15},
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "POST", r.Method)
		assert.Equal(t, "/chat/completions", r.URL.Path)
		assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
		assert.Equal(t, "application/json", r.Header.Get("Content-Type"))

		body, _ := json.Marshal(mockResp)
		w.WriteHeader(http.StatusOK)
		w.Write(body)
	}))
	defer server.Close()

	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:     "test-key",
		BaseURL:    server.URL,
		Model:      "gpt-4",
		MaxRetries: 0,
	})

	req := ChatRequest{
		Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
	}

	resp, err := p.ChatCompletion(context.Background(), req)
	require.NoError(t, err)
	assert.Equal(t, "chatcmpl-test", resp.ID)
	assert.Equal(t, "Mocked response", resp.Choices[0].Message.Content)
}

func TestOpenAIProvider_CreateEmbedding_MockServer(t *testing.T) {
	mockResp := EmbeddingResponse{
		Object: "list",
		Data: []EmbeddingData{
			{
				Object:    "embedding",
				Index:     0,
				Embedding: []float64{0.1, 0.2, 0.3},
			},
		},
		Model: "text-embedding-3-small",
		Usage: TokenUsage{PromptTokens: 3, TotalTokens: 3},
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "POST", r.Method)
		assert.Equal(t, "/embeddings", r.URL.Path)
		assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))

		body, _ := json.Marshal(mockResp)
		w.WriteHeader(http.StatusOK)
		w.Write(body)
	}))
	defer server.Close()

	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:     "test-key",
		BaseURL:    server.URL,
		MaxRetries: 0,
	})

	req := EmbeddingRequest{
		Input: []string{"Hello world"},
	}

	resp, err := p.CreateEmbedding(context.Background(), req)
	require.NoError(t, err)
	assert.Equal(t, 1, len(resp.Data))
	assert.InDelta(t, 0.1, resp.Data[0].Embedding[0], 0.001)
}

func TestOpenAIProvider_ChatCompletion_APIError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte(`{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}`))
	}))
	defer server.Close()

	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:     "bad-key",
		BaseURL:    server.URL,
		MaxRetries: 0,
	})

	req := ChatRequest{
		Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
	}

	resp, err := p.ChatCompletion(context.Background(), req)
	assert.Nil(t, resp)
	require.Error(t, err)

	var apiErr *APIError
	require.True(t, errors.As(err, &apiErr), "error should wrap APIError")
	assert.Equal(t, http.StatusUnauthorized, apiErr.StatusCode)
	assert.Equal(t, "Invalid API key", apiErr.Message)
	assert.Equal(t, "invalid_request_error", apiErr.Type)
}

func TestOpenAIProvider_ChatCompletion_RetryOn5xx(t *testing.T) {
	callCount := 0

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		callCount++
		if callCount < 3 {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(`{"error":{"message":"Internal server error"}}`))
			return
		}

		mockResp := ChatResponse{
			ID:      "chatcmpl-retry",
			Choices: []ChatChoice{{Message: ChatMessage{Role: "assistant", Content: "Success after retry"}}},
		}
		body, _ := json.Marshal(mockResp)
		w.WriteHeader(http.StatusOK)
		w.Write(body)
	}))
	defer server.Close()

	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:     "test-key",
		BaseURL:    server.URL,
		MaxRetries: 3,
		Timeout:    5,
	})

	req := ChatRequest{
		Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
	}

	resp, err := p.ChatCompletion(context.Background(), req)
	require.NoError(t, err)
	assert.Equal(t, "Success after retry", resp.Choices[0].Message.Content)
	assert.Equal(t, 3, callCount)
}

func TestOpenAIProvider_ChatCompletion_NoRetryOn4xx(t *testing.T) {
	callCount := 0

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		callCount++
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(`{"error":{"message":"Bad request"}}`))
	}))
	defer server.Close()

	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:     "test-key",
		BaseURL:    server.URL,
		MaxRetries: 3,
		Timeout:    5,
	})

	req := ChatRequest{
		Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
	}

	resp, err := p.ChatCompletion(context.Background(), req)
	assert.Nil(t, resp)
	require.Error(t, err)
	assert.Equal(t, 1, callCount, "should not retry on 400")
}

func TestOpenAIProvider_DefaultModelApplied(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var reqBody ChatRequest
		json.NewDecoder(r.Body).Decode(&reqBody)
		assert.Equal(t, "gpt-4", reqBody.Model, "default model should be applied when not specified")

		mockResp := ChatResponse{ID: "test"}
		body, _ := json.Marshal(mockResp)
		w.WriteHeader(http.StatusOK)
		w.Write(body)
	}))
	defer server.Close()

	p := NewOpenAIProvider(OpenAIProviderConfig{
		APIKey:  "key",
		BaseURL: server.URL,
		Model:   "gpt-4",
	})

	req := ChatRequest{
		Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
		// Model intentionally left empty
	}

	_, err := p.ChatCompletion(context.Background(), req)
	require.NoError(t, err)
}

// --- SSE stream parsing test ---

func TestOpenAIProvider_ParseSSEStream(t *testing.T) {
	sseData := `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":""}]}

data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":""}]}

data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":""}]}

data: [DONE]

`

	p := NewOpenAIProvider(OpenAIProviderConfig{APIKey: "test"})

	var chunks []StreamChunk
	err := p.parseSSEStream(strings.NewReader(sseData), func(chunk StreamChunk) error {
		chunks = append(chunks, chunk)
		return nil
	})
	require.NoError(t, err)
	assert.Equal(t, 3, len(chunks))
	assert.Equal(t, "assistant", chunks[0].Choices[0].Delta.Role)
	assert.Equal(t, "Hello", chunks[1].Choices[0].Delta.Content)
	assert.Equal(t, " world", chunks[2].Choices[0].Delta.Content)
}

// --- Utility tests ---

func TestParseFloatEmbedding(t *testing.T) {
	raw := []interface{}{float64(0.1), float32(0.2), int(3), int64(4), "5.5"}
	result := ParseFloatEmbedding(raw)

	assert.Equal(t, 5, len(result))
	assert.InDelta(t, 0.1, result[0], 0.001)
	assert.InDelta(t, 0.2, result[1], 0.001)
	assert.InDelta(t, 3.0, result[2], 0.001)
	assert.InDelta(t, 4.0, result[3], 0.001)
	assert.InDelta(t, 5.5, result[4], 0.001)
}

// --- APIError tests ---

func TestAPIError_Error(t *testing.T) {
	err := &APIError{
		StatusCode: 401,
		Message:    "Invalid API key",
		Type:       "invalid_request_error",
		Code:       "invalid_api_key",
	}
	assert.Equal(t, "API error (status 401): Invalid API key", err.Error())
}

func TestIsNonRetriableError(t *testing.T) {
	// 4xx (except 429) should not retry
	assert.True(t, isNonRetriableError(&APIError{StatusCode: 400}))
	assert.True(t, isNonRetriableError(&APIError{StatusCode: 401}))
	assert.True(t, isNonRetriableError(&APIError{StatusCode: 403}))
	assert.True(t, isNonRetriableError(&APIError{StatusCode: 404}))

	// 429 (rate limit) should retry
	assert.False(t, isNonRetriableError(&APIError{StatusCode: 429}))

	// 5xx should retry
	assert.False(t, isNonRetriableError(&APIError{StatusCode: 500}))
	assert.False(t, isNonRetriableError(&APIError{StatusCode: 502}))
	assert.False(t, isNonRetriableError(&APIError{StatusCode: 503}))

	// Non-APIError should not be treated as non-retriable
	assert.False(t, isNonRetriableError(fmt.Errorf("some error")))
}