package gateway

import (
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"aigateway.local/core/internal/apikey"
	auditpkg "aigateway.local/core/internal/audit"
	"aigateway.local/core/internal/provider"
	provideropenai "aigateway.local/core/internal/provider/openai"
)

type principalAuthenticator struct{ principal apikey.Principal }

func (a principalAuthenticator) Authenticate(context.Context, string) error { return nil }
func (a principalAuthenticator) AuthenticatePrincipal(context.Context, string) (apikey.Principal, error) {
	return a.principal, nil
}

type fixedAdmission struct{ decision AdmissionDecision }

func (a fixedAdmission) Allow(context.Context, apikey.Principal, time.Time) (AdmissionDecision, error) {
	return a.decision, nil
}

type recordingTokenQuota struct {
	reservation TokenReservation
	estimate    int64
	actual      int64
}

type fixedRoutingResolver struct {
	adapter ResolvedAdapter
}

type recordingAudit struct{ event auditpkg.Event }

func (r *recordingAudit) Record(event auditpkg.Event) bool { r.event = event; return true }

func (r fixedRoutingResolver) Resolve(string) (ResolvedAdapter, error) { return r.adapter, nil }
func (r fixedRoutingResolver) ModelRoutingEnabled() bool               { return true }
func (r fixedRoutingResolver) ResolveModelRoute(query ModelRouteQuery) (ModelRouteResult, error) {
	if query.Model == "public-chat" {
		return ModelRouteResult{ResolvedAdapter: r.adapter, TargetModel: "upstream-chat", Matched: true, Known: true}, nil
	}
	return ModelRouteResult{}, nil
}

func (q *recordingTokenQuota) Reserve(_ context.Context, _ apikey.Principal, estimate int64, _ time.Time) (TokenReservation, error) {
	q.estimate = estimate
	return q.reservation, nil
}

func (q *recordingTokenQuota) Commit(_ context.Context, _ TokenReservation, actual int64) error {
	q.actual = actual
	return nil
}

func TestProxyRejectsInvalidKey(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("upstream must not be called")
	}))
	defer upstream.Close()
	adapter, err := provideropenai.New(upstream.URL, "upstream-secret")
	if err != nil {
		t.Fatal(err)
	}
	proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1

	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
	request.Header.Set("Authorization", "Bearer wrong")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusUnauthorized {
		t.Fatalf("got %d, want %d", response.Code, http.StatusUnauthorized)
	}
}

func TestProxyRejectsKnownOversizedBody(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("upstream must not be called")
	}))
	defer upstream.Close()
	adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
	proxy := NewProxy(adapter, "gateway-secret", 4, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1

	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("12345"))
	request.Header.Set("Authorization", "Bearer gateway-secret")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("got %d, want %d", response.Code, http.StatusRequestEntityTooLarge)
	}
}

func TestProxyReplacesClientAuthorization(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if got := request.Header.Get("Authorization"); got != "Bearer upstream-secret" {
			t.Fatalf("unexpected upstream authorization: %q", got)
		}
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()
	adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
	proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1

	request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
	request.Header.Set("Authorization", "Bearer gateway-secret")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusOK {
		t.Fatalf("got %d, want %d", response.Code, http.StatusOK)
	}
}

func TestProxyRejectsRateLimitedPrincipalBeforeUpstream(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("upstream must not be called")
	}))
	defer upstream.Close()
	adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
	proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
		APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, RequestsPerMinute: 2,
	}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetAdmissionController(fixedAdmission{decision: AdmissionDecision{
		Allowed: false, Reason: AdmissionRateLimit, Limit: 2, Remaining: 0,
		ResetAt: time.Now().Add(time.Minute), RetryAfter: 30 * time.Second,
	}})

	request := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	request.Header.Set("Authorization", "Bearer gateway-secret")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusTooManyRequests {
		t.Fatalf("got %d, want %d", response.Code, http.StatusTooManyRequests)
	}
	if response.Header().Get("X-RateLimit-Limit") != "2" || response.Header().Get("Retry-After") == "" {
		t.Fatalf("missing rate limit headers: %#v", response.Header())
	}
	if !strings.Contains(response.Body.String(), `"type":"rate_limit_error"`) {
		t.Fatalf("unexpected body: %s", response.Body.String())
	}
}

func TestProxyReconcilesReservedTokensWithUpstreamUsage(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":5,"total_tokens":12}}`))
	}))
	defer upstream.Close()
	adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
	quota := &recordingTokenQuota{reservation: TokenReservation{
		Allowed: true, CounterKey: "tokens", Reserved: 24, Limit: 1000, Remaining: 976, ResetAt: time.Now().Add(time.Hour),
	}}
	proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
		APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, MonthlyTokenQuota: 1000,
	}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
	proxy.SetTokenQuotaController(quota)

	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"test","max_tokens":20}`))
	request.Header.Set("Authorization", "Bearer gateway-secret")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusOK {
		t.Fatalf("got %d, want %d: %s", response.Code, http.StatusOK, response.Body.String())
	}
	if quota.estimate <= 20 {
		t.Fatalf("request budget did not include input tokens: %d", quota.estimate)
	}
	if quota.actual != 12 {
		t.Fatalf("committed %d tokens, want 12", quota.actual)
	}
}

func TestProxyRejectsTokenReservationBeforeUpstream(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("upstream must not be called")
	}))
	defer upstream.Close()
	adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
	quota := &recordingTokenQuota{reservation: TokenReservation{
		Allowed: false, Limit: 100, Remaining: 4, ResetAt: time.Now().Add(time.Hour),
	}}
	proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
		APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, MonthlyTokenQuota: 100,
	}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetTokenQuotaController(quota)

	request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"test","max_output_tokens":20}`))
	request.Header.Set("Authorization", "Bearer gateway-secret")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusTooManyRequests || !strings.Contains(response.Body.String(), `"type":"insufficient_quota"`) {
		t.Fatalf("unexpected response %d: %s", response.Code, response.Body.String())
	}
	if response.Header().Get("X-TokenLimit-Limit") != "100" {
		t.Fatalf("missing token quota headers: %#v", response.Header())
	}
}

func TestProxyRewritesModelAliasBeforeUpstream(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		body, _ := io.ReadAll(request.Body)
		if !strings.Contains(string(body), `"model":"upstream-chat"`) {
			t.Fatalf("model alias was not rewritten: %s", body)
		}
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"choices":[]}`))
	}))
	defer upstream.Close()
	adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
	resolved := ResolvedAdapter{Code: "routed", Adapter: adapter, Capabilities: map[provider.Capability]bool{provider.CapabilityChat: true}}
	proxy := NewDynamicProxy(fixedRoutingResolver{adapter: resolved}, principalAuthenticator{principal: apikey.Principal{Scopes: []string{"gateway:invoke"}}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"public-chat","messages":[]}`))
	request.Header.Set("Authorization", "Bearer gateway-secret")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusOK || response.Header().Get("X-Gateway-Model") != "upstream-chat" || response.Header().Get("X-Gateway-Provider") != "routed" {
		t.Fatalf("unexpected routed response %d %#v: %s", response.Code, response.Header(), response.Body.String())
	}
}

func TestProxyRecordsAuditWithoutBufferingWholeResponse(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"choices":[],"usage":{"prompt_tokens":9,"completion_tokens":4,"total_tokens":13}}`))
	}))
	defer upstream.Close()
	adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
	recorder := &recordingAudit{}
	proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
		APIKeyID: "11111111-1111-4111-8111-111111111111", Scopes: []string{"gateway:invoke"},
	}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
	proxy.SetAllowPrivateProviderURLs(true) // 测试上游监听 127.0.0.1
	proxy.SetAuditRecorder(recorder)
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"audit-model","messages":[]}`))
	request.Header.Set("Authorization", "Bearer gateway-secret")
	response := httptest.NewRecorder()
	proxy.ServeHTTP(response, request)
	if response.Code != http.StatusOK {
		t.Fatalf("unexpected response %d: %s", response.Code, response.Body.String())
	}
	if recorder.event.Model != "audit-model" || recorder.event.ProviderCode != "environment" || recorder.event.StatusCode != 200 || recorder.event.PromptTokens != 9 || recorder.event.CompletionTokens != 4 {
		t.Fatalf("unexpected audit event: %#v", recorder.event)
	}
}