package gateway

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"math/big"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-logr/logr"
	"github.com/go-logr/zapr"
	jwt "github.com/golang-jwt/jwt/v5"
	"github.com/onsi/ginkgo/v2"
	"github.com/onsi/gomega"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/testutil"
	dto "github.com/prometheus/client_model/go"
	jmpapi "github.com/tidbcloud/devops/api/jump"
	itkent "github.com/tidbcloud/devops/entity/itoken"
	jmpent "github.com/tidbcloud/devops/entity/jump"
	"github.com/tidbcloud/devops/pkg/infra/clients/jumppodagentcli"
	"github.com/tidbcloud/devops/pkg/infra/config"
	infralog "github.com/tidbcloud/devops/pkg/infra/log"
	jmpgwmetrics "github.com/tidbcloud/devops/pkg/infra/metrics/jump-gateway"
	"github.com/tidbcloud/devops/pkg/infra/utils/k8s"
	"github.com/tidbcloud/devops/pkg/infra/utils/test"
	"go.uber.org/mock/gomock"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
	"k8s.io/client-go/tools/clientcmd"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
)

func TestJumpController(t *testing.T) {
	gomega.RegisterFailHandler(ginkgo.Fail)
	ginkgo.RunSpecs(t, "Jump Controller Suite", test.ReporterConfig())
}

var _ = ginkgo.Describe("JumpGatewayController proxy", func() {
	var (
		jwtKey   *rsa.PrivateKey
		mockCtrl *gomock.Controller
	)

	ginkgo.BeforeEach(func() {
		var err error
		jwtKey, err = rsa.GenerateKey(rand.Reader, 2048)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		mockCtrl = gomock.NewController(ginkgo.GinkgoT())
		ginkgo.DeferCleanup(mockCtrl.Finish)
	})

	newToken := func(destinations []map[string]any) string {
		claims := jwt.MapClaims{
			"aud":                "jump-gateway",
			"sub":                "alice@pingcap.com",
			"exp":                time.Now().Add(time.Hour).Unix(),
			"user_email":         "alice@pingcap.com",
			"auth_id":            9527,
			"privilege":          "read_only",
			"destinations":       destinations,
			"jump_pod_namespace": "jump-service",
			"jump_pod_name":      "jump-pod-alice",
		}
		token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
		token.Header["kid"] = "kid-1"
		signed, err := token.SignedString(jwtKey)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		return signed
	}

	buildKubeconfig := func(serverURL string) []byte {
		cfg := k8s.NewConfigBuilder().
			AddCluster("gateway_dedicated.123", &clientcmdapi.Cluster{Server: serverURL}).
			AddTokenAuth("gateway-user", "kube-token").
			AddContext("gateway_dedicated.123", clientcmdapi.Context{Cluster: "gateway_dedicated.123", AuthInfo: "gateway-user"}).
			CurrentContext("gateway_dedicated.123").
			Build()
		data, err := clientcmd.Write(*cfg)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		return data
	}

	newMockAgentClient := func(expectFn func(*jumppodagentcli.MockClient)) jumppodagentcli.Client {
		mockClient := jumppodagentcli.NewMockClient(mockCtrl)
		if expectFn != nil {
			expectFn(mockClient)
		}
		return mockClient
	}

	setupAccessLogger := func() *observer.ObservedLogs {
		old := infralog.AccessLogger
		core, logs := observer.New(zapcore.InfoLevel)
		infralog.AccessLogger = zapr.NewLogger(zap.New(core))
		ginkgo.DeferCleanup(func() {
			infralog.AccessLogger = old
		})
		return logs
	}

	newEngine := func(logger logr.Logger, client jumppodagentcli.Client) *gin.Engine {
		jwks := buildJWK(&jwtKey.PublicKey, "kid-1")
		jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_ = json.NewEncoder(w).Encode(map[string]any{
				"keys": []any{jwks},
			})
		}))
		ginkgo.DeferCleanup(jwksSrv.Close)

		auth, err := NewAuth(&config.JumpGatewayAuthConfig{
			DevOpsJWKSURL: jwksSrv.URL,
			Audience:      "jump-gateway",
		})
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		eng := gin.New()
		eng.Use(gin.Recovery())
		eng.Use(auth.Middleware)
		ctrl := NewJumpGatewayController(logger)
		if client != nil {
			ctrl.clientFactory = func(_, _ string) jumppodagentcli.Client {
				return client
			}
		}
		ctrl.RegisterAPI(eng)
		return eng
	}

	ginkgo.It("proxies request with gateway credential", func() {
		var forwardedAuthHeader string

		apiserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			forwardedAuthHeader = r.Header.Get("Authorization")
			_, _ = w.Write([]byte("ok"))
		}))
		defer apiserver.Close()

		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).DoAndReturn(
				func(_ context.Context, req jmpapi.GetKubernetesCredentialRequest) ([]byte, error) {
					gomega.Expect(req.CredIDs).To(gomega.Equal([]jmpent.CredentialID{"gateway_dedicated.123"}))
					return buildKubeconfig(apiserver.URL), nil
				},
			)
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods?x=1", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		body, _ := io.ReadAll(resp.Body)
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK))
		gomega.Expect(string(body)).To(gomega.Equal("ok"))
		gomega.Expect(forwardedAuthHeader).To(gomega.BeEmpty())
	})

	ginkgo.It("returns 403 when destination is not allowed", func() {
		accessLogs := setupAccessLogger()

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), nil))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "999"}},
		}))
		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusForbidden))
		gomega.Expect(accessLogs.Len()).To(gomega.Equal(0))
	})

	ginkgo.It("returns 401 when bearer token is missing", func() {
		httpSrv := httptest.NewServer(newEngine(logr.Discard(), nil))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusUnauthorized))
	})

	ginkgo.It("returns 401 when token is invalid", func() {
		httpSrv := httptest.NewServer(newEngine(logr.Discard(), nil))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer invalid-token")

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusUnauthorized))
	})

	ginkgo.It("returns 400 for invalid encoded credential", func() {
		httpSrv := httptest.NewServer(newEngine(logr.Discard(), nil))
		defer httpSrv.Close()

		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/not!base64/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))
		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusBadRequest))
	})

	ginkgo.It("returns 500 when fetching credential fails", func() {
		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).Return(nil, errors.New("not found"))
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))
		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusInternalServerError))
	})

	ginkgo.It("passes through upstream status", func() {
		apiserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusUnauthorized)
			_, _ = w.Write([]byte("upstream unauthorized"))
		}))
		defer apiserver.Close()

		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).Return(buildKubeconfig(apiserver.URL), nil)
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		body, _ := io.ReadAll(resp.Body)
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusUnauthorized))
		gomega.Expect(bytes.Contains(body, []byte("upstream unauthorized"))).To(gomega.BeTrue())
	})

	ginkgo.It("returns 502 when proxy upstream fails", func() {
		apiserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
		apiserverURL := apiserver.URL
		apiserver.Close()

		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).Return(buildKubeconfig(apiserverURL), nil)
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusBadGateway))
	})

	ginkgo.It("proxies cluster root path to upstream root", func() {
		var gotPath string
		apiserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			gotPath = r.URL.Path
			_, _ = w.Write([]byte("ok-root"))
		}))
		defer apiserver.Close()

		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).Return(buildKubeconfig(apiserver.URL), nil)
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred, nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		body, _ := io.ReadAll(resp.Body)
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK))
		gomega.Expect(string(body)).To(gomega.Equal("ok-root"))
		gomega.Expect(gotPath).To(gomega.Equal("/"))
	})

	ginkgo.It("writes proxy access log for successful upstream response", func() {
		accessLogs := setupAccessLogger()

		apiserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_, _ = w.Write([]byte("ok"))
		}))
		defer apiserver.Close()

		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).Return(buildKubeconfig(apiserver.URL), nil)
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods?x=1", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK))

		gomega.Expect(accessLogs.Len()).To(gomega.Equal(1))
		entry := accessLogs.All()[0]
		fields := entry.ContextMap()
		gomega.Expect(entry.Message).To(gomega.Equal("/api/v1/pods"))
		gomega.Expect(fields["method"]).To(gomega.Equal(http.MethodGet))
		gomega.Expect(fields["status_code"]).To(gomega.BeEquivalentTo(http.StatusOK))
		gomega.Expect(fields["query"]).To(gomega.Equal("x=1"))
		gomega.Expect(fields["proxy_path"]).To(gomega.Equal("/api/v1/pods"))
		gomega.Expect(fields["credential_id"]).To(gomega.BeEquivalentTo("gateway_dedicated.123"))
		gomega.Expect(fields["user_email"]).To(gomega.Equal("alice@pingcap.com"))
		gomega.Expect(fields["auth_id"]).To(gomega.BeEquivalentTo(9527))
		gomega.Expect(fields["jump_pod_namespace"]).To(gomega.Equal("jump-service"))
		gomega.Expect(fields["jump_pod_name"]).To(gomega.Equal("jump-pod-alice"))
		gomega.Expect(fields["trace_id"]).To(gomega.Equal("none"))
		gomega.Expect(fields).To(gomega.HaveKey("latency"))
	})

	ginkgo.It("writes proxy access log for non-2xx upstream response", func() {
		accessLogs := setupAccessLogger()

		apiserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusUnauthorized)
			_, _ = w.Write([]byte("upstream unauthorized"))
		}))
		defer apiserver.Close()

		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).Return(buildKubeconfig(apiserver.URL), nil)
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusUnauthorized))

		gomega.Expect(accessLogs.Len()).To(gomega.Equal(1))
		fields := accessLogs.All()[0].ContextMap()
		gomega.Expect(fields["status_code"]).To(gomega.BeEquivalentTo(http.StatusUnauthorized))
		gomega.Expect(fields["credential_id"]).To(gomega.BeEquivalentTo("gateway_dedicated.123"))
		gomega.Expect(fields["user_email"]).To(gomega.Equal("alice@pingcap.com"))
	})

	ginkgo.It("records proxy metrics for proxied request", func() {
		method := http.MethodGet
		code := "200"
		targetType := "dedicated"

		beforeReqTotal := testutil.ToFloat64(jmpgwmetrics.Proxy.RequestsTotal.WithLabelValues(method, code, targetType))
		beforeReqDur := histogramSampleCount(jmpgwmetrics.Proxy.RequestDuration, map[string]string{
			"method":      method,
			"code":        code,
			"target_type": targetType,
		})

		apiserver := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_, _ = w.Write([]byte("ok"))
		}))
		defer apiserver.Close()

		mockAgentClient := newMockAgentClient(func(m *jumppodagentcli.MockClient) {
			m.EXPECT().GetKubernetesCredential(gomock.Any(), gomock.Any()).Return(buildKubeconfig(apiserver.URL), nil)
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), mockAgentClient))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), method, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "123"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK))

		afterReqTotal := testutil.ToFloat64(jmpgwmetrics.Proxy.RequestsTotal.WithLabelValues(method, code, targetType))
		afterReqDur := histogramSampleCount(jmpgwmetrics.Proxy.RequestDuration, map[string]string{
			"method":      method,
			"code":        code,
			"target_type": targetType,
		})

		gomega.Expect(afterReqTotal).To(gomega.Equal(beforeReqTotal + 1))
		gomega.Expect(afterReqDur).To(gomega.Equal(beforeReqDur + 1))
	})

	ginkgo.It("does not record proxy metrics for pre-proxy rejection", func() {
		method := http.MethodPost
		code := "403"
		targetType := "dedicated"

		beforeReqTotal := testutil.ToFloat64(jmpgwmetrics.Proxy.RequestsTotal.WithLabelValues(method, code, targetType))
		beforeReqDur := histogramSampleCount(jmpgwmetrics.Proxy.RequestDuration, map[string]string{
			"method":      method,
			"code":        code,
			"target_type": targetType,
		})

		httpSrv := httptest.NewServer(newEngine(logr.Discard(), nil))
		defer httpSrv.Close()

		cred := base64.RawURLEncoding.EncodeToString([]byte("gateway_dedicated.123"))
		req, err := http.NewRequestWithContext(context.Background(), method, httpSrv.URL+"/k8s/clusters/"+cred+"/api/v1/pods", nil)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		req.Header.Set("Authorization", "Bearer "+newToken([]map[string]any{
			{"dedicated_cluster": map[string]any{"cluster_id": "999"}},
		}))

		resp, err := http.DefaultClient.Do(req)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		defer resp.Body.Close()
		gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusForbidden))

		afterReqTotal := testutil.ToFloat64(jmpgwmetrics.Proxy.RequestsTotal.WithLabelValues(method, code, targetType))
		afterReqDur := histogramSampleCount(jmpgwmetrics.Proxy.RequestDuration, map[string]string{
			"method":      method,
			"code":        code,
			"target_type": targetType,
		})

		gomega.Expect(afterReqTotal).To(gomega.Equal(beforeReqTotal))
		gomega.Expect(afterReqDur).To(gomega.Equal(beforeReqDur))
	})
})

func buildJWK(pub *rsa.PublicKey, kid string) itkent.JWK {
	return itkent.JWK{
		Kty: "RSA",
		Use: "sig",
		Alg: jwt.SigningMethodRS256.Alg(),
		Kid: kid,
		N:   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
		E:   base64.RawURLEncoding.EncodeToString(bigIntFromInt(pub.E).Bytes()),
	}
}

func bigIntFromInt(v int) *big.Int {
	return big.NewInt(int64(v))
}

func histogramSampleCount(vec *prometheus.HistogramVec, labels map[string]string) uint64 {
	ch := make(chan prometheus.Metric, 32)
	vec.Collect(ch)
	close(ch)

	for metric := range ch {
		dtoMetric := &dto.Metric{}
		err := metric.Write(dtoMetric)
		gomega.Expect(err).ToNot(gomega.HaveOccurred())

		if matchLabels(dtoMetric.GetLabel(), labels) {
			return dtoMetric.GetHistogram().GetSampleCount()
		}
	}

	return 0
}

func matchLabels(pairs []*dto.LabelPair, labels map[string]string) bool {
	for key, expected := range labels {
		found := false
		for _, pair := range pairs {
			if pair.GetName() == key && pair.GetValue() == expected {
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}

	return true
}