package agentk

import (
	"context"
	"fmt"
	"math/rand/v2"
	"net/url"
	"os"
	"time"

	"github.com/ash2k/stager"
	"github.com/go-logr/logr"
	"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/cmd/agent"
	agent2kas_tunnel_agent "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent2kas_tunnel/agent"
	agent2kas_tunnel_agentk "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent2kas_tunnel/agentk"
	agent_configuration_agent "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_configuration/agent"
	agent_configuration_rpc "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_configuration/rpc"
	agent_registrar_agentk "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_registrar/agentk"
	flux_agent "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/flux/agent"
	gitlab_access_agentk "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/gitlab_access/agentk"
	gitlab_access_rpc "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/gitlab_access/rpc"
	google_profiler_agentk "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/google_profiler/agentk"
	kas2agentk_tunnel_agent "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kas2agentk_tunnel/agent"
	kas2agentk_tunnel_router "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kas2agentk_tunnel/router"
	kubernetes_api_agent "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api/agent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modagentk"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	observability_agent "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/observability/agent"
	observability_agentk "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/observability/agentk"
	remote_development_agent "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/remote_development/agent"
	starboard_vulnerability "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/starboard_vulnerability/agent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/errz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"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/retry"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/syncz"
	"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/entity"
	"go.opentelemetry.io/otel"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	core_v1 "k8s.io/api/core/v1"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/kubernetes/scheme"
	client_core_v1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/record"
	"k8s.io/klog/v2"
	"k8s.io/kubectl/pkg/cmd/util"
)

const (
	defaultLogLevel     agentcfg.LogLevelEnum = 0 // whatever is 0 is the default value
	defaultGRPCLogLevel                       = agentcfg.LogLevelEnum_error

	envVarPodNamespace       = "POD_NAMESPACE"
	envVarPodName            = "POD_NAME"
	envVarPodLabelSelector   = "POD_SELECTOR_LABELS"
	envVarServiceAccountName = "SERVICE_ACCOUNT_NAME"
	envVarAgentToken         = "AGENTK_TOKEN"

	getConfigurationInitBackoff   = 10 * time.Second
	getConfigurationMaxBackoff    = 5 * time.Minute
	getConfigurationResetDuration = 10 * time.Minute
	getConfigurationBackoffFactor = 2.0
	getConfigurationJitter        = 5.0

	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
	routingTryNewAgentkInterval = 10 * time.Millisecond
)

type options struct {
	agent.Options

	agentMeta          *entity.AgentkMeta
	configFlags        *genericclioptions.ConfigFlags
	serviceAccountName string

	// kasAddress specifies the address of kas.
	// This field is used as a flag. If set, agentk is in the agentk->kas communication mode.
	// If not set, kas->agentk mode is activated.

	apiListenNetwork string
	apiListenAddress string
	apiCertFile      string
	apiKeyFile       string
	apiJWTFile       string
	apiCACertFile    string
	apiMTLS          bool

	privateAPIListenNetwork string
	privateAPIListenAddress string
	privateAPICertFile      string
	privateAPIKeyFile       string
	privateAPICACertFile    string
	privateAPIJWTFile       string
}

func newOptions() *options {
	return &options{
		Options: agent.Options{
			AgentKey:                   syncz.NewValueHolder[api.AgentKey](),
			GitLabExternalURL:          syncz.NewValueHolder[url.URL](),
			ObservabilityListenNetwork: agent.DefaultObservabilityListenNetwork,
			ObservabilityListenAddress: agent.DefaultObservabilityListenAddress,
		},
		agentMeta: &entity.AgentkMeta{
			Version:            cmd.Version,
			CommitId:           cmd.GitRef, // For backwards compatibility.
			PodNamespace:       os.Getenv(envVarPodNamespace),
			PodName:            os.Getenv(envVarPodName),
			KubernetesVersion:  &entity.KubernetesVersion{},
			GitRef:             cmd.GitRef,
			ExtraTelemetryData: agent.CollectExtraTelemetryData(os.Environ()),
		},
		configFlags:             genericclioptions.NewConfigFlags(true),
		serviceAccountName:      os.Getenv(envVarServiceAccountName),
		apiListenNetwork:        agent.DefaultAPIListenNetwork,
		privateAPIListenNetwork: agent.DefaultPrivateAPIListenNetwork,
		privateAPIListenAddress: agent.DefaultPrivateAPIListenAddress,
	}
}

func (o *options) validate() error {
	err := o.Validator.Validate(o.agentMeta) // validate the constraints in the proto
	if err != nil {
		return fmt.Errorf("invalid agent metadata. Make sure to set %s and %s environment variables correctly: %w",
			envVarPodNamespace, envVarPodName, err)
	}
	return nil
}

// setGlobals shouldn't exist, yet here we are.
func (o *options) setGlobals() {
	grpclog.SetLoggerV2(&grpctool.Logger{Handler: o.GRPCHandler}) // 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((*metric.OtelErrorHandler)(o.Log))
}

func (o *options) run(ctx context.Context) (retErr error) {
	ot, stop, err := o.ConstructObservabilityTools()
	if err != nil {
		return err
	}
	defer errz.SafeCall(stop, &retErr)

	// Construct Kubernetes tools.
	k8sFactory := util.NewFactory(o.configFlags)
	kubeClient, err := k8sFactory.KubernetesClientSet()
	if err != nil {
		return err
	}

	// Construct event recorder
	eventBroadcaster := record.NewBroadcaster()
	eventRecorder := eventBroadcaster.NewRecorder(scheme.Scheme, core_v1.EventSource{Component: agent.AgentName})

	// Construct leader runner
	lr := agent.NewLeaderRunner[*agentcfg.AgentConfiguration](&leaseLeaderElector{
		namespace: o.agentMeta.PodNamespace,
		name: func(ctx context.Context) (string, error) {
			agentKey, err := o.AgentKey.Get(ctx)
			if err != nil {
				return "", err
			}
			// We use agent id as part of lock name so that agentk Pods of different id don't compete with
			// each other. Only Pods with same agent id should compete for a lock. Put differently, agentk Pods
			// with same agent id have the same lock name but with different id have different lock name.
			return fmt.Sprintf("agent-%d-lock", agentKey.ID), nil
		},
		identity:           o.agentMeta.PodName,
		coordinationClient: kubeClient.CoordinationV1(),
		eventRecorder:      eventRecorder,
	})

	rpcAPIFactory := o.NewRPCAPIFactory()

	// ========= things that depend on the mode of operation of the reverse tunnel: kas->agentk / agentk->kas :: START
	var kasConn grpc.ClientConnInterface
	var apiServer grpctool.GRPCServer
	var apiServerStart func(stager.Stage)
	var privateAPIServerStart func(stager.Stage)
	var registerKASAPI func(*grpc.ServiceDesc)
	var tunnelFactory modagentk.Factory
	var runQuerier func(ctx context.Context) error

	if o.KASAddress == "" {
		var err error
		sharedAPI := &agent.AgentSharedAPI{}

		apiSrv, err := agent.NewListenAPIServer( //nolint:contextcheck
			o.Log, ot, rpcAPIFactory,
			o.apiListenNetwork, o.apiListenAddress,
			o.apiCertFile, o.apiKeyFile,
			o.apiJWTFile, api.JWTKAS, api.JWTAgentk,
			o.apiCACertFile, o.apiMTLS,
			o.Validator,
		)
		if err != nil {
			return fmt.Errorf("API Server: %w", err)
		}

		privateSrv, err := agent.NewPrivateAPIServer( //nolint:contextcheck
			o.Log,
			ot,
			modshared.APIToErrReporter(sharedAPI),
			rpcAPIFactory,
			o.Validator,
			o.agentNameVersion(),
			o.privateAPIListenNetwork, o.privateAPIListenAddress,
			o.privateAPICertFile, o.privateAPIKeyFile, o.privateAPICACertFile,
			o.privateAPIJWTFile,
		)
		if err != nil {
			return fmt.Errorf("private API Server: %w", err)
		}
		defer errz.SafeClose(privateSrv, &retErr)

		podLabelSelector := os.Getenv(envVarPodLabelSelector)
		o.Log.Info("Using Pod label selector to find agentk Pods within the namespace",
			logz.PodNamespace(o.agentMeta.PodNamespace), logz.LabelSelector(podLabelSelector))
		gq := kas2agentk_tunnel_router.NewGatewayURLQuerier(
			kubeClient,
			o.agentMeta.PodNamespace,
			podLabelSelector,
			privateSrv.OwnURLScheme,
			privateSrv.OwnURLPort,
		)
		registry := kas2agentk_tunnel_router.NewRegistry(o.Log)
		plugin := &kas2agentk_tunnel_router.Plugin{
			Registry:         registry,
			AgentkPool:       privateSrv.AgentPool,
			GatewayQuerier:   gq,
			API:              sharedAPI,
			OwnPrivateAPIURL: privateSrv.OwnURL,
			Creds:            grpctool.NewTokenCredentials(o.AgentToken, api.AgentTypeKubernetes, true),
			PollConfig: retry.NewPollConfigFactory(routingAttemptInterval, retry.NewExponentialBackoffFactory(
				routingInitBackoff,
				routingMaxBackoff,
				routingResetDuration,
				routingBackoffFactor,
				routingJitter,
			)),
			TryNewGatewayInterval: routingTryNewAgentkInterval,
			TunnelFindTimeout:     routingTunnelFindTimeout,
		}
		kasConn = &tunserver.RoutingClientConn[grpctool.URLTarget]{
			Log:    o.Log,
			API:    sharedAPI,
			Plugin: plugin,
		}
		apiServer = apiSrv.Server
		apiServerStart = apiSrv.Start
		privateAPIServerStart = privateSrv.Start
		registerKASAPI = (&tunserver.Router[grpctool.URLTarget]{
			Plugin:           plugin,
			PrivateAPIServer: privateSrv.Server,
		}).RegisterTunclientAPI
		tunnelFactory = &kas2agentk_tunnel_agent.Factory{
			Registry: registry,
		}
		runQuerier = gq.Run
	} else {
		var err error
		// Construct gRPC connection to gitlab-kas
		realKASConn, err := o.ConstructKASConnection(ot, o.agentNameVersion(), api.AgentTypeKubernetes)
		if err != nil {
			return err
		}
		defer errz.SafeClose(realKASConn, &retErr)

		// Construct in-mem API gRPC server
		apiSrv, err := agent.NewInMemAPIServer(ot, rpcAPIFactory, o.Validator)
		if err != nil {
			return fmt.Errorf("in-mem API Server: %w", err)
		}
		defer errz.SafeClose(apiSrv, &retErr)

		kasConn = realKASConn
		apiServer = apiSrv.Server
		registerKASAPI = func(desc *grpc.ServiceDesc) {}
		apiServerStart = apiSrv.Start
		privateAPIServerStart = func(stage stager.Stage) {}
		tunnelFactory = &agent2kas_tunnel_agentk.Factory{
			Factory: agent2kas_tunnel_agent.Factory{
				APIServerConn: apiSrv.InMemConn,
			},
		}
		runQuerier = func(ctx context.Context) error {
			return nil
		}
	}
	// ========= things that depend on the mode of operation of the reverse tunnel: kas->agentk / agentk->kas :: END

	// Start events processing pipeline.
	loggingWatch := eventBroadcaster.StartStructuredLogging(0)
	defer loggingWatch.Stop()
	eventBroadcaster.StartRecordingToSink(&client_core_v1.EventSinkImpl{Interface: kubeClient.CoreV1().Events("")})
	defer eventBroadcaster.Shutdown()

	instanceID := rand.Int64() //nolint: gosec
	runner := o.newModuleRunner(kasConn)
	accessClient := gitlab_access_rpc.NewGitlabAccessClient(kasConn)
	configForFactory := func(name string) *modagentk.Config {
		return &modagentk.Config{
			Config: modagent.Config{
				Log:        o.Log.With(logz.ModuleName(name)),
				InstanceID: instanceID,
				API: &agent.AgentAPI{
					ModuleName:        name,
					AgentKey:          o.AgentKey,
					GitLabExternalURL: o.GitLabExternalURL,
					Client:            accessClient,
				},
				KASConn:          kasConn,
				APIServer:        apiServer,
				AgentName:        agent.AgentName,
				AgentNameVersion: o.agentNameVersion(),
				Validator:        o.Validator,
			},
			AgentMeta:          o.agentMeta,
			K8sUtilFactory:     k8sFactory,
			RegisterKASAPI:     registerKASAPI,
			ServiceAccountName: o.serviceAccountName,
		}
	}

	// Start things up. Stages are shut down in reverse order.
	return syncz.RunLayers(ctx,
		syncz.LayerGoFunc(func(ctx context.Context) error {
			// Start leader runner.
			lr.Run(ctx)
			return nil
		}),
		syncz.SiblingLayers{
			// Start querier. Tunnel routing depends on it so it should stop after servers.
			syncz.LayerGoFunc(runQuerier),
			&agent.FactoriesLayer[*modagentk.Config, *agentcfg.AgentConfiguration]{
				Log:    o.Log,
				Config: configForFactory,
				LR:     lr,
				MR:     runner,
				Factories: []modagentk.Factory{
					&agent_configuration_agent.Factory{},
					&gitlab_access_agentk.Factory{},
					&google_profiler_agentk.Factory{},
					&kubernetes_api_agent.Factory{},
					&observability_agentk.Factory{
						Factory: observability_agent.Factory{
							LogLevel:      o.LogLevel,
							GRPCLogLevel:  o.GRPCLogLevel,
							Gatherer:      ot.Reg,
							Registerer:    ot.Reg,
							ListenNetwork: o.ObservabilityListenNetwork,
							ListenAddress: o.ObservabilityListenAddress,
							CertFile:      o.ObservabilityCertFile,
							KeyFile:       o.ObservabilityKeyFile,
						},
						DefaultGRPCLogLevel: defaultGRPCLogLevel,
					},
					&remote_development_agent.Factory{},
					&starboard_vulnerability.Factory{},
				},
			},
		},
		syncz.LayerFunc(func(stage stager.Stage) {
			// Start servers.
			privateAPIServerStart(stage)
			apiServerStart(stage)
		}),
		&agent.FactoriesLayer[*modagentk.Config, *agentcfg.AgentConfiguration]{
			Log:    o.Log,
			Config: configForFactory,
			LR:     lr,
			MR:     runner,
			Factories: []modagentk.Factory{
				// Uses kas connection.
				&agent_registrar_agentk.Factory{},
				// Uses kas connection.
				&flux_agent.Factory{},
				// Uses kas connection.
				tunnelFactory,
			},
		},
		// Start configuration refresh.
		syncz.LayerGoFunc(runner.RunConfigurationRefresh),
	)
}

func (o *options) newModuleRunner(kasConn grpc.ClientConnInterface) *agent.ModuleRunner[agent_configuration_rpc.ConfigurationData, *agentcfg.AgentConfiguration] {
	return &agent.ModuleRunner[agent_configuration_rpc.ConfigurationData, *agentcfg.AgentConfiguration]{
		Log: o.Log,
		Watcher: &agent_configuration_rpc.ConfigurationWatcher{
			Log:    o.Log,
			Client: agent_configuration_rpc.NewAgentConfigurationClient(kasConn),
			PollConfig: retry.NewPollConfigFactory(0, retry.NewExponentialBackoffFactory(
				getConfigurationInitBackoff,
				getConfigurationMaxBackoff,
				getConfigurationResetDuration,
				getConfigurationBackoffFactor,
				getConfigurationJitter,
			)),
			ConfigPreProcessor: func(data agent_configuration_rpc.ConfigurationData) error {
				err := o.AgentKey.Set(api.AgentKey{ID: data.Config.AgentId, Type: api.AgentTypeKubernetes})
				if err != nil {
					return err
				}
				u, err := url.Parse(data.Config.GitlabExternalUrl)
				if err != nil {
					return fmt.Errorf("unable to parse configured GitLab External URL %q: %w", data.Config.GitlabExternalUrl, err)
				}
				return o.GitLabExternalURL.Set(*u)
			},
		},
		Data2Config: func(data agent_configuration_rpc.ConfigurationData) (*agentcfg.AgentConfiguration, []any) {
			return data.Config, []any{logz.CommitID(data.CommitID)}
		},
	}
}

// agentNameVersion return a string to use as User-Agent or Server HTTP header.
func (o *options) agentNameVersion() string {
	return fmt.Sprintf("%s/%s/%s", agent.AgentName, o.agentMeta.Version, o.agentMeta.GitRef)
}

func NewCommand() *cobra.Command {
	o := newOptions()
	c := &cobra.Command{
		Use:   "agentk",
		Short: "GitLab Agent for Kubernetes",
		Args:  cobra.NoArgs,
		Run: func(c *cobra.Command, args []string) {
			cmd.LogAndExitOnError(o.Log, o.Complete(envVarAgentToken, defaultLogLevel.String(), defaultGRPCLogLevel.String()))
			cmd.LogAndExitOnError(o.Log, o.validate())
			o.setGlobals()
			cmd.LogAndExitOnError(o.Log, o.run(c.Context()))
		},
		SilenceErrors: true, // Don't want Cobra to print anything, we handle error printing/logging ourselves.
		SilenceUsage:  true,
	}

	agent.AddCommonFlagsToCommand(c, &o.Options)

	f := c.Flags()

	f.StringVar(&o.apiListenNetwork, "api-listen-network", o.apiListenNetwork, "API network to listen on")
	f.StringVar(&o.apiListenAddress, "api-listen-address", o.apiListenAddress, "API address to listen on")
	f.StringVar(&o.apiCertFile, "api-cert-file", o.apiCertFile, "File with X.509 certificate in PEM format for API endpoint TLS")
	f.StringVar(&o.apiKeyFile, "api-key-file", o.apiKeyFile, "File with X.509 key in PEM format for API endpoint TLS")
	f.StringVar(&o.apiJWTFile, "api-jwt-file", o.apiJWTFile, "Base64-encoded EdDSA public key to validate JWT tokens from kas. Used for API endpoint")
	f.StringVar(&o.apiCACertFile, "api-client-ca-cert-file", o.apiCACertFile, "File with X.509 certificate authority certificate in PEM format. Used for verifying client cert of KAS (GitLab Kubernetes Agent Server). This is used for kas->agentk mTLS")
	f.BoolVar(&o.apiMTLS, "api-mtls", o.apiMTLS, "Used to enabled kas->agentk mTLS. Defaults to true when --api-client-ca-cert-file is specified")
	c.MarkFlagsRequiredTogether("api-cert-file", "api-key-file")
	c.MarkFlagsMutuallyExclusive("api-jwt-file", "api-client-ca-cert-file")

	f.StringVar(&o.privateAPIListenNetwork, "private-api-listen-network", o.privateAPIListenNetwork, "Private API network to listen on")
	f.StringVar(&o.privateAPIListenAddress, "private-api-listen-address", o.privateAPIListenAddress, "Private API address to listen on")
	f.StringVar(&o.privateAPICertFile, "private-api-cert-file", o.privateAPICertFile, "File with X.509 certificate in PEM format for private API endpoint TLS")
	f.StringVar(&o.privateAPIKeyFile, "private-api-key-file", o.privateAPIKeyFile, "File with X.509 key in PEM format for private API endpoint TLS")
	f.StringVar(&o.privateAPICACertFile, "private-api-ca-cert-file", o.privateAPICACertFile, "File with X.509 certificate authority certificate in PEM format for private API endpoint TLS")
	f.StringVar(&o.privateAPIJWTFile, "private-api-jwt-file", o.privateAPIJWTFile, "Base64-encoded secret for JWT token creation and validation for agentk to agentk private API access")
	c.MarkFlagsRequiredTogether("private-api-cert-file", "private-api-key-file")

	o.configFlags.AddFlags(f)
	c.MarkFlagsOneRequired("kas-address", "api-listen-address")
	c.MarkFlagsMutuallyExclusive("kas-address", "api-listen-address")

	return c
}
