package openai

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

// chatProviderAuthServer builds a provider-pool tunnel server with the given
// provider auth forwarding config, mirroring chatSidebandServer's catalog so
// pool-model routes over the provider tunnel path.
func chatProviderAuthServer(frames chan *iop.ProviderTunnelFrame, auth config.EdgeOpenAIProviderAuthConf) (*Server, *providerFakeRunService) {
	fake := &providerFakeRunService{
		poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
		tunnelFrames:     frames,
	}
	srv := NewServer(config.EdgeOpenAIConf{ProviderAuth: auth}, fake, nil)
	srv.SetModelCatalog([]config.ModelCatalogEntry{
		{ID: "pool-model", Providers: map[string]string{"prov-1": "served-model"}},
	})
	return srv, fake
}

func staticOKTunnelFrames(body string) chan *iop.ProviderTunnelFrame {
	frames := make(chan *iop.ProviderTunnelFrame, 3)
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(body)}
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
	close(frames)
	return frames
}

// TestChatProviderTunnelForwardsProviderAuthHeader verifies SDD S03: a raw user
// token supplied in the configured provider auth request header is forwarded to
// the provider tunnel request as the configured target header, formatted with
// the configured scheme, while already-prefixed values are preserved.
func TestChatProviderTunnelForwardsProviderAuthHeader(t *testing.T) {
	auth := config.EdgeOpenAIProviderAuthConf{
		Enabled:      true,
		FromHeader:   "X-IOP-Provider-Authorization",
		TargetHeader: "Authorization",
		Scheme:       "Bearer",
		Required:     true,
	}
	for _, tc := range []struct {
		name       string
		callerAuth string
		wantHeader string
	}{
		{name: "raw token gets scheme", callerAuth: "user-token", wantHeader: "Bearer user-token"},
		{name: "already prefixed preserved", callerAuth: "Bearer user-token", wantHeader: "Bearer user-token"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			srv, fake := chatProviderAuthServer(staticOKTunnelFrames(`{"ok":true}`), auth)
			req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
				"model":"pool-model",
				"messages":[{"role":"user","content":"hello"}]
			}`))
			req.Header.Set("X-IOP-Provider-Authorization", tc.callerAuth)
			w := httptest.NewRecorder()
			srv.handleChatCompletions(w, req)

			if w.Code != http.StatusOK {
				t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
			}
			reqs := fake.tunnelReqsSnapshot()
			if len(reqs) != 1 {
				t.Fatalf("expected 1 tunnel dispatch, got %d", len(reqs))
			}
			// The handler sets auth headers on the tunnel request in the
			// tunnel path of handleChatCompletionsProviderPool.
			if got := reqs[0].Headers["Authorization"]; got != tc.wantHeader {
				t.Fatalf("provider Authorization header: got %q, want %q", got, tc.wantHeader)
			}
		})
	}
}

// TestChatProviderTunnelProviderAuthRequired verifies SDD S03: a missing
// required provider auth header is rejected. For provider-pool tunnel path,
// auth is checked after SubmitProviderPool returns. The tunnel handle is closed
// and a 400 is returned (SDD D01).
func TestChatProviderTunnelProviderAuthRequired(t *testing.T) {
	auth := config.EdgeOpenAIProviderAuthConf{
		Enabled:      true,
		FromHeader:   "X-IOP-Provider-Authorization",
		TargetHeader: "Authorization",
		Scheme:       "Bearer",
		Required:     true,
	}
	srv, fake := chatProviderAuthServer(staticOKTunnelFrames(`{"ok":true}`), auth)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"pool-model",
		"messages":[{"role":"user","content":"hello"}]
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	// Auth required check happens in the tunnel path of handleChatCompletionsProviderPool.
	// Since auth is missing, we get 400 and no tunnel request is sent (SDD S03 pre-dispatch).
	if w.Code != http.StatusBadRequest {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if strings.Contains(w.Body.String(), "user-token") {
		t.Fatalf("error body must not echo a raw token: %s", w.Body.String())
	}
	// No tunnel request should be recorded: PrepareTunnel failed before dispatch.
	if got := len(fake.tunnelReqsSnapshot()); got != 0 {
		t.Fatalf("missing required provider auth must not dispatch tunnel: got %d tunnel requests", got)
	}
}

// TestChatCompletionsProviderPoolTunnelProviderAuthForwardedBeforeDispatch verifies
// SDD S03: when provider auth is configured and present, the auth header lands on
// the actual tunnel request that the service sends, not via late handle mutation.
// The auth check runs inside SubmitProviderPool via PrepareTunnel before the
// tunnel request is built and dispatched.
func TestChatCompletionsProviderPoolTunnelProviderAuthForwardedBeforeDispatch(t *testing.T) {
	auth := config.EdgeOpenAIProviderAuthConf{
		Enabled:      true,
		FromHeader:   "X-IOP-Provider-Authorization",
		TargetHeader: "Authorization",
		Scheme:       "Bearer",
		Required:     true,
	}
	srv, fake := chatProviderAuthServer(staticOKTunnelFrames(`{"ok":true}`), auth)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"pool-model",
		"messages":[{"role":"user","content":"hello"}]
	}`))
	req.Header.Set("X-IOP-Provider-Authorization", "user-token")
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	reqs := fake.tunnelReqsSnapshot()
	if len(reqs) != 1 {
		t.Fatalf("expected 1 tunnel dispatch, got %d", len(reqs))
	}
	// Auth header must be on the recorded tunnel request (i.e. pre-dispatch),
	// not on a late handle mutation.
	if got := reqs[0].Headers["Authorization"]; got != "Bearer user-token" {
		t.Fatalf("tunnel request Authorization header: got %q, want %q", got, "Bearer user-token")
	}
}

// TestChatCompletionsProviderPoolOllamaSelectionAllowsMissingProviderAuth verifies
// SDD D01: a missing required provider auth does NOT block the normalized
// (Ollama/CLI) execution path. Only tunnel dispatch requires the auth header.
func TestChatCompletionsProviderPoolOllamaSelectionAllowsMissingProviderAuth(t *testing.T) {
	auth := config.EdgeOpenAIProviderAuthConf{
		Enabled:      true,
		FromHeader:   "X-IOP-Provider-Authorization",
		TargetHeader: "Authorization",
		Scheme:       "Bearer",
		Required:     true,
	}
	// Ollama path: poolDispatchPath=normalized, no tunnel frames needed.
	fake := &providerFakeRunService{
		poolDispatchPath: string(edgeservice.ProviderPoolPathNormalized),
	}
	catalog := []config.ModelCatalogEntry{
		{ID: "ollama-model", Providers: map[string]string{"prov-ollama": "qwen3.6:35b"}},
	}
	srv := NewServer(config.EdgeOpenAIConf{ProviderAuth: auth}, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"ollama-model",
		"messages":[{"role":"user","content":"hello"}]
	}`))
	// No X-IOP-Provider-Authorization header — should still succeed.
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	// No tunnel dispatch for Ollama-normalized path.
	reqs := fake.tunnelReqsSnapshot()
	if len(reqs) != 0 {
		t.Fatalf("expected 0 tunnel dispatches for Ollama path, got %d", len(reqs))
	}
}

// TestChatProviderTunnelProviderAuthDoesNotReplaceInboundAuth verifies the
// provider token is sourced only from the configured provider auth header, not
// from the inbound IOP Authorization header, keeping IOP inbound auth and
// outbound provider auth separate.
func TestChatProviderTunnelProviderAuthDoesNotReplaceInboundAuth(t *testing.T) {
	auth := config.EdgeOpenAIProviderAuthConf{
		Enabled:      true,
		FromHeader:   "X-IOP-Provider-Authorization",
		TargetHeader: "Authorization",
		Scheme:       "Bearer",
		Required:     true,
	}
	srv, fake := chatProviderAuthServer(staticOKTunnelFrames(`{"ok":true}`), auth)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"pool-model",
		"messages":[{"role":"user","content":"hello"}]
	}`))
	// Inbound IOP auth and the configured provider auth header carry different
	// tokens; the provider header must come from the provider auth header only.
	req.Header.Set("Authorization", "Bearer iop-inbound-token")
	req.Header.Set("X-IOP-Provider-Authorization", "provider-token")
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	reqs := fake.tunnelReqsSnapshot()
	if len(reqs) != 1 {
		t.Fatalf("expected 1 tunnel dispatch, got %d", len(reqs))
	}
	// The tunnel auth headers should come from the provider auth header only,
	// not from the inbound IOP Authorization header.
	if got := reqs[0].Headers["Authorization"]; got != "Bearer provider-token" {
		t.Fatalf("provider Authorization must derive from provider auth header, got %q", got)
	}
}

// TestResponsesProviderTunnelForwardsProviderAuthHeader verifies the provider
// auth helper is applied to the Responses tunnel: a raw request-header token is
// forwarded as the configured target header, and a missing required token is
// rejected before dispatch.
func TestResponsesProviderTunnelForwardsProviderAuthHeader(t *testing.T) {
	auth := config.EdgeOpenAIProviderAuthConf{
		Enabled:      true,
		FromHeader:   "X-IOP-Provider-Authorization",
		TargetHeader: "Authorization",
		Scheme:       "Bearer",
		Required:     true,
	}
	t.Run("forwards raw token with scheme", func(t *testing.T) {
		srv, fake := chatProviderAuthServer(staticOKTunnelFrames(`{"ok":true}`), auth)
		req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"pool-model","input":"hi"}`))
		req.Header.Set("X-IOP-Provider-Authorization", "user-token")
		w := httptest.NewRecorder()
		srv.handleResponses(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
		}
		reqs := fake.tunnelReqsSnapshot()
		if len(reqs) != 1 {
			t.Fatalf("expected 1 tunnel dispatch, got %d", len(reqs))
		}
		if got := reqs[0].Headers["Authorization"]; got != "Bearer user-token" {
			t.Fatalf("provider Authorization header: got %q want %q", got, "Bearer user-token")
		}
	})
	t.Run("missing required token rejected", func(t *testing.T) {
		srv, fake := chatProviderAuthServer(staticOKTunnelFrames(`{"ok":true}`), auth)
		req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"pool-model","input":"hi"}`))
		w := httptest.NewRecorder()
		srv.handleResponses(w, req)

		if w.Code != http.StatusBadRequest {
			t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
		}
		if got := len(fake.tunnelReqsSnapshot()); got != 0 {
			t.Fatalf("missing required provider auth must not dispatch: got %d tunnel requests", got)
		}
	})
}