package main

import (
	"encoding/json"
	"errors"
	"strings"
	"testing"
)

func TestUpstreamErrorMessage(t *testing.T) {
	for _, tt := range []struct{ name, body, key, want string }{
		{"routing", `{"error":{"message":"No endpoints found matching your data policy"}}`, "", "Upstream model service returned HTTP 404. No endpoints found matching your data policy"},
		{"proxy", `{"message":"Route not found"}`, "", "Upstream model service returned HTTP 404. Route not found"},
		{"redaction", `{"error":{"message":"opaque-private-key Bearer abc123 sk-or-v1-example https://user:pass@example.com/?key=secret"},"metadata":{"secret":"must not appear"}}`, "opaque-private-key", "Upstream model service returned HTTP 404. [REDACTED] [REDACTED] [REDACTED] [URL omitted]"},
		{"html", `<html>private upstream page</html>`, "", "Upstream model service returned HTTP 404. The response was not a structured JSON error."},
		{"empty", `{}`, "", "Upstream model service returned HTTP 404. No error message was returned."},
		{"oversized", strings.Repeat("x", 16*1024+1), "", "Upstream model service returned HTTP 404. No readable error details were returned."},
	} {
		t.Run(tt.name, func(t *testing.T) {
			if got := upstreamErrorMessage(strings.NewReader(tt.body), 404, tt.key); got != tt.want {
				t.Fatalf("got %q, want %q", got, tt.want)
			}
		})
	}
	message := upstreamErrorMessage(strings.NewReader(`{"error":{"message":"`+strings.Repeat("é", 2000)+`"}}`), 404, "")
	if len([]rune(message)) > 1100 {
		t.Fatal("upstream message was not bounded")
	}
}

func TestSanitizeProviderPayload(t *testing.T) {
	payload := map[string]any{
		"model":                 homarrProviderModelID,
		"models":                []string{"attacker/model"},
		"provider":              map[string]any{"order": []string{"attacker"}},
		"route":                 "fallback",
		"plugins":               []string{"web"},
		"transforms":            []string{"middle-out"},
		"max_tokens":            1_000_000,
		"max_completion_tokens": 1_000_000,
		"n":                     100,
		"parallel_tool_calls":   true,
		"extra_body":            map[string]any{"provider": map[string]any{"order": []string{"attacker"}}},
		"extra_headers":         map[string]any{"Authorization": "Bearer attacker"},
		"metadata":              map[string]any{"user": "private-user-id"},
		"audio":                 map[string]any{"format": "wav"},
		"modalities":            []string{"text", "audio"},
		"logprobs":              true,
		"top_logprobs":          20,
		"prediction":            map[string]any{"content": strings.Repeat("x", 100)},
		"service_tier":          "priority",
		"reasoning":             map[string]any{"effort": "high", "max_tokens": 1_000_000},
		"user":                  "private-user-id",
		"tools": []any{
			map[string]any{"type": "function", "function": map[string]any{"name": "board_list"}},
			map[string]any{
				"type":       "openrouter:web_search",
				"parameters": map[string]any{"max_results": 1000, "max_uses": 1000},
			},
			map[string]any{"type": "openrouter:web_search"},
		},
		"messages": []any{map[string]any{"role": "user", "content": "hello"}},
	}
	if err := sanitizeProviderPayload(payload, "mock/team-selected-model"); err != nil {
		t.Fatal(err)
	}
	for _, field := range clientControlledRoutingFields {
		if field == "provider" {
			continue
		}
		if _, exists := payload[field]; exists {
			t.Fatalf("client-controlled routing field %q was forwarded", field)
		}
	}
	for _, field := range clientControlledCostFields {
		if _, exists := payload[field]; exists {
			t.Fatalf("client-controlled cost field %q was forwarded", field)
		}
	}
	if payload["model"] != "mock/team-selected-model" {
		t.Fatalf("unexpected upstream model: %v", payload["model"])
	}
	if payload["max_completion_tokens"] != maxChatOutputTokens || payload["n"] != 1 || payload["parallel_tool_calls"] != false {
		t.Fatalf("unsafe generation controls were forwarded: %#v", payload)
	}
	if _, exists := payload["max_tokens"]; exists {
		t.Fatal("deprecated client max_tokens was forwarded")
	}
	if _, exists := payload["user"]; exists {
		t.Fatal("client user identifier was forwarded")
	}
	if _, exists := payload["metadata"]; exists {
		t.Fatal("client metadata was forwarded")
	}
	reasoning := payload["reasoning"].(map[string]any)
	if _, exists := reasoning["max_tokens"]; exists || reasoning["effort"] != "high" {
		t.Fatalf("reasoning budget was not bounded safely: %#v", reasoning)
	}
	tools := payload["tools"].([]any)
	if len(tools) != 2 {
		t.Fatalf("unexpected sanitized tool count: %d", len(tools))
	}
	webSearch := tools[1].(map[string]any)
	parameters := webSearch["parameters"].(map[string]any)
	if parameters["max_results"] != maxWebSearchResults ||
		parameters["max_uses"] != maxWebSearchUses ||
		parameters["max_total_results"] != maxWebSearchTotal ||
		parameters["max_characters"] != maxWebSearchCharacters ||
		parameters["search_context_size"] != webSearchContextSize {
		t.Fatalf("web search limits were not enforced: %#v", parameters)
	}
	usage, err := json.Marshal(payload["usage"])
	if err != nil || string(usage) != `{"include":true}` {
		t.Fatalf("unexpected usage options: %s, %v", usage, err)
	}
	privacy := payload["provider"].(map[string]any)
	for _, field := range []string{"zdr", "data_collection"} {
		if _, exists := privacy[field]; exists {
			t.Fatalf("privacy policy must be left to OpenRouter guardrails: %#v", privacy)
		}
	}
}

func TestSanitizeProviderPayloadPinsDefaultModelQuality(t *testing.T) {
	payload := map[string]any{
		"reasoning_effort":  "none",
		"reasoning":         map[string]any{"effort": "low", "enabled": false, "exclude": true},
		"include_reasoning": false,
		"temperature":       0.3,
		"top_p":             0.5,
		"provider":          map[string]any{"only": []string{"azure"}, "allow_fallbacks": true, "zdr": true, "data_collection": "allow"},
		"stream":            true,
		"stream_options":    map[string]any{"include_usage": false},
	}
	if err := sanitizeProviderPayload(payload, defaultOpenRouterModelID); err != nil {
		t.Fatal(err)
	}

	if payload["model"] != "openai/gpt-6-luna" {
		t.Fatalf("unexpected default model: %#v", payload["model"])
	}
	reasoning := payload["reasoning"].(map[string]any)
	if reasoning["effort"] != "max" || reasoning["exclude"] != false {
		t.Fatalf("default model did not use visible max reasoning: %#v", reasoning)
	}
	for _, field := range []string{"reasoning_effort", "include_reasoning", "temperature", "top_p"} {
		if _, exists := payload[field]; exists {
			t.Fatalf("conflicting or unsupported field %q was forwarded", field)
		}
	}
	if payload["stream_options"].(map[string]any)["include_usage"] != true {
		t.Fatal("streamed token usage must be included")
	}
	preferences := payload["provider"].(map[string]any)
	if _, exists := preferences["only"]; exists {
		t.Fatalf("provider allowlists must not restrict eligible endpoints: %#v", preferences)
	}
	if preferences["allow_fallbacks"] != false {
		t.Fatalf("default model must not fall back to another provider: %#v", preferences)
	}
	for _, field := range []string{"zdr", "data_collection"} {
		if _, exists := preferences[field]; exists {
			t.Fatalf("client privacy overrides must not be forwarded: %#v", preferences)
		}
	}
	for _, field := range []string{"order", "quantizations"} {
		if _, exists := preferences[field]; exists {
			t.Fatalf("obsolete DeepInfra routing field %q was forwarded", field)
		}
	}
}

func TestSanitizeProviderPayloadPinsLunaReasoning(t *testing.T) {
	payload := map[string]any{
		"reasoning":        map[string]any{"effort": "xhigh"},
		"reasoning_effort": "xhigh",
	}
	if err := sanitizeProviderPayload(payload, lunaOpenRouterModelID); err != nil {
		t.Fatal(err)
	}

	if _, exists := payload["reasoning"]; exists {
		t.Fatal("Luna must use the top-level reasoning_effort field")
	}
	if payload["reasoning_effort"] != "high" {
		t.Fatalf("Luna did not use high reasoning: %#v", payload["reasoning_effort"])
	}
	preferences := payload["provider"].(map[string]any)
	if _, exists := preferences["order"]; exists {
		t.Fatalf("Luna must not inherit DeepInfra routing: %#v", preferences)
	}
}

func TestSanitizeProviderPayloadRejectsUnsupportedServerTools(t *testing.T) {
	payload := map[string]any{"tools": []any{map[string]any{"type": "openrouter:computer"}}}
	if err := sanitizeProviderPayload(payload, "mock/team-selected-model"); err == nil {
		t.Fatal("expected unsupported server tool to be rejected")
	}
}

func TestSafeUpstreamStatus(t *testing.T) {
	for _, status := range []int{401, 403} {
		if actual := safeUpstreamStatus(status); actual != 502 {
			t.Fatalf("expected upstream %d to become 502, got %d", status, actual)
		}
	}
	if actual := safeUpstreamStatus(429); actual != 429 {
		t.Fatalf("expected upstream 429 to be preserved, got %d", actual)
	}
}

func TestValidateProviderInput(t *testing.T) {
	payload := map[string]any{
		"messages": []any{map[string]any{
			"role": "user",
			"content": []any{map[string]any{
				"type":      "image_url",
				"image_url": map[string]any{"url": "data:image/png;base64," + strings.Repeat("a", 1_300_000)},
			}},
		}},
	}
	if err := validateProviderInput(payload); err != nil {
		t.Fatalf("expected a supported Homarr image request, got %v", err)
	}
	payload["messages"] = []any{map[string]any{"role": "user", "content": strings.Repeat("x", maxChatTextBytes-100)}}
	if err := validateProviderInput(payload); err != nil {
		t.Fatalf("expected text at the conservative context boundary to be accepted, got %v", err)
	}
	payload["messages"] = []any{map[string]any{"role": "user", "content": strings.Repeat("x", maxChatTextBytes+1)}}
	if err := validateProviderInput(payload); !errors.Is(err, errInputTooLarge) {
		t.Fatalf("expected oversized text input to be rejected, got %v", err)
	}
	payload["messages"] = []any{map[string]any{
		"role": "user",
		"content": []any{map[string]any{
			"type":      "image_url",
			"image_url": map[string]any{"url": "data:image/png;base64," + strings.Repeat("a", maxChatImageDataBytes)},
		}},
	}}
	if err := validateProviderInput(payload); !errors.Is(err, errTooManyImages) {
		t.Fatalf("expected oversized image input to be rejected, got %v", err)
	}
	payload["messages"] = []any{map[string]any{
		"role": "user",
		"content": []any{map[string]any{
			"type":      "image_url",
			"image_url": map[string]any{"url": "https://attacker.example/expensive-image.png"},
		}},
	}}
	if err := validateProviderInput(payload); !errors.Is(err, errRemoteImage) {
		t.Fatalf("expected a remote image URL to be rejected, got %v", err)
	}
	payload["messages"] = []any{map[string]any{
		"role": "user",
		"content": []any{
			map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
			map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
			map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
			map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
			map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
		},
	}}
	if err := validateProviderInput(payload); err != nil {
		t.Fatalf("expected five uploaded images to be accepted, got %v", err)
	}
	payload["messages"].([]any)[0].(map[string]any)["content"] = append(
		payload["messages"].([]any)[0].(map[string]any)["content"].([]any),
		map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
	)
	if err := validateProviderInput(payload); !errors.Is(err, errTooManyImages) {
		t.Fatalf("expected six uploaded images to be rejected, got %v", err)
	}
}

func TestProviderInFlightLimits(t *testing.T) {
	provider := &homarrProvider{inFlightByUser: make(map[string]int)}
	if !provider.acquireRequest("one") {
		t.Fatal("expected the first user request to be admitted")
	}
	if !provider.acquireRequest("one") {
		t.Fatal("expected the second user request to be admitted")
	}
	if provider.acquireRequest("one") {
		t.Fatal("expected the third user request to be rejected")
	}
	provider.releaseRequest("one")
	provider.releaseRequest("one")
	for index := range maxGlobalInFlight {
		if !provider.acquireRequest(string(rune(index + 1))) {
			t.Fatalf("expected global request %d to be admitted", index)
		}
	}
	if provider.acquireRequest("overflow") {
		t.Fatal("expected the global in-flight limit to reject overflow")
	}
}

func TestCompletionModelAlias(t *testing.T) {
	body, err := aliasCompletionBody([]byte(`{"id":"one","model":"private/model","choices":[]}`))
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(body), "private/model") || !strings.Contains(string(body), `"model":"homarr/model"`) {
		t.Fatalf("private upstream model leaked in JSON: %s", body)
	}
	var stream strings.Builder
	_, err = copyAliasedSSE(
		&stream,
		strings.NewReader("data: {\"id\":\"one\",\"model\":\"private/model\",\"choices\":[]}\n\ndata: [DONE]\n\n"),
		1024,
	)
	if err != nil || strings.Contains(stream.String(), "private/model") ||
		!strings.Contains(stream.String(), `"model":"homarr/model"`) {
		t.Fatalf("private upstream model leaked in SSE: %s, %v", stream.String(), err)
	}
}

func TestProviderEnvironment(t *testing.T) {
	t.Setenv("OPENROUTER_API_KEY", "test-key")
	t.Setenv("HOMARR_AI_OPENROUTER_BASE_URL", "https://router.example/v1/")
	t.Setenv("HOMARR_AI_OPENROUTER_MODEL", "mock/team-selected-model")
	t.Setenv("HOMARR_AI_DAILY_REQUEST_LIMIT", "75")
	t.Setenv("HOMARR_AI_GLOBAL_DAILY_REQUEST_LIMIT", "900")
	provider, err := newHomarrProviderFromEnv()
	if err != nil {
		t.Fatal(err)
	}
	if provider.apiKey != "test-key" || provider.baseURL != "https://router.example/v1" ||
		provider.modelID != "mock/team-selected-model" || provider.dailyLimit != 75 || provider.globalDailyLimit != 900 {
		t.Fatalf("unexpected provider config: %#v", provider)
	}

	t.Setenv("HOMARR_AI_DAILY_REQUEST_LIMIT", "0")
	if _, err := newHomarrProviderFromEnv(); err == nil {
		t.Fatal("expected invalid daily limit to fail")
	}

	t.Setenv("HOMARR_AI_DAILY_REQUEST_LIMIT", "50")
	for _, unsafeURL := range []string{
		"https://user:password@router.example/v1",
		"https://router.example/v1?key=secret",
		"https://router.example/v1#fragment",
	} {
		t.Setenv("HOMARR_AI_OPENROUTER_BASE_URL", unsafeURL)
		if _, err := newHomarrProviderFromEnv(); err == nil {
			t.Fatalf("expected unsafe upstream URL %q to fail", unsafeURL)
		}
	}
	t.Setenv("HOMARR_AI_OPENROUTER_BASE_URL", "http://router.example/v1")
	if _, err := newHomarrProviderFromEnv(); err == nil {
		t.Fatal("expected an insecure upstream to fail without an explicit development opt-in")
	}
}

func TestReadBoundedBody(t *testing.T) {
	body, err := readBoundedBody(strings.NewReader("1234"), 4)
	if err != nil || string(body) != "1234" {
		t.Fatalf("unexpected bounded body: %q, %v", body, err)
	}
	if _, err := readBoundedBody(strings.NewReader("12345"), 4); err != errProviderResponseTooLarge {
		t.Fatalf("expected response size error, got %v", err)
	}
}

func TestCopyBounded(t *testing.T) {
	var target strings.Builder
	written, err := copyBounded(&target, strings.NewReader("1234"), 4)
	if err != nil || written != 4 || target.String() != "1234" {
		t.Fatalf("unexpected bounded copy: %q, %d, %v", target.String(), written, err)
	}

	target.Reset()
	written, err = copyBounded(&target, strings.NewReader("12345"), 4)
	if err != errProviderResponseTooLarge || written != 4 || target.String() != "1234" {
		t.Fatalf("expected bounded copy failure, got %q, %d, %v", target.String(), written, err)
	}
}