package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/apr/concurrency"
	"github.com/tidbcloud/AgentTeam/internal/apr/agentteambridge"
	"github.com/tidbcloud/AgentTeam/internal/apr/runtime"
	"github.com/tidbcloud/AgentTeam/internal/store/organization"
	"github.com/tidbcloud/AgentTeam/internal/workspace/owner"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

type fakeProxyKeyAuthenticator struct{}

func (fakeProxyKeyAuthenticator) AuthenticateOrganizationLLMProxyKey(_ context.Context, token string) (organization.OrganizationProviderProxyAccess, error) {
	return organization.OrganizationProviderProxyAccess{
		KeyID:                  "llmk_1",
		OrganizationID:         "org_1",
		OrganizationProviderID: token,
		Source:                 "persistent_organization",
	}, nil
}

type fakeWorkspaceProxyKeyAuthenticator struct{}

func (fakeWorkspaceProxyKeyAuthenticator) AuthenticateOrganizationLLMProxyKey(_ context.Context, token string) (organization.OrganizationProviderProxyAccess, error) {
	return organization.OrganizationProviderProxyAccess{
		WorkspaceID:        "ws_1",
		AllowedProviderIDs: []string{"opvd_1"},
		AllowedModels:      []string{"zhipu/glm-5.1"},
		Source:             "signed_workspace",
	}, nil
}

type fakeUsageTracker struct {
	mu          sync.Mutex
	decision    *organization.OrganizationLLMQuotaDecision
	decisionErr error
	checks      []struct {
		Access organization.OrganizationProviderProxyAccess
		Model  string
	}
	records []organization.OrganizationLLMUsageRecord
}

func (tracker *fakeUsageTracker) CheckOrganizationLLMQuota(_ context.Context, access organization.OrganizationProviderProxyAccess, model string) (organization.OrganizationLLMQuotaDecision, error) {
	tracker.mu.Lock()
	defer tracker.mu.Unlock()
	tracker.checks = append(tracker.checks, struct {
		Access organization.OrganizationProviderProxyAccess
		Model  string
	}{Access: access, Model: model})
	if tracker.decisionErr != nil {
		return organization.OrganizationLLMQuotaDecision{}, tracker.decisionErr
	}
	if tracker.decision != nil {
		return *tracker.decision, nil
	}
	return organization.OrganizationLLMQuotaDecision{Allowed: true}, nil
}

func (tracker *fakeUsageTracker) RecordOrganizationLLMUsage(_ context.Context, record organization.OrganizationLLMUsageRecord) error {
	tracker.mu.Lock()
	defer tracker.mu.Unlock()
	tracker.records = append(tracker.records, record)
	return nil
}

func TestProxyAuthorizationRejectsOwnerMismatch(t *testing.T) {
	gin.SetMode(gin.ReleaseMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{
		APRDispatcher:         runtime.NewDispatcher(),
		ProxyKeyAuthenticator: fakeWorkspaceProxyKeyAuthenticator{},
		ManagedWorkspaces:     owner.NewStaticSet([]string{"wsp_other"}),
	})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/ws_1/llm/v1/models", nil)
	req.Header.Set("Authorization", "Bearer token")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.Contains(t, rec.Body.String(), "workspace_owner_mismatch")
}

func newDispatcher(upstreamURL string, httpClient *http.Client, providers []struct {
	name      string
	wireAPIs  []string
	models    []string
	token     string
	chatPath  string
	msgPath   string
	respPath  string
	embedPath string
}, routes []runtime.ModelRoute) *runtime.Dispatcher {
	dispatcher := runtime.NewDispatcher()
	for _, item := range providers {
		modelCache := agentteambridge.NewModelCache(agentteambridge.NewStaticModelProvider(item.models))
		_ = modelCache.Initialize()
		dispatcher.AddProvider(agentteambridge.AdaptProviderForTests(agentteambridge.NewForwardProvider(agentteambridge.ForwardProviderConfig{
			Name:               item.name,
			BaseURL:            upstreamURL,
			WireAPIs:           item.wireAPIs,
			TokenPool:          agentteambridge.NewTokenPool([]agentteambridge.TokenState{{Token: item.token, Enabled: true}}, &agentteambridge.RoundRobinSelector{}),
			ModelCache:         modelCache,
			HTTPClient:         httpClient,
			ChatEndpoint:       endpointOrNil(upstreamURL, item.chatPath),
			AnthropicEndpoint:  endpointOrNil(upstreamURL, item.msgPath),
			ResponsesEndpoint:  endpointOrNil(upstreamURL, item.respPath),
			EmbeddingsEndpoint: endpointOrNil(upstreamURL, item.embedPath),
		})))
	}
	dispatcher.SetRoutes(routes)
	dispatcher.Sync()
	return dispatcher
}

func endpointOrNil(baseURL string, path string) *agentteambridge.WireEndpoint {
	if path == "" {
		return nil
	}
	return &agentteambridge.WireEndpoint{BaseURL: baseURL, Path: path}
}

type openAIChatCompletionChunk struct {
	Choices []struct {
		Index int `json:"index"`
		Delta struct {
			ToolCalls []struct {
				Index    int    `json:"index"`
				ID       string `json:"id"`
				Type     string `json:"type"`
				Function struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `json:"function"`
			} `json:"tool_calls"`
		} `json:"delta"`
	} `json:"choices"`
}

func collectOpenAIChunkDataLines(t *testing.T, body string) []openAIChatCompletionChunk {
	t.Helper()
	lines := strings.Split(body, "\n")
	chunks := make([]openAIChatCompletionChunk, 0, len(lines))
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || !strings.HasPrefix(trimmed, "data: ") || trimmed == "data: [DONE]" {
			continue
		}
		var chunk openAIChatCompletionChunk
		require.NoError(t, json.Unmarshal([]byte(strings.TrimPrefix(trimmed, "data: ")), &chunk))
		chunks = append(chunks, chunk)
	}
	return chunks
}

func TestProxyEndpointsRouteThroughRegisteredProvider(t *testing.T) {
	tests := []struct {
		name          string
		path          string
		requestModel  string
		providerName  string
		providerModel string
		upstreamPath  string
	}{
		{name: "chat completions", path: "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", requestModel: "anthropic/claude-sonnet-4.6", providerName: "copilot", providerModel: "claude-sonnet-4.6", upstreamPath: "/chat/completions"},
		{name: "messages", path: "/api/v1/workspaces/wsp_1/llm/v1/messages", requestModel: "anthropic/claude-sonnet-4.6", providerName: "copilot", providerModel: "claude-sonnet-4.6", upstreamPath: "/v1/messages"},
		{name: "responses", path: "/api/v1/workspaces/wsp_1/llm/v1/responses", requestModel: "openai/gpt-4o-mini", providerName: "openai-main", providerModel: "gpt-4o-mini", upstreamPath: "/responses"},
		{name: "embeddings", path: "/api/v1/workspaces/wsp_1/llm/v1/embeddings", requestModel: "openai/text-embedding-3-small", providerName: "openai-main", providerModel: "text-embedding-3-small", upstreamPath: "/embeddings"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var received struct {
				Path  string
				Model string
				Auth  string
			}
			upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
				body, _ := io.ReadAll(request.Body)
				var payload map[string]any
				_ = json.Unmarshal(body, &payload)
				received.Path = request.URL.Path
				received.Model, _ = payload["model"].(string)
				received.Auth = request.Header.Get("Authorization")
				writer.Header().Set("Content-Type", "application/json")
				_, _ = writer.Write([]byte(`{"ok":true}`))
			}))
			defer upstream.Close()

			dispatcher := newDispatcher(upstream.URL, upstream.Client(),
				[]struct {
					name      string
					wireAPIs  []string
					models    []string
					token     string
					chatPath  string
					msgPath   string
					respPath  string
					embedPath string
				}{
					{name: "copilot", wireAPIs: []string{agentteambridge.WireAIOpenAI, agentteambridge.WireAPIAnthropic}, models: []string{"claude-sonnet-4.6"}, token: "copilot-token"},
					{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI, agentteambridge.WireAPIResponses}, models: []string{"gpt-4o-mini", "text-embedding-3-small"}, token: "openai-token"},
				},
				[]runtime.ModelRoute{
					{Vendor: "anthropic", Groups: []runtime.RouteGroup{{Models: []string{"claude-sonnet-4.6"}, Providers: []string{"copilot"}}}},
					{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini", "text-embedding-3-small"}, Providers: []string{"openai-main"}}}},
				},
			)

			gin.SetMode(gin.TestMode)
			router := gin.New()
			RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

			body := bytes.NewBufferString(`{"model":"` + tt.requestModel + `","input":"hi","messages":[{"role":"user","content":"hi"}]}`)
			request := httptest.NewRequest(http.MethodPost, tt.path, body)
			request.Header.Set("Content-Type", "application/json")
			request.Header.Set("Authorization", "Bearer "+tt.providerName)
			recorder := httptest.NewRecorder()
			router.ServeHTTP(recorder, request)

			if recorder.Code != http.StatusOK {
				t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
			}
			if received.Path != tt.upstreamPath {
				t.Fatalf("upstream path = %q want %q", received.Path, tt.upstreamPath)
			}
			if received.Model != tt.providerModel {
				t.Fatalf("rewritten model = %q want %q", received.Model, tt.providerModel)
			}
		})
	}
}

func TestMessagesFallbackToChatCompletions(t *testing.T) {
	var receivedPath string
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		receivedPath = request.URL.Path
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"id":"chatcmpl_1","choices":[{"message":{"role":"assistant","content":"hello from chat"}}],"usage":{"prompt_tokens":11,"completion_tokens":7}}`))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "copilot", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"claude-sonnet-4.6"}, token: "copilot-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "anthropic", Groups: []runtime.RouteGroup{{Models: []string{"claude-sonnet-4.6"}, Providers: []string{"copilot"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/messages", bytes.NewBufferString(`{"model":"anthropic/claude-sonnet-4.6","max_tokens":64,"messages":[{"role":"user","content":"hello"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer copilot")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	if receivedPath != "/chat/completions" {
		t.Fatalf("upstream path = %q want %q", receivedPath, "/chat/completions")
	}
}

func TestProxyEndpointsAcceptProviderAPIKeyHeaders(t *testing.T) {
	tests := []struct {
		name   string
		header string
		value  string
	}{
		{name: "x-api-key", header: "x-api-key", value: "copilot"},
		{name: "api-key", header: "api-key", value: "copilot"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var receivedPath string
			upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
				receivedPath = request.URL.Path
				writer.Header().Set("Content-Type", "application/json")
				_, _ = writer.Write([]byte(`{"ok":true}`))
			}))
			defer upstream.Close()

			dispatcher := newDispatcher(upstream.URL, upstream.Client(),
				[]struct {
					name      string
					wireAPIs  []string
					models    []string
					token     string
					chatPath  string
					msgPath   string
					respPath  string
					embedPath string
				}{
					{name: "copilot", wireAPIs: []string{agentteambridge.WireAPIAnthropic}, models: []string{"claude-sonnet-4.6"}, token: "copilot-token", msgPath: "/v1/messages"},
				},
				[]runtime.ModelRoute{
					{Vendor: "anthropic", Groups: []runtime.RouteGroup{{Models: []string{"claude-sonnet-4.6"}, Providers: []string{"copilot"}}}},
				},
			)

			gin.SetMode(gin.TestMode)
			router := gin.New()
			RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

			request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/messages", bytes.NewBufferString(`{"model":"anthropic/claude-sonnet-4.6","max_tokens":64,"messages":[{"role":"user","content":"hello"}]}`))
			request.Header.Set("Content-Type", "application/json")
			request.Header.Set(tt.header, tt.value)
			recorder := httptest.NewRecorder()
			router.ServeHTTP(recorder, request)

			if recorder.Code != http.StatusOK {
				t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
			}
			if receivedPath != "/v1/messages" {
				t.Fatalf("upstream path = %q want %q", receivedPath, "/v1/messages")
			}
		})
	}
}

func TestWorkspaceScopedProxyAllowsFullModelIDsWithoutVendorRouteMatch(t *testing.T) {
	var received struct {
		Path  string
		Model string
		Auth  string
	}
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		body, _ := io.ReadAll(request.Body)
		var payload map[string]any
		_ = json.Unmarshal(body, &payload)
		received.Path = request.URL.Path
		received.Model, _ = payload["model"].(string)
		received.Auth = request.Header.Get("Authorization")
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "opvd_1", wireAPIs: []string{agentteambridge.WireAPIAnthropic}, models: []string{"zhipu/glm-5.1"}, token: "upstream-token", msgPath: "/v1/messages"},
		},
		[]runtime.ModelRoute{
			{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"zhipu/glm-5.1"}, Providers: []string{"opvd_1"}}}},
		},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeWorkspaceProxyKeyAuthenticator{}})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/llm/v1/messages", bytes.NewBufferString(`{"model":"zhipu/glm-5.1","max_tokens":64,"messages":[{"role":"user","content":"hello"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer lpi_test")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	if received.Path != "/v1/messages" {
		t.Fatalf("upstream path = %q want %q", received.Path, "/v1/messages")
	}
	if received.Model != "zhipu/glm-5.1" {
		t.Fatalf("rewritten model = %q want %q", received.Model, "zhipu/glm-5.1")
	}
	if received.Auth != "Bearer upstream-token" {
		t.Fatalf("auth = %q want %q", received.Auth, "Bearer upstream-token")
	}
}

func TestResponsesFallbackToChatCompletions(t *testing.T) {
	var receivedPath string
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		receivedPath = request.URL.Path
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"id":"chatcmpl_2","choices":[{"message":{"role":"assistant","content":"hello from responses fallback"}}],"usage":{"prompt_tokens":5,"completion_tokens":9}}`))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/responses", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","input":"hello"}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer openai-main")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	if receivedPath != "/chat/completions" {
		t.Fatalf("upstream path = %q want %q", receivedPath, "/chat/completions")
	}
}

func TestResponsesFallbackToChatCompletionsWhenResponsesEndpointMissing(t *testing.T) {
	var receivedPaths []string
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		receivedPaths = append(receivedPaths, request.URL.Path)
		if request.URL.Path == "/responses" {
			writer.Header().Set("Content-Type", "application/json")
			writer.WriteHeader(http.StatusNotFound)
			_, _ = writer.Write([]byte(`{"error":"not found","path":"/responses"}`))
			return
		}
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"id":"chatcmpl_2","choices":[{"message":{"role":"assistant","content":"hello from responses fallback"}}],"usage":{"prompt_tokens":5,"completion_tokens":9}}`))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI, agentteambridge.WireAPIResponses}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions", respPath: "/responses"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/responses", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","input":"hello"}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer openai-main")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	if len(receivedPaths) != 2 {
		t.Fatalf("received paths = %v want 2 requests", receivedPaths)
	}
	if receivedPaths[0] != "/responses" || receivedPaths[1] != "/chat/completions" {
		t.Fatalf("received paths = %v want [/responses /chat/completions]", receivedPaths)
	}
}

func TestCountTokensEndpoint(t *testing.T) {
	dispatcher := newDispatcher("", nil,
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "copilot", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"claude-sonnet-4.6"}, token: "copilot-token"},
		},
		[]runtime.ModelRoute{{Vendor: "anthropic", Groups: []runtime.RouteGroup{{Models: []string{"claude-sonnet-4.6"}, Providers: []string{"copilot"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

	body := `{"model":"anthropic/claude-sonnet-4.6","messages":[{"role":"user","content":"hello world"}]}`
	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/messages/count_tokens", bytes.NewBufferString(body))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer copilot")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	var payload countTokensResponse
	if err := json.Unmarshal(recorder.Body.Bytes(), &payload); err != nil {
		t.Fatalf("invalid response json: %v", err)
	}
	if payload.InputTokens <= 0 {
		t.Fatalf("input_tokens = %d want > 0", payload.InputTokens)
	}
}

func TestUsageTrackerRecordsChatCompletionUsage(t *testing.T) {
	tracker := &fakeUsageTracker{}
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"id":"chatcmpl_1","choices":[{"message":{"role":"assistant","content":"hello"}}],"usage":{"prompt_tokens":11,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":3}}}`))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}, UsageTracker: tracker})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer openai-main")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	if len(tracker.checks) != 1 {
		t.Fatalf("quota checks = %d want 1", len(tracker.checks))
	}
	if tracker.checks[0].Model != "gpt-4o-mini" {
		t.Fatalf("quota model = %q want %q", tracker.checks[0].Model, "gpt-4o-mini")
	}
	if len(tracker.records) != 1 {
		t.Fatalf("usage records = %d want 1", len(tracker.records))
	}
	record := tracker.records[0]
	if record.ProviderID != "openai-main" {
		t.Fatalf("provider_id = %q want %q", record.ProviderID, "openai-main")
	}
	if record.InputTokens != 11 || record.CachedInputTokens != 3 || record.OutputTokens != 7 {
		t.Fatalf("usage = (%d,%d,%d) want (11,3,7)", record.InputTokens, record.CachedInputTokens, record.OutputTokens)
	}
	if record.ErrorType != "" {
		t.Fatalf("error_type = %q want empty", record.ErrorType)
	}
}

func TestUsageTrackerBlocksQuotaExceededRequests(t *testing.T) {
	tracker := &fakeUsageTracker{
		decision: &organization.OrganizationLLMQuotaDecision{
			Allowed:         false,
			DailyQuotaUSD:   1.25,
			CurrentSpentUSD: 1.50,
		},
	}
	upstreamHits := 0
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		upstreamHits++
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}, UsageTracker: tracker})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer openai-main")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusTooManyRequests {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	if upstreamHits != 0 {
		t.Fatalf("upstream hits = %d want 0", upstreamHits)
	}
	if len(tracker.records) != 1 {
		t.Fatalf("usage records = %d want 1", len(tracker.records))
	}
	if tracker.records[0].ErrorType != "rate_limit" {
		t.Fatalf("error_type = %q want %q", tracker.records[0].ErrorType, "rate_limit")
	}
}

func TestUsageTrackerRecordsStreamingUsage(t *testing.T) {
	tracker := &fakeUsageTracker{}
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n"))
		_, _ = writer.Write([]byte("data: {\"usage\":{\"prompt_tokens\":13,\"completion_tokens\":5,\"prompt_tokens_details\":{\"cached_tokens\":2}}}\n\n"))
		_, _ = writer.Write([]byte("data: [DONE]\n\n"))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}, UsageTracker: tracker})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer openai-main")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d body=%s", recorder.Code, recorder.Body.String())
	}
	if len(tracker.records) != 1 {
		t.Fatalf("usage records = %d want 1", len(tracker.records))
	}
	record := tracker.records[0]
	if record.InputTokens != 13 || record.CachedInputTokens != 2 || record.OutputTokens != 5 {
		t.Fatalf("usage = (%d,%d,%d) want (13,2,5)", record.InputTokens, record.CachedInputTokens, record.OutputTokens)
	}
}

func TestChatCompletionsStreamingNormalizesMissingToolCallID(t *testing.T) {
	tracker := &fakeUsageTracker{}
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"type\":\"function\",\"function\":{\"name\":\"search\",\"arguments\":\"{\\\"q\\\":\\\"hel\"}}]}}]}\n\n"))
		_, _ = writer.Write([]byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"lo\\\"}\"}}]}}]}\n\n"))
		_, _ = writer.Write([]byte("data: {\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":7}}\n\n"))
		_, _ = writer.Write([]byte("data: [DONE]\n\n"))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}, UsageTracker: tracker})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer openai-main")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
	chunks := collectOpenAIChunkDataLines(t, recorder.Body.String())
	require.Len(t, chunks, 3)
	require.Len(t, chunks[0].Choices, 1)
	require.Len(t, chunks[0].Choices[0].Delta.ToolCalls, 1)
	require.NotEmpty(t, chunks[0].Choices[0].Delta.ToolCalls[0].ID)
	require.Equal(t, chunks[0].Choices[0].Delta.ToolCalls[0].ID, chunks[1].Choices[0].Delta.ToolCalls[0].ID)
	require.Len(t, tracker.records, 1)
	require.EqualValues(t, 11, tracker.records[0].InputTokens)
	require.EqualValues(t, 7, tracker.records[0].OutputTokens)
}

func TestChatCompletionsStreamingNormalizesConflictingToolCallIDs(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_alpha\",\"type\":\"function\",\"function\":{\"name\":\"search\",\"arguments\":\"{\\\"q\\\":\\\"hel\"}}]}}]}\n\n"))
		_, _ = writer.Write([]byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_beta\",\"function\":{\"arguments\":\"lo\\\"}\"}}]}}]}\n\n"))
		_, _ = writer.Write([]byte("data: [DONE]\n\n"))
	}))
	defer upstream.Close()

	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer openai-main")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
	chunks := collectOpenAIChunkDataLines(t, recorder.Body.String())
	require.Len(t, chunks, 2)
	require.Equal(t, "call_alpha", chunks[0].Choices[0].Delta.ToolCalls[0].ID)
	require.Equal(t, "call_alpha", chunks[1].Choices[0].Delta.ToolCalls[0].ID)
}

func TestConcurrencyLimitReturns429(t *testing.T) {
	t.Cleanup(func() {
		concurrency.DefaultLimiter().Update(concurrency.Config{Default: 5})
	})

	started := make(chan struct{}, 1)
	releaseUpstream := make(chan struct{})
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		started <- struct{}{}
		<-releaseUpstream
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	concurrency.DefaultLimiter().Update(concurrency.Config{Default: 1})
	dispatcher := newDispatcher(upstream.URL, upstream.Client(),
		[]struct {
			name      string
			wireAPIs  []string
			models    []string
			token     string
			chatPath  string
			msgPath   string
			respPath  string
			embedPath string
		}{
			{name: "openai-main", wireAPIs: []string{agentteambridge.WireAIOpenAI}, models: []string{"gpt-4o-mini"}, token: "openai-token", chatPath: "/chat/completions"},
		},
		[]runtime.ModelRoute{{Vendor: "openai", Groups: []runtime.RouteGroup{{Models: []string{"gpt-4o-mini"}, Providers: []string{"openai-main"}}}}},
	)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

	firstDone := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer openai-main")
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)
		firstDone <- rec
	}()
	<-started

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	secondReq := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi again"}]}`)).WithContext(ctx)
	secondReq.Header.Set("Content-Type", "application/json")
	secondReq.Header.Set("Authorization", "Bearer openai-main")
	secondRec := httptest.NewRecorder()
	router.ServeHTTP(secondRec, secondReq)
	if secondRec.Code != http.StatusTooManyRequests {
		t.Fatalf("status = %d body=%s", secondRec.Code, secondRec.Body.String())
	}
	close(releaseUpstream)
	firstRec := <-firstDone
	if firstRec.Code != http.StatusOK {
		t.Fatalf("first status = %d body=%s", firstRec.Code, firstRec.Body.String())
	}
}