package router

import (
	"encoding/base64"
	"net/http/httptest"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/go-resty/resty/v2"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/inference-service/internal/config"
	"github.com/tidbcloud/inference-service/internal/service"
)

// setupTestServer creates a test server using the real application's route setup
func setupTestServer(t *testing.T) (*httptest.Server, *resty.Client, *service.UsageService) {
	// Set gin to test mode
	gin.SetMode(gin.TestMode)

	// Create temporary directory for usage data
	tempDir := t.TempDir()

	// Load test configuration
	err := config.LoadConfig("")
	require.NoError(t, err)

	// Override storage directory for testing
	cfg := config.Global()
	cfg.Storage.UsageDataDir = tempDir

	// Initialize services exactly like the real application
	bedrockService := service.NewBedrockService(
		cfg.EmbeddingProviders.AWSBedrock.APIToken,
		cfg.EmbeddingProviders.AWSBedrock.Region,
	)

	usageService, err := service.NewUsageService()
	require.NoError(t, err)

	// Setup router with middleware exactly like the real application
	router := NewRouter(bedrockService, usageService)
	httpServer := httptest.NewServer(router)

	// Create resty client
	client := resty.New().SetBaseURL(httpServer.URL)

	return httpServer, client, usageService
}

func TestEmbeddingOptWhitelist(t *testing.T) {
	server, client, usageService := setupTestServer(t)
	defer server.Close()
	defer usageService.Close()

	// Test with Cohere model but try to pass Titan-specific parameters
	// This should work but the disallowed parameters should be filtered out
	requestBody := map[string]any{
		"model":      "cohere/embed-multilingual-v3",
		"texts":      []string{"Hello world"},
		"input_type": "search_document", // allowed for Cohere
		"truncate":   "END",             // allowed for Cohere
		"dimensions": 512,               // NOT allowed for Cohere (Titan only)
		"normalize":  true,              // NOT allowed for Cohere (Titan only)
	}

	var response map[string]any
	resp, err := client.R().
		SetBody(requestBody).
		SetResult(&response).
		SetError(&response).
		Post("/api/v1/inference/embeddings/test-tenant")

	require.NoError(t, err)
	require.Equal(t, 400, resp.StatusCode())
	errorMsg, ok := response["error"].(string)
	require.True(t, ok)
	require.Contains(t, errorMsg, `is not allowed for model`)
}

func TestInvalidModel(t *testing.T) {
	server, client, usageService := setupTestServer(t)
	defer server.Close()
	defer usageService.Close()

	requestBody := map[string]interface{}{
		"model": "invalid-model",
		"texts": []string{"Hello world"},
	}

	var response service.ErrorResponse
	resp, err := client.R().
		SetBody(requestBody).
		SetResult(&response).
		SetError(&response).
		Post("/api/v1/inference/embeddings/test-tenant")

	require.NoError(t, err)
	require.Equal(t, 400, resp.StatusCode())
	require.Equal(t, `Unknown model 'invalid-model'`, response.Error)
}

func TestEmptyTexts(t *testing.T) {
	server, client, usageService := setupTestServer(t)
	defer server.Close()
	defer usageService.Close()

	requestBody := map[string]interface{}{
		"model": "amazon/titan-embed-text-v1",
		"texts": []string{},
	}

	var response service.ErrorResponse
	resp, err := client.R().
		SetBody(requestBody).
		SetResult(&response).
		SetError(&response).
		Post("/api/v1/inference/embeddings/test-tenant")

	require.NoError(t, err)
	require.Equal(t, 400, resp.StatusCode())
	require.Equal(t, `Texts array cannot be empty`, response.Error)
}

func TestTenantQuotaUsage(t *testing.T) {
	server, client, usageService := setupTestServer(t)
	defer server.Close()
	defer usageService.Close()

	var response service.QuotaUsageResponse
	resp, err := client.R().
		SetResult(&response).
		Get("/api/v1/inference/quota_usage/test-tenant")

	require.NoError(t, err)
	require.Equal(t, 200, resp.StatusCode())

	require.Equal(t, "test-tenant", response.TenantID)
	require.GreaterOrEqual(t, response.CurrentUsage, 0.0)
	require.Greater(t, response.QuotaLimit, 0.0)
	require.False(t, response.ResetDate.IsZero())
}

func TestAllQuotaUsage(t *testing.T) {
	server, client, usageService := setupTestServer(t)
	defer server.Close()
	defer usageService.Close()

	var response service.AllQuotaUsageResponse
	resp, err := client.R().
		SetResult(&response).
		Get("/api/v1/inference/quota_usage")

	require.NoError(t, err)
	require.Equal(t, 200, resp.StatusCode())
	require.NotNil(t, response.Tenants)
	require.GreaterOrEqual(t, len(response.Tenants), 0)
}

func TestMissingTenantID(t *testing.T) {
	server, client, usageService := setupTestServer(t)
	defer server.Close()
	defer usageService.Close()

	requestBody := map[string]interface{}{
		"model": "amazon/titan-embed-text-v1",
		"texts": []string{"Hello world"},
	}

	resp, err := client.R().
		SetBody(requestBody).
		Post("/api/v1/inference/embeddings/")

	require.NoError(t, err)
	require.Equal(t, 404, resp.StatusCode())
}

func TestInvalidJSON(t *testing.T) {
	server, client, usageService := setupTestServer(t)
	defer server.Close()
	defer usageService.Close()

	var response service.ErrorResponse
	resp, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetBody("invalid json").
		SetResult(&response).
		SetError(&response).
		Post("/api/v1/inference/embeddings/test-tenant")

	require.NoError(t, err)
	require.Equal(t, 400, resp.StatusCode())
	require.Contains(t, response.Error, "Invalid request body")
}

// TestServiceAPIKeyAuthentication tests that service API key authentication is enforced
func TestServiceAPIKeyAuthentication(t *testing.T) {
	// Test cases for service API key authentication
	testCases := []struct {
		name               string
		configServiceKey   string
		authHeader         string
		expectedStatusCode int
		expectSuccess      bool
	}{
		{
			name:               "NoAPIKeyConfigured_ShouldAllowAccess",
			configServiceKey:   "",
			authHeader:         "",
			expectedStatusCode: 200,
			expectSuccess:      true,
		},
		{
			name:               "APIKeyConfigured_NoAuthHeader_ShouldDeny",
			configServiceKey:   "test-service-key-123",
			authHeader:         "",
			expectedStatusCode: 401,
			expectSuccess:      false,
		},
		{
			name:               "APIKeyConfigured_InvalidFormat_ShouldDeny",
			configServiceKey:   "test-service-key-123",
			authHeader:         "Basic dGVzdA==",
			expectedStatusCode: 401,
			expectSuccess:      false,
		},
		{
			name:               "APIKeyConfigured_WrongKey_ShouldDeny",
			configServiceKey:   "test-service-key-123",
			authHeader:         "Bearer wrong-key",
			expectedStatusCode: 401,
			expectSuccess:      false,
		},
		{
			name:               "APIKeyConfigured_CorrectKey_ShouldAllow",
			configServiceKey:   "test-service-key-123",
			authHeader:         "Bearer test-service-key-123",
			expectedStatusCode: 200,
			expectSuccess:      true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			server, client, usageService := setupTestServer(t)
			defer server.Close()
			defer usageService.Close()

			// Override the service API key configuration
			cfg := config.Global()
			originalServiceKey := cfg.Server.CheckServiceAPIKey
			cfg.Server.CheckServiceAPIKey = tc.configServiceKey
			defer func() {
				cfg.Server.CheckServiceAPIKey = originalServiceKey
			}()

			// Prepare request body (valid request to focus on auth testing)
			requestBody := map[string]interface{}{
				"model": "amazon/titan-embed-text-v2",
				"texts": []string{"Test authentication"},
			}

			// Create request with or without auth header
			req := client.R().SetBody(requestBody)
			if tc.authHeader != "" {
				req.SetHeader("Authorization", tc.authHeader)
			}

			var response map[string]interface{}
			resp, err := req.
				SetResult(&response).
				SetError(&response).
				Post("/api/v1/inference/embeddings/test-tenant")

			require.NoError(t, err)
			require.Equal(t, tc.expectedStatusCode, resp.StatusCode())

			if tc.expectSuccess {
				// For successful requests, verify we get a proper embedding response
				data, ok := response["embeddings"].([]any)
				require.True(t, ok)
				require.Len(t, data, 1)
			} else {
				// For failed requests, verify we get proper error response
				errorObj, ok := response["error"].(string)
				require.True(t, ok)
				require.True(t, len(errorObj) > 0)
			}

			t.Logf("✅ Authentication test case %s passed", tc.name)
		})
	}
}

// TestModelsWorking tests these models are working as expected (smoke tests)
func TestModelsWorking(t *testing.T) {
	// Test cases for each model with appropriate parameters
	testCases := []struct {
		name             string
		model            string
		texts            []string
		otherOpts        map[string]interface{}
		expectDimensions int
		expectError      string
	}{
		{
			name:             "TitanV2_Basic",
			model:            "amazon/titan-embed-text-v2",
			texts:            []string{"Titan v2 test"},
			otherOpts:        map[string]interface{}{},
			expectDimensions: 1024,
		},
		{
			name:  "TitanV2_WithDimensions",
			model: "amazon/titan-embed-text-v2",
			texts: []string{"Advanced embedding test"},
			otherOpts: map[string]interface{}{
				"dimensions": 256,
			},
			expectDimensions: 256,
		},
		{
			name:  "TitanV2_Multiple",
			model: "amazon/titan-embed-text-v2",
			texts: []string{"Advanced embedding test", "hello world", "abc"},
			otherOpts: map[string]interface{}{
				"dimensions": 256,
			},
			expectDimensions: 256,
		},
		{
			name:        "CohereEnglish_NoArgs",
			model:       "cohere/embed-english-v3",
			texts:       []string{"English text embedding"},
			otherOpts:   map[string]interface{}{},
			expectError: "required key [input_type] not found",
		},
		{
			name:  "CohereEnglish_Basic",
			model: "cohere/embed-english-v3",
			texts: []string{"English text embedding"},
			otherOpts: map[string]interface{}{
				"input_type": "search_query",
			},
			expectDimensions: 1024,
		},
		{
			name:  "CohereEnglish_WithTruncate",
			model: "cohere/embed-english-v3",
			texts: []string{"This is a longer text that might need truncation for testing purposes"},
			otherOpts: map[string]interface{}{
				"input_type": "search_query",
				"truncate":   "END",
			},
			expectDimensions: 1024,
		},
		{
			name:  "CohereMultilingual_Basic",
			model: "cohere/embed-multilingual-v3",
			texts: []string{"Hello world", "Hola mundo", "Bonjour le monde"},
			otherOpts: map[string]interface{}{
				"input_type": "search_document",
			},
			expectDimensions: 1024,
		},
		{
			name:  "CohereMultilingual_Classification",
			model: "cohere/embed-multilingual-v3",
			texts: []string{"Classification task example"},
			otherOpts: map[string]interface{}{
				"input_type": "classification",
				"truncate":   "START",
			},
			expectDimensions: 1024,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			server, client, usageService := setupTestServer(t)
			defer server.Close()
			defer usageService.Close()

			requestBody := map[string]interface{}{
				"model": tc.model,
				"texts": tc.texts,
			}
			for key, value := range tc.otherOpts {
				requestBody[key] = value
			}

			var response map[string]interface{}
			resp, err := client.R().
				SetBody(requestBody).
				SetResult(&response).
				SetError(&response).
				Post("/api/v1/inference/embeddings/test-tenant")

			require.NoError(t, err)

			// Check if the request succeeded or failed gracefully
			if tc.expectError != "" {
				if resp.StatusCode() == 200 {
					t.Errorf("Expected error for model %s but got none", tc.model)
					return
				}

				errorMsg, ok := response["error"].(string)
				require.True(t, ok)
				require.Contains(t, errorMsg, tc.expectError)
				return
			} else {
				require.Equal(t, 200, resp.StatusCode(), "Response body: %s", resp.String())
			}

			data, ok := response["embeddings"].([]any)
			require.True(t, ok)
			require.Len(t, data, len(tc.texts))

			// Validate each embedding
			for _, item := range data {
				embeddingBase64, ok := item.(string)
				require.True(t, ok)

				decoded, err := base64.StdEncoding.DecodeString(embeddingBase64)
				require.NoError(t, err)

				require.Equal(t, tc.expectDimensions*4, len(decoded))
			}

			t.Logf("✅ Model %s test passed", tc.model)
		})
	}
}