package server

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"net"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	gapi "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api/rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/cache"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/ioz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/jwttool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/nettool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/prototool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/redistool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/tlstool"
	otelmetric "go.opentelemetry.io/otel/metric"
	"google.golang.org/grpc"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/apimachinery/pkg/util/sets"
)

const (
	k8sAPIRequestCountKnownMetric = "k8s_api_proxy_request"
	// `ci_access` metric names
	k8sAPIProxyRequestsViaCIAccessMetricName             = "k8s_api_proxy_requests_via_ci_access"
	k8sAPIProxyRequestsUniqueAgentsViaCIAccessMetricName = "k8s_api_proxy_requests_unique_agents_via_ci_access"
	// `user_access` metric names
	k8sAPIProxyRequestsViaUserAccessMetricName             = "k8s_api_proxy_requests_via_user_access"
	k8sAPIProxyRequestsUniqueAgentsViaUserAccessMetricName = "k8s_api_proxy_requests_unique_agents_via_user_access"
	// PAT access metric names
	k8sAPIProxyRequestsViaPatAccessMetricName             = "k8s_api_proxy_requests_via_pat_access"
	k8sAPIProxyRequestsUniqueAgentsViaPatAccessMetricName = "k8s_api_proxy_requests_unique_agents_via_pat_access"
	// Event names
	k8sAPIProxyRequestsUniqueUsersViaCIAccessEventName   = "k8s_api_proxy_requests_unique_users_via_ci_access"
	k8sAPIProxyRequestsUniqueUsersViaUserAccessEventName = "k8s_api_proxy_requests_unique_users_via_user_access"
	k8sAPIProxyRequestsUniqueUsersViaPatAccessEventName  = "k8s_api_proxy_requests_unique_users_via_pat_access"

	k8sAPIProxyAggregatedWatchCountMetricName = "k8s_api_proxy_aggregated_watch"

	requestsTotalCounterName   = "k8s_api_proxy_requests"
	requestsSuccessCounterName = "k8s_api_proxy_requests_success"
	requestsErrorCounterName   = "k8s_api_proxy_requests_error"
)

type Factory struct {
	WebSocketTokenJWTSecretFile string
	FFCheckCounter              otelmetric.Int64Counter
}

func (f *Factory) New(config *modserver.Config) (modserver.Module, error) {
	k8sAPI := config.Config.Agent.KubernetesApi
	listenCfg := k8sAPI.Listen
	certFile := listenCfg.CertificateFile
	keyFile := listenCfg.KeyFile
	var listener func(context.Context) (net.Listener, error)

	tlsConfig, err := tlstool.MaybeServerConfig(certFile, keyFile)
	if err != nil {
		return nil, err
	}
	if tlsConfig != nil {
		// This enables HTTP/2 when the listener is configured with TLS.
		tlsConfig.NextProtos = []string{httpz.TLSNextProtoH2, httpz.TLSNextProtoH1}
		listener = func(ctx context.Context) (net.Listener, error) {
			// stdlib's HTTP server requires tls.Conn for HTTP/2, so we cannot wrap connections here.
			// See https://pkg.go.dev/net/http#Server.Serve.
			return nettool.TLSListenWithOSTCPKeepAlive(ctx, *listenCfg.Network, listenCfg.Address, tlsConfig)
		}
	} else {
		listener = func(ctx context.Context) (net.Listener, error) {
			lis, lisErr := nettool.ListenWithOSTCPKeepAlive(ctx, *listenCfg.Network, listenCfg.Address)
			if lisErr != nil {
				return nil, lisErr
			}
			// Proxied traffic goes via the private API server, so use max connection age from it.
			wrappedLis, lisErr := config.ListenerMetrics(lis, "kubernetes_api", config.Config.PrivateApi.Listen.MaxConnectionAge.AsDuration(), false)
			if lisErr != nil {
				_ = lis.Close()
				return nil, lisErr
			}
			return wrappedLis, nil
		}
	}
	var allowedOriginURLs []string
	if u := config.Config.Gitlab.GetExternalUrl(); u != "" {
		allowedOriginURLs = append(allowedOriginURLs, u)
	}
	allowedAgentCacheTTL := k8sAPI.AllowedAgentCacheTtl.AsDuration()
	allowedAgentCacheErrorTTL := k8sAPI.AllowedAgentCacheErrorTtl.AsDuration()
	tracer := config.TraceProvider.Tracer(kubernetes_api.ModuleName)

	aggregatedWatchCounter, err := config.Meter.Int64UpDownCounter(
		k8sAPIProxyAggregatedWatchCountMetricName,
		otelmetric.WithDescription("Number of currently running aggregated watch proxy requests"))
	if err != nil {
		return nil, err
	}

	// Kubernetes API proxy secret
	webSocketTokenJWTSecret, err := ioz.LoadSHA3_512Base64Secret(config.Log, f.WebSocketTokenJWTSecretFile)
	if err != nil {
		return nil, err
	}
	wst := &webSocketToken{
		validator:        config.Validator,
		jwtSigningMethod: jwttool.SigningMethodHS3_512,
		jwtSecret:        webSocketTokenJWTSecret,
		ffCheckCounter:   f.FFCheckCounter,
	}

	requestsTotalCounter, err := config.Meter.Int64Counter(
		requestsTotalCounterName,
		otelmetric.WithDescription("The total number of Kubernetes API proxy requests"),
	)
	if err != nil {
		return nil, err
	}

	requestsSuccessCounter, err := config.Meter.Int64Counter(
		requestsSuccessCounterName,
		otelmetric.WithDescription("The number of successful Kubernetes API proxy requests"),
	)
	if err != nil {
		return nil, err
	}

	requestsErrorCounter, err := config.Meter.Int64Counter(
		requestsErrorCounterName,
		otelmetric.WithDescription("The number of Kubernetes API proxy requests with errors"),
	)
	if err != nil {
		return nil, err
	}

	m := &module{
		log: config.Log,
		proxy: kubernetesAPIProxy{
			log: config.Log,
			api: config.API,
			makeRequest: func(ctx context.Context, agentKey api.AgentKey) (grpctool.HTTPRequestClient, error) {
				return rpc.NewKubernetesApiClient(config.AgentConnPool(agentKey)).
					MakeRequest(ctx, grpc.WaitForReady(true))
			},
			watchGraph: func(ctx context.Context, agentKey api.AgentKey, r *rpc.WatchGraphRequest) (grpc.ServerStreamingClient[rpc.WatchGraphResponse], error) {
				client := rpc.NewKubernetesApiClient(config.AgentConnPool(agentKey))
				if r.Roots == nil {
					return client.WatchGraph(ctx, r, grpc.WaitForReady(true))
				} else {
					return client.WatchGraphWithRoots(ctx, r, grpc.WaitForReady(true))
				}
			},
			gitLabClient:      config.GitLabClient,
			allowedOriginURLs: allowedOriginURLs,
			allowedAgentsCache: cache.NewWithError[string, *gapi.AllowedAgentsForJob](
				allowedAgentCacheTTL,
				allowedAgentCacheErrorTTL,
				&redistool.ErrCacher[string]{
					Log:          config.Log,
					ErrRep:       modshared.APIToErrReporter(config.API),
					Client:       config.RedisClient,
					ErrMarshaler: prototool.ProtoErrMarshaler{},
					KeyToRedisKey: func(jobToken string) string {
						// Hash half of the token. Even if that hash leaks, it's not a big deal.
						// We do the same in api.AgentToken2key().
						n := len(jobToken) / 2
						tokenHash := sha256.Sum256([]byte(jobToken[:n]))
						tokenHashStr := base64.StdEncoding.EncodeToString(tokenHash[:])
						return config.Config.Redis.KeyPrefix + ":allowed_agents_errs:" + tokenHashStr
					},
				},
				tracer,
				gapi.IsCacheableError,
			),
			authorizeProxyUserCache: cache.NewWithError[proxyUserCacheKey, *gapi.AuthorizeProxyUserResponse](
				allowedAgentCacheTTL,
				allowedAgentCacheErrorTTL,
				&redistool.ErrCacher[proxyUserCacheKey]{
					Log:           config.Log,
					ErrRep:        modshared.APIToErrReporter(config.API),
					Client:        config.RedisClient,
					ErrMarshaler:  prototool.ProtoErrMarshaler{},
					KeyToRedisKey: getAuthorizedProxyUserCacheKey(config.Config.Redis.KeyPrefix),
				},
				tracer,
				gapi.IsCacheableError,
			),
			utRequestCounter:           config.UsageTracker.RegisterCounter(k8sAPIRequestCountKnownMetric),
			utCIAccessRequestCounter:   config.UsageTracker.RegisterCounter(k8sAPIProxyRequestsViaCIAccessMetricName),
			utCIAccessAgentsCounter:    config.UsageTracker.RegisterUniqueCounter(k8sAPIProxyRequestsUniqueAgentsViaCIAccessMetricName),
			utCIAccessEventTracker:     config.EventTracker.RegisterEvent(k8sAPIProxyRequestsUniqueUsersViaCIAccessEventName),
			utUserAccessRequestCounter: config.UsageTracker.RegisterCounter(k8sAPIProxyRequestsViaUserAccessMetricName),
			utUserAccessAgentsCounter:  config.UsageTracker.RegisterUniqueCounter(k8sAPIProxyRequestsUniqueAgentsViaUserAccessMetricName),
			utUserAccessEventTracker:   config.EventTracker.RegisterEvent(k8sAPIProxyRequestsUniqueUsersViaUserAccessEventName),
			utPatAccessRequestCounter:  config.UsageTracker.RegisterCounter(k8sAPIProxyRequestsViaPatAccessMetricName),
			utPatAccessAgentsCounter:   config.UsageTracker.RegisterUniqueCounter(k8sAPIProxyRequestsUniqueAgentsViaPatAccessMetricName),
			utPatAccessEventTracker:    config.EventTracker.RegisterEvent(k8sAPIProxyRequestsUniqueUsersViaPatAccessEventName),
			requestsTotalCounter:       requestsTotalCounter,
			requestsSuccessCounter:     requestsSuccessCounter,
			requestsErrorCounter:       requestsErrorCounter,
			responseSerializer:         serializer.NewCodecFactory(runtime.NewScheme()),
			traceProvider:              config.TraceProvider,
			tracePropagator:            config.TracePropagator,
			meterProvider:              config.MeterProvider,
			validator:                  config.Validator,
			aggregatedWatchCounter:     aggregatedWatchCounter,
			serverName:                 config.KASNameVersion,
			serverVersion:              config.Version,
			serverVia:                  "gRPC/1.0 " + config.KASNameVersion,
			urlPathPrefix:              k8sAPI.UrlPathPrefix,
			listenerGracePeriod:        listenCfg.ListenGracePeriod.AsDuration(),
			shutdownGracePeriod:        listenCfg.ShutdownGracePeriod.AsDuration(),
			gitlabReleasesList:         config.GitLabReleasesList,
			httpServerTracing:          config.HTTPServerTracing,
			webSocketToken:             wst,
			allowedResponseHeaders:     defaultAllowedResponseHeaders.Union(sets.New(k8sAPI.ExtraAllowedResponseHeaders...)),
		},
		listener: listener,
	}
	config.RegisterAgentAPI(&rpc.KubernetesApi_ServiceDesc)
	return m, nil
}

func (f *Factory) Name() string {
	return kubernetes_api.ModuleName
}

func getAuthorizedProxyUserCacheKey(redisKeyPrefix string) redistool.KeyToRedisKey[proxyUserCacheKey] {
	return func(key proxyUserCacheKey) string {
		// Hash half of the token. Even if that hash leaks, it's not a big deal.
		// We do the same in api.AgentToken2key().
		n := len(key.accessKey) / 2

		// Use delimiters between fields to ensure hash of "ab" + "c" is different from "a" + "bc".
		h := sha256.New()
		id := make([]byte, 8)
		binary.LittleEndian.PutUint64(id, uint64(key.agentKey.ID)) //nolint: gosec
		h.Write(id)
		// Don't need a delimiter here because id is fixed size in bytes
		h.Write([]byte(key.accessType))
		h.Write([]byte{11}) // delimiter
		h.Write([]byte(key.accessKey[:n]))
		h.Write([]byte{11}) // delimiter
		h.Write([]byte(key.csrfToken))
		tokenHash := h.Sum(nil)
		tokenHashStr := base64.StdEncoding.EncodeToString(tokenHash)
		return redisKeyPrefix + ":auth_proxy_user_errs:" + tokenHashStr
	}
}
