package kas

import (
	"bytes"
	"cmp"
	"context"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"fmt"
	"log/slog"
	"math/rand/v2"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"buf.build/go/protovalidate"
	"github.com/ash2k/stager"
	"github.com/getsentry/sentry-go"
	"github.com/go-logr/logr"
	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/redis/rueidis"
	"github.com/redis/rueidis/rueidisotel"
	"github.com/spf13/cobra"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/cmd"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitaly"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitaly/vendored/client"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab"
	gapi "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab/api"
	agent2kas_tunnel_router "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent2kas_tunnel/router"
	agent2kas_tunnel_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent2kas_tunnel/server"
	agent_configuration_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_configuration/server"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_registrar"
	agent_registrar_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_registrar/server"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_tracker"
	agent_tracker_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_tracker/server"
	autoflow_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/autoflow/server"
	autoflow_codec_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/autoflow_codec_server/server"
	configuration_project_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/configuration_project/server"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/event_tracker"
	event_tracker_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/event_tracker/server"
	flux_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/flux/server"
	gitlab_access_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/gitlab_access/server"
	google_profiler_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/google_profiler/server"
	kas2agentk_tunnel_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kas2agentk_tunnel/server"
	kubernetes_api_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api/server"
	managed_resources_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/managed_resources/server"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	notifications_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/notifications/server"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/observability"
	observability_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/observability/server"
	server_info_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/server_info/server"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/usage_metrics"
	usage_metrics_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/usage_metrics/server"
	workspaces_server "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/workspaces/server"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/cache"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/errz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/featureflag"
	"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/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/metric"
	"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/retry"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/syncz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/tlstool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/version"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/tunserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/agentcfg"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/kascfg"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	promexp "go.opentelemetry.io/otel/exporters/prometheus"
	otelmetric "go.opentelemetry.io/otel/metric"
	"go.opentelemetry.io/otel/propagation"
	metricsdk "go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/resource"
	tracesdk "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.34.0"
	"go.opentelemetry.io/otel/trace"
	"go.opentelemetry.io/otel/trace/noop"
	"golang.org/x/time/rate"
	"google.golang.org/grpc"
	_ "google.golang.org/grpc/encoding/gzip" // Install the gzip compressor
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/stats"
	"google.golang.org/protobuf/encoding/protojson"
	"k8s.io/klog/v2"
)

const (
	routingAttemptInterval   = 50 * time.Millisecond
	routingInitBackoff       = 100 * time.Millisecond
	routingMaxBackoff        = 1 * time.Second
	routingResetDuration     = 10 * time.Second
	routingBackoffFactor     = 2.0
	routingJitter            = 5.0
	routingTunnelFindTimeout = 20 * time.Second
	routingCachePeriod       = 5 * time.Minute
	routingTryNewKASInterval = 10 * time.Millisecond

	kasName = "gitlab-kas"

	kasTracerName = "kas"
	kasMeterName  = "kas"

	gitlabBuildInfoGaugeMetricName               = "gitlab_build_info"
	kasVersionAttr                 attribute.Key = "version"
	kasGitRefAttr                  attribute.Key = "git_ref"

	delayedShutdownGoroutineDumpInterval   = 15 * time.Minute
	delayedShutdownGoroutineDumpEnvVarName = "KAS_DEBUG_DELAYED_SHUTDOWN"
)

type options struct {
	log               *slog.Logger
	grpcLog           *grpctool.Logger
	configurationFile string
	hub               *sentry.Hub
	validator         protovalidate.Validator
	configuration     *kascfg.ConfigurationFile
}

func newOptions() *options {
	return &options{}
}

func (o *options) complete(ctx context.Context) error {
	var err error
	o.configuration, err = LoadConfigurationFile(ctx, o.configurationFile)
	if err != nil {
		return err
	}
	ApplyDefaultsToKASConfigurationFile(o.configuration)
	o.validator, err = protovalidate.New()
	if err != nil {
		return err
	}
	err = o.validator.Validate(o.configuration)
	if err != nil {
		return fmt.Errorf("config validate: %w", err)
	}
	obs := o.configuration.Observability
	o.log, o.grpcLog, err = loggerFromConfig(obs.Logging)
	if err != nil {
		return err
	}
	// Sentry
	o.hub, err = constructSentryHub(obs.Sentry.Dsn, obs.Sentry.Environment)
	if err != nil {
		return fmt.Errorf("sentry: %w", err)
	}
	return nil
}

// setGlobals shouldn't exist, yet here we are.
func (o *options) setGlobals() {
	grpclog.SetLoggerV2(o.grpcLog) // pipe gRPC logs into slog
	logrLogger := logr.FromSlogHandler(o.log.Handler())
	// Kubernetes uses klog so here we pipe all logs from it to our logger via an adapter.
	klog.SetLogger(logrLogger)
	otel.SetLogger(logrLogger)
	otel.SetErrorHandler(&otelErrorHandler{
		log: o.log,
		hub: o.hub,
	})
}

func (o *options) startGoroutineDump(ctx context.Context) {
	if os.Getenv(delayedShutdownGoroutineDumpEnvVarName) == "true" {
		context.AfterFunc(ctx, func() {
			o.log.Info("Starting to collect goroutine dumps to investigate delayed shutdown", slog.Int64("interval_sec", int64(delayedShutdownGoroutineDumpInterval.Seconds())))
			startStacktraceCollector(o.log, o.hub, delayedShutdownGoroutineDumpInterval)
		})
	}
}

func (o *options) run(ctx context.Context) (retErr error) {
	o.log.Info("Running KAS", kasServerNameField())
	defer o.log.Info("KAS shutdown done, exiting")
	defer o.hub.Client().Close()
	defer o.hub.Flush(time.Second)

	// Observability
	ot, stop, err := o.constructObservabilityTools(ctx)
	if err != nil {
		return err
	}
	defer errz.SafeCall(stop, &retErr)

	// Feature Flag check counter
	ffCheckCounter, err := featureflag.NewCheckCounter(ot.meter)
	if err != nil {
		return err
	}

	// GitLab REST client
	gitLabClient, err := o.constructGitLabClient(ot, ffCheckCounter) //nolint:contextcheck
	if err != nil {
		return err
	}

	// Redis
	redisClient, err := o.constructRedisClient(ot)
	if err != nil {
		return fmt.Errorf("redis client: %w", err)
	}
	defer redisClient.Close()
	ot.probeRegistry.RegisterReadinessProbe("redis", constructRedisReadinessProbe(redisClient))

	// Agent tracker
	agentTracker, err := o.constructAgentTracker(ErrReporterWithoutAgentInfos(o.hub), redisClient, ot.meter)
	if err != nil {
		return err
	}

	agentInfoResolver := constructAgentInfoResolver(agentTracker)

	srvAPI := newServerAPI(o.log, o.hub, redisClient, agentInfoResolver)

	errRep := modshared.APIToErrReporter(srvAPI)
	grpcServerErrorReporter := &serverErrorReporter{log: o.log, errReporter: errRep}

	instanceID := rand.Int64() //nolint: gosec
	listenMetric := &nettool.ListenerMetrics{
		Log: o.log,
		NewSet: (&connectionSetFactory{
			instanceID: instanceID,
			redis:      redisClient,
			keyPrefix:  o.configuration.Redis.KeyPrefix,
			m:          ot.meter,
		}).newConnectionSet,
		GCPeriod: 10 * time.Minute,
		ErrRep:   errRep,
	}

	// RPC API factory
	rpcAPIFactory, agentRPCAPIFactory := o.constructRPCAPIFactory(
		errRep,
		gitLabClient,
		redisClient,
		ot.tracer,
		agentInfoResolver,
		ffCheckCounter,
	)

	// Server for handling API requests from other kas instances
	privateAPISrv, err := newPrivateAPIServer(o.log, ot, errRep, o.configuration, rpcAPIFactory, grpcServerErrorReporter, o.validator, listenMetric) //nolint:contextcheck
	if err != nil {
		return fmt.Errorf("private API server: %w", err)
	}
	defer errz.SafeClose(privateAPISrv, &retErr)

	// Server for handling agentk requests
	agentSrv, err := newAgentServer(o.log, ot, o.configuration, srvAPI, redisClient, //nolint: contextcheck
		kas2agentk_tunnel_server.AgentRPCAPIFactory(agentRPCAPIFactory), agentRPCAPIFactory, o.validator,
		privateAPISrv.ownURL, grpcServerErrorReporter, listenMetric)
	if err != nil {
		return fmt.Errorf("agent server: %w", err)
	}
	defer errz.SafeClose(agentSrv, &retErr)

	// Server for handling external requests e.g. from GitLab
	apiSrv, err := newAPIServer(o.log, ot, o.configuration, rpcAPIFactory, grpcServerErrorReporter, o.validator, listenMetric) //nolint: contextcheck
	if err != nil {
		return fmt.Errorf("API server: %w", err)
	}

	// Kas to agentk router
	pollConfig := retry.NewPollConfigFactory(routingAttemptInterval, retry.NewExponentialBackoffFactory(
		routingInitBackoff,
		routingMaxBackoff,
		routingResetDuration,
		routingBackoffFactor,
		routingJitter,
	))
	gatewayQuerier := agent2kas_tunnel_router.NewAggregatingQuerier(o.log, agentSrv.tunnelRegistry, srvAPI, ot.tracer, routingAttemptInterval, routingCachePeriod)
	routerPlugin, err := agent2kas_tunnel_router.NewPlugin(
		srvAPI,
		privateAPISrv.kasPool,
		gatewayQuerier,
		constructAgentFinder(agentTracker),
		agentSrv.tunnelRegistry,
		ot.tracer,
		ot.meter,
		privateAPISrv.ownURL,
		pollConfig,
		routingTryNewKASInterval,
		routingTunnelFindTimeout,
	)
	if err != nil {
		return err
	}
	kasToAgentRouter := &tunserver.Router[nettool.MultiURL]{
		Plugin:           routerPlugin,
		PrivateAPIServer: privateAPISrv.server,
	}

	// Gitaly client
	gitalyClientPool, err := o.constructGitalyPool(ot) //nolint: contextcheck
	if err != nil {
		return err
	}
	defer errz.SafeClose(gitalyClientPool, &retErr)

	// kas->agentk connection pool
	agentPool := newAgentConnPool(&tunserver.RoutingClientConn[nettool.MultiURL]{
		Log:    o.log,
		API:    srvAPI,
		Plugin: routerPlugin,
	})

	// Usage tracker
	usageTracker := usage_metrics.NewUsageTracker()

	// Event tracker
	eventTracker := event_tracker.NewEventTracker()

	poolWrapper := &gitaly.Pool{
		ClientPool: gitalyClientPool,
	}

	v, err := version.NewVersion(cmd.Version)
	if err != nil {
		return err
	}

	// Module factories
	configForFactory := func(name string) *modserver.Config {
		return &modserver.Config{
			Log:                o.log.With(logz.ModuleName(name)),
			InstanceID:         instanceID,
			API:                srvAPI,
			Config:             o.configuration,
			GitLabClient:       gitLabClient,
			UsageTracker:       usageTracker,
			EventTracker:       eventTracker,
			AgentServer:        agentSrv.server,
			APIServer:          apiSrv.server,
			RegisterAgentAPI:   kasToAgentRouter.RegisterTunclientAPI,
			AgentConnPool:      agentPool.Get,
			ListenerMetrics:    listenMetric.Wrap,
			Gitaly:             poolWrapper,
			TraceProvider:      ot.tp,
			TracePropagator:    ot.p,
			MeterProvider:      ot.mp,
			Meter:              ot.meter,
			StreamClientProm:   ot.streamClientProm,
			UnaryClientProm:    ot.unaryClientProm,
			RedisClient:        redisClient,
			KASName:            kasName,
			KASNameVersion:     kasServerName(),
			Version:            v,
			GitRef:             cmd.GitRef,
			ProbeRegistry:      ot.probeRegistry,
			GitLabReleasesList: GitLabReleasesList,
			Validator:          o.validator,
			HTTPServerTracing:  ot.httpServerTracing,
			HTTPClientTracing:  ot.httpClientTracing,
			GRPCServerTracing:  ot.grpcServerTracing,
			GRPCClientTracing:  ot.grpcClientTracing,
		}
	}
	// Start things up. Layers are shut down in reverse order.
	return syncz.RunLayers(ctx,
		syncz.LayerFunc(func(stage stager.Stage) {
			// Start things that modules use.
			stage.Go(agentTracker.Run)
			stage.Go(gatewayQuerier.Run)
			stage.Go(func(ctx context.Context) error {
				listenMetric.RunGC(ctx)
				return nil
			})
		}),
		&syncz.LogWhenDoneLayer{
			Log: o.log,
			Msg: "Stopping agent tracker and gateway querier",
		},
		&factoriesLayer{
			log:    o.log,
			config: configForFactory,
			factories: []modserver.Factory{
				&agent_configuration_server.Factory{
					AgentConfigDefaulters: []func(*agentcfg.ConfigurationFile){
						kubernetes_api_server.ApplyDefaultsToAgentConfig,
					},
				},
				&agent_registrar_server.Factory{AgentRegisterer: agentTracker},
				&agent_tracker_server.Factory{AgentQuerier: agentTracker},
				&agent2kas_tunnel_server.Factory{TunnelHandler: agentSrv.tunnelRegistry},
				&configuration_project_server.Factory{},
				&managed_resources_server.Factory{},
				&event_tracker_server.Factory{EventTracker: eventTracker},
				&flux_server.Factory{},
				&gitlab_access_server.Factory{},
				&google_profiler_server.Factory{},
				&notifications_server.Factory{PublishEvent: srvAPI.publishEvent, SubscribeToEvents: srvAPI.subscribeToEvents},
				&observability_server.Factory{Gatherer: ot.reg, Registerer: ot.reg},
				&server_info_server.Factory{},
				&usage_metrics_server.Factory{UsageTracker: usageTracker},
				&workspaces_server.Factory{},
				&autoflow_server.Factory{
					// TODO these are going to go away once we move to OTEL for metrics.
					CSH:              ot.csh,
					StreamClientProm: ot.streamClientProm,
					UnaryClientProm:  ot.unaryClientProm,
				},
				&autoflow_codec_server.Factory{},
			},
		},
		syncz.SiblingLayers{
			// kas2agentk_tunnel_server starts as a sibling to the servers since it owns and runs the
			// in-memory agent API server for kas->agentk tunnels. Those tunnels can take a while to
			// wrap up handling in-flight requests so it makes sense to shut them down concurrently with
			// the other servers, which too can take a while to wrap up.
			syncz.LayerFunc(func(stage stager.Stage) { //nolint:contextcheck
				// Start gRPC servers.
				agentSrv.Start(stage)
				apiSrv.Start(stage)
				privateAPISrv.Start(stage)
			}),
			&factoriesLayer{
				log:    o.log,
				config: configForFactory,
				factories: []modserver.Factory{
					&kas2agentk_tunnel_server.Factory{
						AgentConnPool:        agentPool,
						AgentInMemServer:     agentSrv.inMemServer,
						AgentInMemListener:   agentSrv.inMemListener,
						AgentInMemServerConn: agentSrv.inMemConn,
						AgentInMemAuxCancel:  agentSrv.inMemAuxCancel,
					},
				},
			},
		},
		&syncz.LogWhenDoneLayer{
			Log: o.log,
			Msg: "Stopping servers",
		},
		&factoriesLayer{
			log:    o.log,
			config: configForFactory,
			factories: []modserver.Factory{
				// Uses agent connection (config.AgentConnPool) so depends on:
				// - private API server for kas <- agentk tunnels.
				// - kas2agentk_tunnel_server module for kas -> agentk tunnels.
				// Hence, must stop before those two.
				&kubernetes_api_server.Factory{
					WebSocketTokenJWTSecretFile: o.configuration.Agent.KubernetesApi.WebsocketTokenSecretFile,
					FFCheckCounter:              ffCheckCounter,
				},
			},
		},
		&syncz.LogWhenDoneLayer{
			Log: o.log,
			Msg: "Starting KAS shutdown",
		},
	)
}

func (o *options) constructRPCAPIFactory(errRep errz.ErrReporter, gitLabClient gitlab.ClientInterface, redisClient rueidis.Client, dt trace.Tracer, agentInfoResolver modserver.AgentInfoResolver, ffCheckCounter otelmetric.Int64Counter) (modshared.RPCAPIFactory, modserver.AgentRPCAPIFactory) {
	aCfg := o.configuration.Agent
	f := serverRPCAPIFactory{
		log:               o.log,
		sentryHub:         o.hub,
		agentInfoResolver: agentInfoResolver,
		ffCheckCounter:    ffCheckCounter,
	}
	fAgent := serverAgentRPCAPIFactory{
		rpcAPIFactory: f.New,
		gitLabClient:  gitLabClient,
		agentInfoCache: cache.NewWithError[api.AgentTokenWithType, api.AgentInfo](
			aCfg.InfoCacheTtl.AsDuration(),
			aCfg.InfoCacheErrorTtl.AsDuration(),
			&redistool.ErrCacher[api.AgentTokenWithType]{
				Log:          o.log,
				ErrRep:       errRep,
				Client:       redisClient,
				ErrMarshaler: prototool.ProtoErrMarshaler{},
				KeyToRedisKey: func(key api.AgentTokenWithType) string {
					tokenHash := api.AgentToken2key(key.Token)
					tokenHashStr := base64.StdEncoding.EncodeToString(tokenHash)
					return o.configuration.Redis.KeyPrefix + ":agent_info_errs:" + tokenHashStr
				},
			},
			dt,
			gapi.IsCacheableError,
		),
	}
	return f.New, fAgent.New
}

func (o *options) constructAgentTracker(errRep errz.ErrReporter, redisClient rueidis.Client, m otelmetric.Meter) (*agent_tracker.RedisTracker, error) {
	t, err := agent_tracker.NewRedisTracker(
		o.log,
		errRep,
		o.validator,
		redisClient,
		o.configuration.Redis.KeyPrefix+":agent_tracker2",
		agent_registrar.RegisterAttemptInterval*3, // TTL
		agent_registrar.RegisterAttemptInterval*2, // GC
		m,
	)
	return t, err
}

func (o *options) constructGitLabClient(ot *obsTools, ffCheckCounter otelmetric.Int64Counter) (*gitlab.Client, error) {
	cfg := o.configuration

	gitLabURL, err := url.Parse(cfg.Gitlab.Address)
	if err != nil {
		return nil, err
	}
	// TLS cert for talking to GitLab/Workhorse.
	clientTLSConfig, err := tlstool.ClientConfigWithCACert(cfg.Gitlab.CaCertificateFile)
	if err != nil {
		return nil, err
	}
	// Secret for JWT signing
	decodedAuthSecret, err := ioz.LoadSHA256Base64Secret(o.log, o.configuration.Gitlab.AuthenticationSecretFile)
	if err != nil {
		return nil, fmt.Errorf("GitLab authentication secret: %w", err)
	}
	var limiter httpz.Limiter
	limiter = rate.NewLimiter(
		rate.Limit(cfg.Gitlab.ApiRateLimit.RefillRatePerSecond),
		int(cfg.Gitlab.ApiRateLimit.BucketSize),
	)
	limiter, err = metric.NewWaitLimiterInstrumentation(
		"gitlab_client",
		cfg.Gitlab.ApiRateLimit.RefillRatePerSecond,
		"{refill/s}",
		ot.tracer,
		ot.meter,
		limiter,
	)
	if err != nil {
		return nil, err
	}
	return gitlab.NewClient(
		gitLabURL,
		decodedAuthSecret,
		gitlab.WithTextMapPropagator(ot.p),
		gitlab.WithTracerProvider(ot.MaybeTraceProvider(ot.httpClientTracing)),
		gitlab.WithMeterProvider(ot.mp),
		gitlab.WithUserAgent(kasServerName()),
		gitlab.WithTLSConfig(clientTLSConfig),
		gitlab.WithRateLimiter(limiter),
		gitlab.WithValidator(o.validator),
		gitlab.WithFeatureFlagCheckCounter(ffCheckCounter),
	), nil
}

func (o *options) constructGitalyPool(ot *obsTools) (*client.Pool, error) {
	g := o.configuration.Gitaly
	var globalGitalyRPCLimiter grpctool.ClientLimiter
	globalGitalyRPCLimiter = rate.NewLimiter(
		rate.Limit(g.GlobalApiRateLimit.RefillRatePerSecond),
		int(g.GlobalApiRateLimit.BucketSize),
	)
	globalGitalyRPCLimiter, err := metric.NewWaitLimiterInstrumentation(
		"gitaly_client_global",
		g.GlobalApiRateLimit.RefillRatePerSecond,
		"{refill/s}",
		ot.tracer,
		ot.meter,
		globalGitalyRPCLimiter,
	)
	if err != nil {
		return nil, err
	}
	return client.NewPoolWithOptions(
		client.WithDialOptions(
			grpc.WithUserAgent(kasServerName()),
			grpc.WithStatsHandler(ot.csh),
			grpc.WithStatsHandler(otelgrpc.NewClientHandler(
				otelgrpc.WithTracerProvider(ot.MaybeTraceProvider(ot.grpcClientTracing)),
				otelgrpc.WithMeterProvider(ot.mp),
				otelgrpc.WithPropagators(ot.p),
				otelgrpc.WithMessageEvents(otelgrpc.ReceivedEvents, otelgrpc.SentEvents),
			)),
			grpc.WithSharedWriteBuffer(true),
			grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(api.GRPCMaxMessageSize)),
			// In https://gitlab.com/groups/gitlab-org/-/epics/8971, we added DNS discovery support to Praefect. This was
			// done by making two changes:
			// - Configure client-side round-robin load-balancing in client dial options. We added that as a default option
			// inside gitaly client in gitaly client since v15.9.0
			// - Configure DNS resolving. Due to some technical limitations, we don't use gRPC's built-in DNS resolver.
			// Instead, we implement our own DNS resolver. This resolver is exposed via the following configuration.
			// Afterward, workhorse can detect and handle DNS discovery automatically. The user needs to setup and set
			// Gitaly address to something like "dns:gitaly.service.dc1.consul"
			client.WithGitalyDNSResolver(client.DefaultDNSResolverBuilderConfig()),
			// Don't put interceptors here as order is important. Put them below.
		),
		client.WithDialer(func(ctx context.Context, address string, dialOptions []grpc.DialOption) (*grpc.ClientConn, error) {
			var perServerGitalyRPCLimiter grpctool.ClientLimiter
			perServerGitalyRPCLimiter = rate.NewLimiter(
				rate.Limit(g.PerServerApiRateLimit.RefillRatePerSecond),
				int(g.PerServerApiRateLimit.BucketSize))
			perServerGitalyRPCLimiter, err := metric.NewWaitLimiterInstrumentation( //nolint: contextcheck
				"gitaly_client_"+address,
				g.GlobalApiRateLimit.RefillRatePerSecond,
				"{refill/s}",
				ot.tracer,
				ot.meter,
				perServerGitalyRPCLimiter,
			)
			if err != nil {
				return nil, err
			}
			opts := []grpc.DialOption{
				grpc.WithChainStreamInterceptor(
					ot.streamClientProm,
					grpctool.StreamClientLimitingInterceptor(globalGitalyRPCLimiter),
					grpctool.StreamClientLimitingInterceptor(perServerGitalyRPCLimiter),
				),
				grpc.WithChainUnaryInterceptor(
					ot.unaryClientProm,
					grpctool.UnaryClientLimitingInterceptor(globalGitalyRPCLimiter),
					grpctool.UnaryClientLimitingInterceptor(perServerGitalyRPCLimiter),
				),
			}
			opts = append(opts, dialOptions...)
			return client.DialContext(ctx, address, opts)
		}),
	), nil
}

func (o *options) constructRedisClient(ot *obsTools) (rueidis.Client, error) {
	cfg := o.configuration.Redis
	dialTimeout := cfg.DialTimeout.AsDuration()
	writeTimeout := cfg.WriteTimeout.AsDuration()
	var err error
	var tlsConfig *tls.Config
	if cfg.Tls != nil && cfg.Tls.Enabled {
		tlsConfig, err = tlstool.ClientConfigWithCACertKeyPair(cfg.Tls.CaCertificateFile, cfg.Tls.CertificateFile, cfg.Tls.KeyFile)
		if err != nil {
			return nil, err
		}
	}
	var password string
	switch {
	case cfg.PasswordFile != "" && cfg.Password != "":
		return nil, errors.New("only one of password or password_file can be specified")
	case cfg.PasswordFile != "":
		passwordBytes, err := os.ReadFile(cfg.PasswordFile)
		if err != nil {
			return nil, err
		}
		password = string(passwordBytes)
	case cfg.Password != "":
		password = cfg.Password
	}
	opts := rueidis.ClientOption{
		Dialer: net.Dialer{
			Timeout: dialTimeout,
		},
		TLSConfig:         tlsConfig,
		Username:          cfg.Username,
		Password:          password,
		ClientName:        kasServerName(),
		ConnWriteTimeout:  writeTimeout,
		DisableCache:      true,
		SelectDB:          int(cfg.DatabaseIndex),
		PipelineMultiplex: int(cfg.PipelineMultiplex),
	}
	if cfg.Network == "unix" {
		opts.DialFn = redistool.UnixDialer
	}
	switch v := cfg.RedisConfig.(type) {
	case *kascfg.RedisCF_Server:
		opts.InitAddress = []string{v.Server.Address}
		if opts.TLSConfig != nil {
			opts.TLSConfig.ServerName, _, _ = strings.Cut(v.Server.Address, ":")
		}
	case *kascfg.RedisCF_Sentinel:
		opts.InitAddress = v.Sentinel.Addresses
		var sentinelPassword string
		if v.Sentinel.SentinelPasswordFile != "" {
			sentinelPasswordBytes, err := os.ReadFile(v.Sentinel.SentinelPasswordFile)
			if err != nil {
				return nil, err
			}
			sentinelPassword = string(sentinelPasswordBytes)
		}
		opts.Sentinel = rueidis.SentinelOption{
			Dialer:    opts.Dialer,
			TLSConfig: opts.TLSConfig,
			MasterSet: v.Sentinel.MasterName,
			Username:  cmp.Or(v.Sentinel.Username, cfg.Username),
			Password:  sentinelPassword,
		}
	default:
		// This should never happen
		return nil, fmt.Errorf("unexpected Redis config type: %T", cfg.RedisConfig)
	}
	// Instrument Redis client with tracing and metrics.
	return rueidisotel.NewClient(opts,
		rueidisotel.WithTracerProvider(ot.MaybeTraceProvider(ot.redisTracing)),
		rueidisotel.WithMeterProvider(ot.mp),
	)
}

func constructTracingExporter(ctx context.Context, otlpEndpoint, otlpCaCertificateFile, otlpTokenSecretFile string) (tracesdk.SpanExporter, error) {
	u, err := url.Parse(otlpEndpoint)
	if err != nil {
		return nil, fmt.Errorf("parsing tracing url %s failed: %w", otlpEndpoint, err)
	}

	var otlpOptions []otlptracehttp.Option

	switch u.Scheme {
	case "https":
	case "http":
		otlpOptions = append(otlpOptions, otlptracehttp.WithInsecure())
	default:
		return nil, fmt.Errorf("unsupported schema of tracing url %q, only `http` and `https` are permitted", u.Scheme)
	}

	otlpOptions = append(otlpOptions, otlptracehttp.WithEndpoint(u.Host))
	otlpOptions = append(otlpOptions, otlptracehttp.WithURLPath(u.Path))

	if otlpTokenSecretFile != "" {
		token, err := os.ReadFile(otlpTokenSecretFile) //nolint: gosec, govet
		if err != nil {
			return nil, fmt.Errorf("unable to read OTLP token from %q: %w", otlpTokenSecretFile, err)
		}
		token = bytes.TrimSpace(token)
		headers := map[string]string{
			"Private-Token": string(token),
		}

		otlpOptions = append(otlpOptions, otlptracehttp.WithHeaders(headers))
	}

	tlsConfig, err := tlstool.ClientConfigWithCACert(otlpCaCertificateFile)
	if err != nil {
		return nil, err
	}
	otlpOptions = append(otlpOptions, otlptracehttp.WithTLSClientConfig(tlsConfig))

	return otlptracehttp.New(ctx, otlpOptions...)
}

func (o *options) constructOTELMeterProvider(r *resource.Resource, reg prometheus.Registerer) (*metricsdk.MeterProvider, func() error, error) {
	otelPromExp, err := promexp.New(promexp.WithRegisterer(reg))
	if err != nil {
		return nil, nil, err
	}
	mp := metricsdk.NewMeterProvider(metricsdk.WithReader(otelPromExp), metricsdk.WithResource(r))
	mpStop := func() error {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return mp.Shutdown(ctx)
	}
	return mp, mpStop, nil
}

func (o *options) constructOTELTracingTools(ctx context.Context, r *resource.Resource) (trace.TracerProvider, propagation.TextMapPropagator, func() error, error) {
	if !o.isTracingEnabled() {
		return noop.NewTracerProvider(), propagation.NewCompositeTextMapPropagator(), func() error { return nil }, nil
	}

	// Exporter must be constructed right before TracerProvider as it's started implicitly so needs to be stopped,
	// which TracerProvider does in its Shutdown() method.
	exporter, err := constructTracingExporter(
		ctx,
		o.configuration.Observability.Tracing.GetOtlpEndpoint(),
		o.configuration.Observability.Tracing.GetOtlpCaCertificateFile(),
		o.configuration.Observability.Tracing.GetOtlpTokenSecretFile(),
	)
	if err != nil {
		return nil, nil, nil, err
	}

	tp := tracesdk.NewTracerProvider(
		tracesdk.WithResource(r),
		tracesdk.WithBatcher(exporter),
	)
	p := propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})
	tpStop := func() error { //nolint:contextcheck
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return tp.Shutdown(ctx)
	}
	return tp, p, tpStop, nil
}

func (o *options) isTracingEnabled() bool {
	return o.configuration.Observability.Tracing != nil
}

func (o *options) constructObservabilityTools(ctx context.Context) (ot *obsTools, stop func() error, retErr error) {
	// Metrics
	reg := prometheus.NewPedanticRegistry()
	goCollector := collectors.NewGoCollector()
	procCollector := collectors.NewProcessCollector(collectors.ProcessCollectorOpts{})
	srvProm := grpc_prometheus.NewServerMetrics()
	clientProm := grpc_prometheus.NewClientMetrics()
	err := metric.Register(reg, goCollector, procCollector, srvProm, clientProm)
	if err != nil {
		return nil, nil, err
	}

	// OTEL resource
	r, err := constructOTELResource()
	if err != nil {
		return nil, nil, err
	}

	// OTEL metrics
	mp, mpStop, err := o.constructOTELMeterProvider(r, reg) //nolint: contextcheck
	if err != nil {
		return nil, nil, err
	}
	defer callOnError(mpStop, &retErr)
	dm := mp.Meter(kasMeterName)
	ssh, err := grpctool.NewServerRequestsInFlightStatsHandler(dm)
	if err != nil {
		return nil, nil, err
	}
	csh, err := grpctool.NewClientRequestsInFlightStatsHandler(dm)
	if err != nil {
		return nil, nil, err
	}
	err = gitlabBuildInfoGauge(dm) //nolint: contextcheck
	if err != nil {
		return nil, nil, err
	}

	// OTEL Tracing
	tp, p, tpStop, err := o.constructOTELTracingTools(ctx, r)
	if err != nil {
		return nil, nil, err
	}
	defer callOnError(tpStop, &retErr)

	// may be nil
	tracing := o.configuration.Observability.Tracing

	return &obsTools{
			probeRegistry:     observability.NewProbeRegistry(),
			reg:               reg,
			tracer:            tp.Tracer(kasTracerName),
			meter:             dm,
			tp:                tp,
			mp:                mp,
			p:                 p,
			ssh:               ssh,
			csh:               csh,
			streamProm:        srvProm.StreamServerInterceptor(),
			unaryProm:         srvProm.UnaryServerInterceptor(),
			streamClientProm:  clientProm.StreamClientInterceptor(),
			unaryClientProm:   clientProm.UnaryClientInterceptor(),
			httpServerTracing: tracing.GetHttpServerEnabled(),
			httpClientTracing: tracing.GetHttpClientEnabled(),
			grpcServerTracing: tracing.GetGrpcServerEnabled(),
			grpcClientTracing: tracing.GetGrpcClientEnabled(),
			redisTracing:      tracing.GetRedisEnabled(),
		}, func() error {
			return errors.Join(
				mpStop(),
				tpStop(),
			)
		}, nil
}

func NewCommand() *cobra.Command {
	o := newOptions()
	c := &cobra.Command{
		Use:   "kas",
		Short: "GitLab Kubernetes Agent Server",
		Args:  cobra.NoArgs,
		Run: func(c *cobra.Command, args []string) {
			ctx := c.Context()
			cmd.LogAndExitOnError(o.log, o.complete(ctx))
			o.setGlobals()
			o.startGoroutineDump(ctx)
			cmd.LogAndExitOnError(o.log, o.run(ctx))
		},
		SilenceErrors: true, // Don't want Cobra to print anything, we handle error printing/logging ourselves.
		SilenceUsage:  true,
	}
	c.Flags().StringVar(&o.configurationFile, "configuration-file", o.configurationFile, "Configuration file to use (YAML)")
	cobra.CheckErr(c.MarkFlagRequired("configuration-file"))

	return c
}

func loggerFromConfig(loggingCfg *kascfg.LoggingCF) (*slog.Logger, *grpctool.Logger, error) {
	var level, grpcLevel slog.Level
	err := level.UnmarshalText([]byte(loggingCfg.Level.String()))
	if err != nil {
		return nil, nil, err
	}
	err = grpcLevel.UnmarshalText([]byte(loggingCfg.GrpcLevel.String()))
	if err != nil {
		return nil, nil, err
	}
	lockedWriter := &logz.LockedWriter{
		Writer: os.Stderr,
	}
	log := slog.New(slog.NewJSONHandler(lockedWriter, &slog.HandlerOptions{
		Level: level,
	}))
	grpcLog := &grpctool.Logger{
		Handler: slog.NewJSONHandler(lockedWriter, &slog.HandlerOptions{
			Level: grpcLevel,
		}),
	}

	return log, grpcLog, nil
}

func constructSentryHub(dsn, environment string) (*sentry.Hub, error) {
	dialer := net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}
	sentryClient, err := sentry.NewClient(sentry.ClientOptions{
		Dsn:         dsn, // empty DSN disables Sentry transport
		SampleRate:  1,   // no sampling
		Release:     cmd.Version,
		Environment: environment,
		HTTPTransport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           dialer.DialContext,
			TLSClientConfig:       tlstool.ClientConfig(),
			MaxIdleConns:          10,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ResponseHeaderTimeout: 20 * time.Second,
			ForceAttemptHTTP2:     true,
		},
	})
	if err != nil {
		return nil, err
	}
	scope := sentry.NewScope()
	scope.SetTag("git_ref", cmd.GitRef)
	return sentry.NewHub(sentryClient, scope), nil
}

type otelErrorHandler struct {
	log *slog.Logger
	hub *sentry.Hub
}

func (h *otelErrorHandler) Handle(err error) {
	h.log.LogAttrs(context.Background(), slog.LevelError, "OpenTelemetry error", logz.Error(err))
	capture(context.Background(), h.hub, "", modshared.NoAgentKey, "OpenTelemetry error", err, nil)
}

type obsTools struct {
	probeRegistry    *observability.ProbeRegistry
	reg              *prometheus.Registry
	tracer           trace.Tracer
	meter            otelmetric.Meter
	tp               trace.TracerProvider
	mp               otelmetric.MeterProvider
	p                propagation.TextMapPropagator
	ssh              stats.Handler
	csh              stats.Handler
	streamProm       grpc.StreamServerInterceptor
	unaryProm        grpc.UnaryServerInterceptor
	streamClientProm grpc.StreamClientInterceptor
	unaryClientProm  grpc.UnaryClientInterceptor

	// Tracing enabled flags
	httpServerTracing bool
	httpClientTracing bool
	grpcServerTracing bool
	grpcClientTracing bool
	redisTracing      bool
}

func (ot *obsTools) MaybeTraceProvider(flag bool) trace.TracerProvider {
	if flag {
		return ot.tp
	}
	return noop.NewTracerProvider()
}

func callOnError(f func() error, err *error) {
	if *err != nil {
		_ = f()
	}
}

func constructRedisReadinessProbe(redisClient rueidis.Client) observability.Probe {
	return func(ctx context.Context) error {
		pingCmd := redisClient.B().Ping().Build()
		err := redisClient.Do(ctx, pingCmd).Error()
		if err != nil {
			return fmt.Errorf("redis: %w", err)
		}
		return nil
	}
}

func constructOTELResource() (*resource.Resource, error) {
	return resource.Merge(
		resource.Default(),
		resource.NewWithAttributes(
			semconv.SchemaURL,
			semconv.ServiceName(kasName),
			semconv.ServiceVersion(cmd.Version),
		),
	)
}

type connectedAgentkInfoMarshaler agent_tracker.ConnectedAgentkInfo

func (i *connectedAgentkInfoMarshaler) MarshalJSON() ([]byte, error) {
	return protojson.Marshal((*agent_tracker.ConnectedAgentkInfo)(i))
}

type connectedAgentwInfoMarshaler agent_tracker.ConnectedAgentwInfo

func (i *connectedAgentwInfoMarshaler) MarshalJSON() ([]byte, error) {
	return protojson.Marshal((*agent_tracker.ConnectedAgentwInfo)(i))
}

type redisAgentInfoResolver struct {
	agentQuerier agent_tracker.Querier
}

func (r redisAgentInfoResolver) Get(ctx context.Context, agentKey api.AgentKey) (map[string]any, error) {
	var infos []any

	//nolint:exhaustive
	switch agentKey.Type {
	case api.AgentTypeKubernetes:
		for info := range r.agentQuerier.GetAgentkConnectionsByID(ctx, agentKey.ID) {
			infos = append(infos, (*connectedAgentkInfoMarshaler)(info))
		}
	case api.AgentTypeWorkspace:
		for info := range r.agentQuerier.GetAgentwConnectionsByID(ctx, agentKey.ID) {
			infos = append(infos, (*connectedAgentwInfoMarshaler)(info))
		}
	default:
		return nil, errors.New("unknown agent type")
	}

	return map[string]any{
		"agents": infos,
	}, nil
}

func constructAgentInfoResolver(agentQuerier agent_tracker.Querier) modserver.AgentInfoResolver {
	return redisAgentInfoResolver{agentQuerier: agentQuerier}
}

type agentFinder struct {
	agentQuerier agent_tracker.Querier
}

func (a *agentFinder) AgentLastConnected(ctx context.Context, agentKey api.AgentKey) (time.Time, error) {
	var lastConnectedAt time.Time

	//nolint:exhaustive
	switch agentKey.Type {
	case api.AgentTypeKubernetes:
		for info := range a.agentQuerier.GetAgentkConnectionsByID(ctx, agentKey.ID) {
			t := info.ConnectedAt.AsTime()
			if t.After(lastConnectedAt) {
				lastConnectedAt = t
			}
		}
	case api.AgentTypeWorkspace:
		for info := range a.agentQuerier.GetAgentwConnectionsByID(ctx, agentKey.ID) {
			t := info.ConnectedAt.AsTime()
			if t.After(lastConnectedAt) {
				lastConnectedAt = t
			}
		}
	default:
		return time.Time{}, errors.New("unknown agent type")
	}

	return lastConnectedAt, nil
}

func constructAgentFinder(agentQuerier agent_tracker.Querier) *agentFinder {
	return &agentFinder{agentQuerier: agentQuerier}
}

func gitlabBuildInfoGauge(m otelmetric.Meter) error {
	g, err := m.Int64Gauge(gitlabBuildInfoGaugeMetricName, otelmetric.WithDescription("Current build info for this GitLab Service"))
	if err != nil {
		return err
	}
	g.Record(context.Background(), 1, otelmetric.WithAttributeSet(attribute.NewSet(
		kasVersionAttr.String(cmd.Version),
		kasGitRefAttr.String(cmd.GitRef),
	)))
	return nil
}

func kasServerName() string {
	return fmt.Sprintf("%s/%s/%s", kasName, cmd.Version, cmd.GitRef)
}

func kasServerNameField() slog.Attr {
	return slog.String("kas", kasServerName())
}

var (
	_ redistool.RPCAPI = (*tokenLimiterAPI)(nil)
)

type tokenLimiterAPI struct {
	rpcAPI modserver.AgentRPCAPI
}

func (a *tokenLimiterAPI) Log() *slog.Logger {
	return a.rpcAPI.Log()
}

func (a *tokenLimiterAPI) HandleProcessingError(msg string, err error) {
	a.rpcAPI.HandleProcessingError(a.rpcAPI.Log(), msg, err)
}

func (a *tokenLimiterAPI) RequestKey() []byte {
	return api.AgentToken2key(a.rpcAPI.AgentTokenWithType().Token)
}

var (
	_ grpctool.ServerErrorReporter = (*serverErrorReporter)(nil)
)

// serverErrorReporter implements the grpctool.ServerErrorReporter interface
// in order to report unknown grpc status code errors.
// In this case the errz.ErrReporter is used as a proxy for the modshared.RPCAPI,
// which logs and captures errors in Sentry.
type serverErrorReporter struct {
	log         *slog.Logger
	errReporter errz.ErrReporter
}

func (r *serverErrorReporter) Report(ctx context.Context, fullMethod string, err error) {
	r.errReporter.HandleProcessingError(ctx, r.log, fmt.Sprintf("Unknown gRPC error in %q", fullMethod), err)
}
