package kas

import (
	"context"
	"crypto/tls"
	"log/slog"
	"net"
	"sync"
	"time"

	"buf.build/go/protovalidate"
	"github.com/ash2k/stager"
	"github.com/coder/websocket"
	"github.com/redis/rueidis"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	agent2kas_tunnel_router "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent2kas_tunnel/router"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modserver"
	"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/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/redistool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/tlstool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/wstunnel"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/kascfg"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/rpc"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tool/grpcz"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel/propagation"
	"google.golang.org/grpc"
	"google.golang.org/grpc/backoff"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
)

// agentServer is the agent API server. It handles requests from agentk.
// Note that while both the network listening and in-memory servers are constructed here, only the former one
// is started and shut down here. The in-memory server is operated by other code.
type agentServer struct {
	log            *slog.Logger
	listenCfg      *kascfg.ListenAgentCF
	listenMetric   *nettool.ListenerMetrics
	tlsConfig      *tls.Config
	tunnelRegistry *agent2kas_tunnel_router.Registry
	ready          func()

	server grpctool.GRPCServer

	listenServer    *grpc.Server
	listenAuxCancel context.CancelFunc

	inMemServer    *grpc.Server
	inMemListener  net.Listener
	inMemConn      *grpc.ClientConn
	inMemAuxCancel context.CancelFunc
}

func newAgentServer(log *slog.Logger, ot *obsTools, cfg *kascfg.ConfigurationFile, srvAPI modserver.API,
	redisClient rueidis.Client, inMemFactory, factory modserver.AgentRPCAPIFactory, v protovalidate.Validator,
	ownPrivateAPIURL nettool.MultiURL, grpcServerErrorReporter grpctool.ServerErrorReporter, listenMetric *nettool.ListenerMetrics) (*agentServer, error) {
	listenCfg := cfg.Agent.Listen
	tlsConfig, err := tlstool.MaybeServerConfig(listenCfg.CertificateFile, listenCfg.KeyFile)
	if err != nil {
		return nil, err
	}
	// Tunnel registry
	tracker, err := agent2kas_tunnel_router.NewRedisTracker(redisClient, v, cfg.Redis.KeyPrefix+":tunnel_tracker2", ownPrivateAPIURL, ot.meter)
	if err != nil {
		return nil, err
	}
	tunnelRegistry, err := agent2kas_tunnel_router.NewRegistry(
		log,
		srvAPI,
		ot.tracer,
		ot.meter,
		cfg.Agent.RedisConnInfoRefresh.AsDuration(),
		cfg.Agent.RedisConnInfoGc.AsDuration(),
		cfg.Agent.RedisConnInfoTtl.AsDuration(),
		tracker,
	)
	if err != nil {
		return nil, err
	}

	agentConnectionLimiter, err := metric.NewAllowLimiterInstrumentation(
		"agent_connection",
		float64(listenCfg.ConnectionsPerTokenPerMinute),
		"{connection/token/m}",
		ot.tracer,
		ot.meter,
		redistool.NewTokenLimiter(
			redisClient,
			cfg.Redis.KeyPrefix+":agent_limit",
			uint64(listenCfg.ConnectionsPerTokenPerMinute),
			func(ctx context.Context) redistool.RPCAPI {
				return &tokenLimiterAPI{
					rpcAPI: modserver.AgentRPCAPIFromContext(ctx),
				}
			},
		),
	)
	if err != nil {
		return nil, err
	}

	// In-memory gRPC client->listener pipe
	listener := grpcz.NewPipeDialListener()

	inMemConn, err := newInMemConn(ot, v, listener.DialContext)
	if err != nil {
		return nil, err
	}

	traceContextProp := propagation.TraceContext{} // only want trace id, not baggage from external clients/agents
	sharedOpts := []grpc.ServerOption{
		grpc.StatsHandler(otelgrpc.NewServerHandler(
			otelgrpc.WithTracerProvider(ot.MaybeTraceProvider(ot.grpcServerTracing)),
			otelgrpc.WithMeterProvider(ot.mp),
			otelgrpc.WithPropagators(traceContextProp),
			otelgrpc.WithMessageEvents(otelgrpc.ReceivedEvents, otelgrpc.SentEvents),
		)),
		grpc.StatsHandler(ot.ssh),
		grpc.SharedWriteBuffer(true),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
			MinTime:             20 * time.Second,
			PermitWithoutStream: true,
		}),
		grpc.MaxRecvMsgSize(api.GRPCMaxMessageSize),
		grpc.ForceServerCodecV2(rpc.RawCodecWithProtoFallback{}),
	}

	perServerOpts := func(auxCtx context.Context, rpcAPIFactory modserver.AgentRPCAPIFactory) []grpc.ServerOption {
		keepaliveOpt, sh := grpcz.MaxConnectionAge2GRPCKeepalive(auxCtx, listenCfg.MaxConnectionAge.AsDuration())
		return []grpc.ServerOption{
			keepaliveOpt,
			grpc.StatsHandler(sh),
			grpc.ChainStreamInterceptor(
				ot.streamProm, // 1. measure all invocations
				modserver.StreamAgentRPCAPIInterceptor(rpcAPIFactory),                  // 2. inject RPC API
				grpctool.StreamServerValidatingInterceptor(v),                          // x. wrap with validator
				grpctool.StreamServerLimitingInterceptor(agentConnectionLimiter),       //nolint:contextcheck
				grpctool.StreamServerErrorReporterInterceptor(grpcServerErrorReporter), //nolint:contextcheck
			),
			grpc.ChainUnaryInterceptor(
				ot.unaryProm, // 1. measure all invocations
				modserver.UnaryAgentRPCAPIInterceptor(rpcAPIFactory), // 2. inject RPC API
				grpctool.UnaryServerValidatingInterceptor(v),         // x. wrap with validator
				grpctool.UnaryServerLimitingInterceptor(agentConnectionLimiter),
				grpctool.UnaryServerErrorReporterInterceptor(grpcServerErrorReporter),
			),
		}
	}

	inMemAuxCtx, inMemAuxCancel := context.WithCancel(context.Background())
	inMemSrv := grpc.NewServer(append(sharedOpts, perServerOpts(inMemAuxCtx, inMemFactory)...)...)

	listenSrvOpts := sharedOpts
	if !*listenCfg.Websocket && tlsConfig != nil {
		// If we are listening for WebSocket connections, gRPC server doesn't need TLS as it's handled by the
		// HTTP/WebSocket server. Otherwise, we handle it here (if configured).
		listenSrvOpts = append(listenSrvOpts, grpc.Creds(credentials.NewTLS(tlsConfig)))
	}

	listenAuxCtx, listenAuxCancel := context.WithCancel(context.Background())
	listenSrv := grpc.NewServer(append(listenSrvOpts, perServerOpts(listenAuxCtx, factory)...)...)
	return &agentServer{
		log:            log,
		listenCfg:      listenCfg,
		listenMetric:   listenMetric,
		tlsConfig:      tlsConfig,
		tunnelRegistry: tunnelRegistry,
		ready:          ot.probeRegistry.RegisterReadinessToggle("agentServer"),

		server: grpctool.AggregateServer{listenSrv, inMemSrv},

		listenServer:    listenSrv,
		listenAuxCancel: listenAuxCancel,

		inMemServer:    inMemSrv,
		inMemListener:  listener,
		inMemConn:      inMemConn,
		inMemAuxCancel: inMemAuxCancel,
	}, nil
}

func (s *agentServer) Close() error {
	// Closing these here is as a safety net for unfinished startup.
	// Normally they are closed somewhere else.
	// They may return errors when closed the second time, so we ignore them.
	_ = s.inMemConn.Close()     // first close the client
	_ = s.inMemListener.Close() // then close the listener
	return nil
}

func (s *agentServer) Start(stage stager.Stage) {
	var wg sync.WaitGroup
	graceDuration := s.listenCfg.ListenGracePeriod.AsDuration()
	registryCtx, registryCancel := context.WithCancel(context.Background())
	stage.Go(func(ctx context.Context) error {
		return s.tunnelRegistry.Run(registryCtx) // use a separate ctx to stop when the server starts stopping
	})
	stage.GoWhenDone(func() error {
		// We first want gRPC servers to send GOAWAY and only then return from the RPC handlers.
		// So we delay signaling the handlers and registry.
		// See https://github.com/grpc/grpc-go/issues/6830 for more background.
		time.Sleep(graceDuration + time.Second)
		s.listenAuxCancel() // Signal running RPC handlers to stop.
		wg.Wait()           // Wait for the server to stop.
		// Now that the agent server has stopped, there can be no new tunnels. We can stop the registry.
		registryCancel()
		return nil
	})
	wg.Add(1)
	grpctool.StartServer(stage, s.listenServer,
		func(ctx context.Context) (net.Listener, error) {
			var lis net.Listener
			var err error
			maxConnAge := s.listenCfg.MaxConnectionAge.AsDuration()
			isTLSEnabled := s.tlsConfig != nil
			if *s.listenCfg.Websocket { // Explicitly handle TLS for a WebSocket server
				if isTLSEnabled {
					s.tlsConfig.NextProtos = []string{httpz.TLSNextProtoH2, httpz.TLSNextProtoH1} // h2 for gRPC, http/1.1 for WebSocket
					lis, err = nettool.TLSListenWithOSTCPKeepAlive(ctx, *s.listenCfg.Network, s.listenCfg.Address, s.tlsConfig)
				} else {
					lis, err = nettool.ListenWithOSTCPKeepAlive(ctx, *s.listenCfg.Network, s.listenCfg.Address)
				}
				if err != nil {
					return nil, err
				}
				lisWrapped, wrapErr := s.listenMetric.Wrap(lis, "agent_api", maxConnAge, isTLSEnabled)
				if wrapErr != nil {
					_ = lis.Close()
					return nil, wrapErr
				}
				wsWrapper := wstunnel.ListenerWrapper{
					AcceptOptions: websocket.AcceptOptions{
						CompressionMode: websocket.CompressionDisabled,
					},
					// TODO set timeouts
					ReadLimit:  api.WebSocketMaxMessageSize,
					ServerName: kasServerName(),
				}
				lis = wsWrapper.Wrap(lisWrapped, isTLSEnabled)
			} else {
				lis, err = nettool.ListenWithOSTCPKeepAlive(ctx, *s.listenCfg.Network, s.listenCfg.Address)
				if err != nil {
					return nil, err
				}
				lisWrapped, err := s.listenMetric.Wrap(lis, "agent_api", maxConnAge, false)
				if err != nil {
					_ = lis.Close()
					return nil, err
				}
				lis = lisWrapped
			}
			addr := lis.Addr()
			s.log.Info("Agentk API endpoint is up",
				logz.NetNetworkFromAddr(addr),
				logz.NetAddressFromAddr(addr),
				logz.IsWebSocket(*s.listenCfg.Websocket),
			)

			s.ready()

			return lis, nil
		},
		func() {
			time.Sleep(graceDuration) // This delays closing the listener.
		},
		wg.Done, // the server has stopped and drained all connections.
	)
}

func newInMemConn(ot *obsTools, v protovalidate.Validator, dialer func(context.Context, string) (net.Conn, error)) (*grpc.ClientConn, error) {
	// Construct in-memory connection to the agent API gRPC server
	return grpc.NewClient("passthrough:agent-server",
		grpc.WithSharedWriteBuffer(true),
		// Default gRPC parameters are good, no need to change them at the moment.
		// Specify them explicitly for discoverability.
		// See https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
		grpc.WithConnectParams(grpc.ConnectParams{
			Backoff:           backoff.DefaultConfig,
			MinConnectTimeout: 20 * time.Second, // matches the default gRPC value.
		}),
		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.WithStatsHandler(ot.csh),
		grpc.WithUserAgent(kasServerName()),
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time:                55 * time.Second,
			PermitWithoutStream: true,
		}),
		grpc.WithChainStreamInterceptor(
			ot.streamClientProm,
			grpctool.StreamClientValidatingInterceptor(v),
		),
		grpc.WithChainUnaryInterceptor(
			ot.unaryClientProm,
			grpctool.UnaryClientValidatingInterceptor(v),
		),
		grpc.WithContextDialer(dialer),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(api.GRPCMaxMessageSize)),
	)
}
