package openai

import (
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"iop/apps/edge/internal/authprojection"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

func sha256Hex(raw string) string {
	sum := sha256.Sum256([]byte(raw))
	return hex.EncodeToString(sum[:])
}

func managedProjectionFixture(generation uint64, now time.Time, ttl time.Duration, rawToken string) *iop.PrincipalProjection {
	projection := &iop.PrincipalProjection{
		Generation: generation, IssuedAtUnixNano: now.UnixNano(), ExpiresAtUnixNano: now.Add(ttl).UnixNano(),
	}
	if rawToken != "" {
		projection.Tokens = []*iop.ProjectedPrincipalToken{{
			TokenDigestSha256: sha256Hex(rawToken), PrincipalRef: "managed-principal",
			PrincipalAlias: "managed-alias", TokenRef: "managed-token", TokenRevision: generation,
		}}
	}
	return projection
}

func setManagedPrincipalProjection(srv *Server, projection authprojection.Reader) {
	srv.SetCredentialPlaneManaged(true)
	srv.SetPrincipalProjection(projection)
}

func TestManagedPrincipalMetadataAndStaticCompatibility(t *testing.T) {
	clock := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := authprojection.NewCache(authprojection.DefaultLimits(), func() time.Time { return clock })
	cfg := config.EdgeOpenAIConf{PrincipalTokens: []config.OpenAIPrincipalTokenConf{{
		TokenHashSHA256: sha256Hex("static-token"), PrincipalRef: "static-principal", TokenRef: "static-token-ref",
	}}}
	legacyServer := NewServer(cfg, &fakeRunService{}, nil)
	legacyServer.SetPrincipalProjection(cache)

	var got openAIPrincipal
	legacyHandlerCalls := 0
	legacyHandler := legacyServer.withAuth(func(w http.ResponseWriter, r *http.Request) {
		legacyHandlerCalls++
		got, _ = principalFromContext(r.Context())
		w.WriteHeader(http.StatusNoContent)
	})

	staticReq := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	staticReq.Header.Set("Authorization", "Bearer static-token")
	staticW := httptest.NewRecorder()
	legacyHandler(staticW, staticReq)
	if staticW.Code != http.StatusNoContent || got.PrincipalRef != "static-principal" || got.Source != principalSourceToken {
		t.Fatalf("unmanaged static compatibility: status=%d principal=%+v", staticW.Code, got)
	}

	if err := cache.Apply(managedProjectionFixture(1, clock, time.Minute, "managed-token-value")); err != nil {
		t.Fatal(err)
	}
	staticAfterProjection := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	staticAfterProjection.Header.Set("Authorization", "Bearer static-token")
	staticAfterProjectionW := httptest.NewRecorder()
	legacyHandler(staticAfterProjectionW, staticAfterProjection)
	if staticAfterProjectionW.Code != http.StatusNoContent || got.Source != principalSourceToken {
		t.Fatalf("disabled mode consumed installed projection: status=%d principal=%+v", staticAfterProjectionW.Code, got)
	}

	managedServer := NewServer(cfg, &fakeRunService{}, nil)
	setManagedPrincipalProjection(managedServer, cache)
	managedHandlerCalls := 0
	managedHandler := managedServer.withAuth(func(w http.ResponseWriter, r *http.Request) {
		managedHandlerCalls++
		got, _ = principalFromContext(r.Context())
		w.WriteHeader(http.StatusNoContent)
	})
	managedReq := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	managedReq.Header.Set("Authorization", "Bearer managed-token-value")
	managedW := httptest.NewRecorder()
	managedHandler(managedW, managedReq)
	if managedW.Code != http.StatusNoContent {
		t.Fatalf("managed status=%d body=%s", managedW.Code, managedW.Body.String())
	}
	if got.PrincipalRef != "managed-principal" || got.PrincipalAlias != "managed-alias" || got.TokenRef != "managed-token" || got.Source != principalSourceProjection {
		t.Fatalf("managed principal metadata: %+v", got)
	}

	staticAfterManaged := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	staticAfterManaged.Header.Set("Authorization", "Bearer static-token")
	staticAfterManagedW := httptest.NewRecorder()
	managedHandler(staticAfterManagedW, staticAfterManaged)
	if staticAfterManagedW.Code != http.StatusUnauthorized {
		t.Fatalf("managed mode accepted static source: status=%d", staticAfterManagedW.Code)
	}
	if legacyHandlerCalls != 2 || managedHandlerCalls != 1 {
		t.Fatalf("handler calls: legacy=%d managed=%d, want 2/1", legacyHandlerCalls, managedHandlerCalls)
	}
}

func TestManagedOpenAIAuthRejectsBeforeHandler(t *testing.T) {
	clock := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := authprojection.NewCache(authprojection.DefaultLimits(), func() time.Time { return clock })
	if err := cache.Apply(managedProjectionFixture(1, clock, time.Minute, "managed-token-value")); err != nil {
		t.Fatal(err)
	}
	srv := NewServer(config.EdgeOpenAIConf{BearerToken: "legacy-token"}, &fakeRunService{}, nil)
	setManagedPrincipalProjection(srv, cache)

	handlerCalls := 0
	handler := srv.withAuth(func(w http.ResponseWriter, _ *http.Request) {
		handlerCalls++
		w.WriteHeader(http.StatusNoContent)
	})
	assertRejected := func(name, token string) {
		t.Helper()
		req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
		if token != "" {
			req.Header.Set("Authorization", "Bearer "+token)
		}
		w := httptest.NewRecorder()
		handler(w, req)
		if w.Code != http.StatusUnauthorized {
			t.Fatalf("%s: status=%d body=%s", name, w.Code, w.Body.String())
		}
	}

	assertRejected("missing", "")
	assertRejected("unregistered", "unknown-token")
	assertRejected("legacy source", "legacy-token")
	if handlerCalls != 0 {
		t.Fatalf("rejected authentication reached handler %d times", handlerCalls)
	}

	if err := cache.Apply(managedProjectionFixture(2, clock, time.Minute, "")); err != nil {
		t.Fatal(err)
	}
	assertRejected("revoked", "managed-token-value")
	clock = clock.Add(time.Minute)
	assertRejected("expired", "managed-token-value")
	if handlerCalls != 0 {
		t.Fatalf("revoked/expired authentication reached handler %d times", handlerCalls)
	}
}

func TestManagedOpenAIAuthFailsClosedBeforeProjectionIsUsable(t *testing.T) {
	clock := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	for _, tc := range []struct {
		name       string
		projection authprojection.Reader
	}{
		{name: "projection unavailable"},
		{name: "projection not yet received", projection: authprojection.NewCache(authprojection.DefaultLimits(), func() time.Time { return clock })},
	} {
		t.Run(tc.name, func(t *testing.T) {
			srv := NewServer(config.EdgeOpenAIConf{
				BearerToken: "legacy-token",
				PrincipalTokens: []config.OpenAIPrincipalTokenConf{{
					TokenHashSHA256: sha256Hex("static-token"), PrincipalRef: "static-principal",
				}},
			}, &fakeRunService{}, nil)
			setManagedPrincipalProjection(srv, tc.projection)
			handlerCalls := 0
			handler := srv.withAuth(func(w http.ResponseWriter, _ *http.Request) {
				handlerCalls++
				w.WriteHeader(http.StatusNoContent)
			})
			for _, token := range []string{"legacy-token", "static-token", "managed-token"} {
				req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
				req.Header.Set("Authorization", "Bearer "+token)
				w := httptest.NewRecorder()
				handler(w, req)
				if w.Code != http.StatusUnauthorized {
					t.Fatalf("token %q status = %d, want 401", token, w.Code)
				}
			}
			if handlerCalls != 0 {
				t.Fatalf("unusable managed projection reached handler %d times", handlerCalls)
			}
		})
	}
}

func TestRoutesResolvePrincipalTokenMapping(t *testing.T) {
	const rawToken = "sk-alice-raw-token"
	cfg := config.EdgeOpenAIConf{
		PrincipalTokens: []config.OpenAIPrincipalTokenConf{
			{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
		},
		Models: []string{"model-a"},
	}
	srv := NewServer(cfg, &fakeRunService{}, nil)

	for _, tc := range []struct {
		name       string
		auth       string
		wantStatus int
	}{
		{name: "missing", wantStatus: http.StatusUnauthorized},
		{name: "wrong token", auth: "Bearer wrong-token", wantStatus: http.StatusUnauthorized},
		{name: "valid hash-mapped token", auth: "Bearer " + rawToken, wantStatus: http.StatusOK},
	} {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
			if tc.auth != "" {
				req.Header.Set("Authorization", tc.auth)
			}
			w := httptest.NewRecorder()

			srv.routes().ServeHTTP(w, req)

			if w.Code != tc.wantStatus {
				t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
			}
		})
	}
}

func TestRoutesPrincipalTokenMappingFallsBackToLegacyBearer(t *testing.T) {
	cfg := config.EdgeOpenAIConf{
		PrincipalTokens: []config.OpenAIPrincipalTokenConf{
			{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex("alice-token"), PrincipalRef: "user:alice"},
		},
		BearerToken: "legacy-secret",
		Models:      []string{"model-a"},
	}
	srv := NewServer(cfg, &fakeRunService{}, nil)

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	req.Header.Set("Authorization", "Bearer legacy-secret")
	w := httptest.NewRecorder()

	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
}

func TestResolvePrincipalAllowsMultipleTokensForSamePrincipal(t *testing.T) {
	cfg := config.EdgeOpenAIConf{
		PrincipalTokens: []config.OpenAIPrincipalTokenConf{
			{TokenRef: "alice-app-a", TokenHashSHA256: sha256Hex("alice-app-a-token"), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
			{TokenRef: "alice-app-b", TokenHashSHA256: sha256Hex("alice-app-b-token"), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
		},
	}

	first, ok := resolvePrincipal(cfg, "Bearer alice-app-a-token")
	if !ok {
		t.Fatal("expected first app token to authenticate")
	}
	if first.PrincipalRef != "user:alice" || first.PrincipalAlias != "alice" || first.TokenRef != "alice-app-a" {
		t.Fatalf("unexpected first principal: %+v", first)
	}

	second, ok := resolvePrincipal(cfg, "Bearer alice-app-b-token")
	if !ok {
		t.Fatal("expected second app token to authenticate")
	}
	if second.PrincipalRef != "user:alice" || second.PrincipalAlias != "alice" || second.TokenRef != "alice-app-b" {
		t.Fatalf("unexpected second principal: %+v", second)
	}
}

func TestHealthzDoesNotRequireBearerTokenWithPrincipalTokensConfigured(t *testing.T) {
	cfg := config.EdgeOpenAIConf{
		PrincipalTokens: []config.OpenAIPrincipalTokenConf{
			{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex("alice-token"), PrincipalRef: "user:alice"},
		},
	}
	srv := NewServer(cfg, &fakeRunService{}, nil)

	req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
	w := httptest.NewRecorder()

	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
}

func TestChatCompletionsAddsPrincipalMetadataFromBearerToken(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	const rawToken = "sk-alice-raw-token"
	cfg := config.EdgeOpenAIConf{
		Adapter: "ollama",
		PrincipalTokens: []config.OpenAIPrincipalTokenConf{
			{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
		},
	}
	srv := NewServer(cfg, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"from-request",
		"messages":[{"role":"user","content":"hi"}],
		"metadata":{"iop_principal_ref":"spoofed","user":"someone-else"}
	}`))
	req.Header.Set("Authorization", "Bearer "+rawToken)
	w := httptest.NewRecorder()

	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if fake.req.Metadata["iop_principal_ref"] != "user:alice" {
		t.Fatalf("iop_principal_ref must come from the authenticated context, not caller metadata: got %q", fake.req.Metadata["iop_principal_ref"])
	}
	if fake.req.Metadata["iop_principal_alias"] != "alice" {
		t.Fatalf("iop_principal_alias: got %q", fake.req.Metadata["iop_principal_alias"])
	}
	if fake.req.Metadata["iop_token_ref"] != "iop-tok-alice" {
		t.Fatalf("iop_token_ref: got %q", fake.req.Metadata["iop_token_ref"])
	}
	if fake.req.Metadata["iop_principal_source"] != principalSourceToken {
		t.Fatalf("iop_principal_source: got %q, want %q", fake.req.Metadata["iop_principal_source"], principalSourceToken)
	}
	if fake.req.Metadata["user"] != "someone-else" {
		t.Fatalf("caller metadata.user should still pass through as opaque metadata (not used as identity): %+v", fake.req.Metadata)
	}
}

func TestRoutesDoNotUseMetadataUserForPrincipal(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	// No principal_tokens and no bearer_token configured: the endpoint is
	// open (unauthenticated), so no caller identity should be resolved even
	// though the caller supplies metadata.user.
	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"from-request",
		"messages":[{"role":"user","content":"hi"}],
		"metadata":{"user":"claimed-identity"}
	}`))
	w := httptest.NewRecorder()

	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if _, ok := fake.req.Metadata["iop_principal_ref"]; ok {
		t.Fatalf("iop_principal_ref must not be derived from metadata.user: %+v", fake.req.Metadata)
	}
	if fake.req.Metadata["iop_principal_source"] != principalSourceUnknown {
		t.Fatalf("iop_principal_source: got %q, want %q", fake.req.Metadata["iop_principal_source"], principalSourceUnknown)
	}
	if fake.req.Metadata["user"] != "claimed-identity" {
		t.Fatalf("caller metadata.user should still pass through as opaque metadata: %+v", fake.req.Metadata)
	}
}

func TestResponsesAddsPrincipalMetadataFromBearerToken(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	const rawToken = "sk-bob-raw-token"
	cfg := config.EdgeOpenAIConf{
		Adapter: "ollama",
		PrincipalTokens: []config.OpenAIPrincipalTokenConf{
			{TokenRef: "iop-tok-bob", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:bob"},
		},
	}
	srv := NewServer(cfg, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
		"model":"from-request",
		"input":"hi"
	}`))
	req.Header.Set("Authorization", "Bearer "+rawToken)
	w := httptest.NewRecorder()

	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if fake.req.Metadata["iop_principal_ref"] != "user:bob" {
		t.Fatalf("iop_principal_ref: got %q, want user:bob", fake.req.Metadata["iop_principal_ref"])
	}
	if fake.req.Metadata["iop_token_ref"] != "iop-tok-bob" {
		t.Fatalf("iop_token_ref: got %q", fake.req.Metadata["iop_token_ref"])
	}
	if fake.req.Metadata["iop_principal_alias"] != "" {
		t.Fatalf("iop_principal_alias should be empty when unset: got %q", fake.req.Metadata["iop_principal_alias"])
	}
}

func TestProviderTunnelAddsPrincipalMetadataFromBearerToken(t *testing.T) {
	fake := &providerFakeRunService{
		tunnelFrames: staticProviderTunnelFrames(`{"ok":true}`),
	}
	catalog := []config.ModelCatalogEntry{{
		ID:        "ornith:35b",
		Providers: map[string]string{"prov-vllm": "Ornith-1.0-35B"},
	}}

	const rawToken = "sk-carol-raw-token"
	cfg := config.EdgeOpenAIConf{
		PrincipalTokens: []config.OpenAIPrincipalTokenConf{
			{TokenRef: "iop-tok-carol", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:carol", PrincipalAlias: "carol"},
		},
	}
	srv := NewServer(cfg, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"ornith:35b",
		"messages":[{"role":"user","content":"hello"}]
	}`))
	req.Header.Set("Authorization", "Bearer "+rawToken)
	w := httptest.NewRecorder()

	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	tunnelReqs := fake.tunnelReqsSnapshot()
	if len(tunnelReqs) != 1 {
		t.Fatalf("expected 1 tunnel dispatch, got %d", len(tunnelReqs))
	}
	meta := tunnelReqs[0].Metadata
	if meta["iop_principal_ref"] != "user:carol" {
		t.Fatalf("iop_principal_ref: got %q, want user:carol", meta["iop_principal_ref"])
	}
	if meta["iop_principal_alias"] != "carol" {
		t.Fatalf("iop_principal_alias: got %q", meta["iop_principal_alias"])
	}
	if meta["iop_token_ref"] != "iop-tok-carol" {
		t.Fatalf("iop_token_ref: got %q", meta["iop_token_ref"])
	}
}

func TestChatCompletionsPrincipalMetadataEmptyWhenUnauthenticatedDirectCall(t *testing.T) {
	// Handlers invoked directly (bypassing withAuth, as most existing tests
	// do) must not fail when no principal was ever resolved for the context.
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"from-request",
		"messages":[{"role":"user","content":"hi"}]
	}`))
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if _, ok := fake.req.Metadata["iop_principal_source"]; ok {
		t.Fatalf("no principal metadata should be added without an authenticated context: %+v", fake.req.Metadata)
	}
}