package kas

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"time"

	"buf.build/go/protovalidate"
	"github.com/ash2k/stager"
	"github.com/golang-jwt/jwt/v5"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"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/errz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"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/nettool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/tlstool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/kascfg"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"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"
)

type privateAPIServer struct {
	log           *slog.Logger
	listenCfg     *kascfg.ListenPrivateApiCF
	listenMetric  *nettool.ListenerMetrics
	ownURL        nettool.MultiURL
	server        grpctool.GRPCServer
	listenServer  *grpc.Server
	inMemServer   *grpc.Server
	inMemListener net.Listener
	kasPool       grpctool.PoolInterface[nettool.MultiURL]
	kasPoolClose  func()
	auxCancel     context.CancelFunc
	ready         func()
}

func newPrivateAPIServer(log *slog.Logger, ot *obsTools, errRep errz.ErrReporter, cfg *kascfg.ConfigurationFile,
	factory modshared.RPCAPIFactory, grpcServerErrorReporter grpctool.ServerErrorReporter, v protovalidate.Validator,
	listenMetric *nettool.ListenerMetrics) (*privateAPIServer, error) {
	listenCfg := cfg.PrivateApi.Listen
	jwtSecret, err := ioz.LoadSHA256Base64Secret(log, listenCfg.AuthenticationSecretFile)
	if err != nil {
		return nil, fmt.Errorf("auth secret file: %w", err)
	}

	ou := privateAPIServerURL{
		log:            log,
		interfaceAddrs: net.InterfaceAddrs,
	}
	ownURL, err := ou.constructFromEnvironment(
		*listenCfg.Network,
		listenCfg.Address,
		listenCfg.IsTLSEnabled(),
	)
	if err != nil {
		return nil, fmt.Errorf("construct own private API multi URL: %w", err)
	}
	log.Info("Using own private API multi URL", logz.URL(ownURL.String()))

	// In-memory gRPC client->listener pipe
	listener := grpctool.NewDialListener()

	// Client pool
	kasPool, err := newKASPool(log, ot, errRep, jwtSecret, ownURL, listenCfg.CaCertificateFile,
		v, listener.DialContext)
	if err != nil {
		return nil, fmt.Errorf("kas pool: %w", err)
	}

	// Server
	auxCtx, auxCancel := context.WithCancel(context.Background())
	server, inMemServer, err := newPrivateAPIServerImpl(auxCtx, ot, cfg, jwtSecret, factory, grpcServerErrorReporter, v)
	if err != nil {
		auxCancel()
		return nil, fmt.Errorf("new server: %w", err)
	}
	return &privateAPIServer{
		log:           log,
		listenCfg:     listenCfg,
		listenMetric:  listenMetric,
		ownURL:        ownURL,
		server:        grpctool.AggregateServer{server, inMemServer},
		listenServer:  server,
		inMemServer:   inMemServer,
		inMemListener: listener,
		kasPool:       kasPool,
		kasPoolClose:  sync.OnceFunc(func() { kasPool.Shutdown(30 * time.Second) }),
		auxCancel:     auxCancel,
		ready:         ot.probeRegistry.RegisterReadinessToggle("privateAPIServer"),
	}, nil
}

func (s *privateAPIServer) Start(stage stager.Stage) {
	stopInMem := make(chan struct{})
	grpctool.StartServer(stage, s.inMemServer,
		func(ctx context.Context) (net.Listener, error) {
			return s.inMemListener, nil
		},
		func() {
			<-stopInMem
			s.kasPoolClose()
		},
		func() {},
	)
	grpctool.StartServer(stage, s.listenServer,
		func(ctx context.Context) (net.Listener, error) {
			lis, err := nettool.ListenWithOSTCPKeepAlive(ctx, *s.listenCfg.Network, s.listenCfg.Address)
			if err != nil {
				return nil, err
			}
			lisWrapped, err := s.listenMetric.Wrap(lis, "private_api", s.listenCfg.MaxConnectionAge.AsDuration())
			if err != nil {
				_ = lis.Close()
				return nil, err
			}
			addr := lisWrapped.Addr()
			s.log.Info("Private API endpoint is up",
				logz.NetNetworkFromAddr(addr),
				logz.NetAddressFromAddr(addr),
			)
			s.ready()
			return lisWrapped, nil
		},
		func() {
			time.Sleep(s.listenCfg.ListenGracePeriod.AsDuration())
			// We first want gRPC server to send GOAWAY and only then return from the RPC handlers.
			// So we delay signaling the handlers.
			// See https://github.com/grpc/grpc-go/issues/6830 for more background.
			// Start a goroutine in a second and...
			time.AfterFunc(time.Second, func() {
				close(stopInMem) // ... signal the in-memory server to stop.
				s.auxCancel()    // ... signal running RPC handlers to stop.
			})
		},
		func() {},
	)
}

func (s *privateAPIServer) Close() error {
	s.kasPoolClose()               // first close the client (if not closed already)
	return s.inMemListener.Close() // then close the listener (if not closed already)
}

func newPrivateAPIServerImpl(auxCtx context.Context, ot *obsTools, cfg *kascfg.ConfigurationFile,
	jwtSecret []byte, factory modshared.RPCAPIFactory, grpcServerErrorReporter grpctool.ServerErrorReporter,
	v protovalidate.Validator) (*grpc.Server, *grpc.Server, error) {
	listenCfg := cfg.PrivateApi.Listen
	credsOpt, err := grpctool.MaybeTLSCreds(listenCfg.CertificateFile, listenCfg.KeyFile)
	if err != nil {
		return nil, nil, err
	}

	jwtAuther := grpctool.NewHMACJWTAuther(jwtSecret, api.JWTKAS, api.JWTKAS, func(ctx context.Context) *slog.Logger {
		return modshared.RPCAPIFromContext[modserver.RPCAPI](ctx).Log()
	})

	keepaliveOpt, sh := grpctool.MaxConnectionAge2GRPCKeepalive(auxCtx, listenCfg.MaxConnectionAge.AsDuration())
	sharedOpts := []grpc.ServerOption{
		keepaliveOpt,
		grpc.StatsHandler(otelgrpc.NewServerHandler(
			otelgrpc.WithTracerProvider(ot.MaybeTraceProvider(ot.grpcServerTracing)),
			otelgrpc.WithMeterProvider(ot.mp),
			otelgrpc.WithPropagators(ot.p),
			otelgrpc.WithMessageEvents(otelgrpc.ReceivedEvents, otelgrpc.SentEvents),
		)),
		grpc.StatsHandler(ot.ssh),
		grpc.StatsHandler(sh),
		grpc.SharedWriteBuffer(true),
		grpc.ChainStreamInterceptor(
			ot.streamProm, // 1. measure all invocations
			modshared.StreamRPCAPIInterceptor(factory),                             // 2. inject RPC API
			jwtAuther.StreamServerInterceptor,                                      // 3. auth and maybe log
			grpctool.StreamServerValidatingInterceptor(v),                          // x. wrap with validator
			grpctool.StreamServerErrorReporterInterceptor(grpcServerErrorReporter), //nolint:contextcheck
		),
		grpc.ChainUnaryInterceptor(
			ot.unaryProm, // 1. measure all invocations
			modshared.UnaryRPCAPIInterceptor(factory),    // 2. inject RPC API
			jwtAuther.UnaryServerInterceptor,             // 3. auth and maybe log
			grpctool.UnaryServerValidatingInterceptor(v), // x. wrap with validator
			grpctool.UnaryServerErrorReporterInterceptor(grpcServerErrorReporter),
		),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
			MinTime:             20 * time.Second,
			PermitWithoutStream: true,
		}),
		grpc.ForceServerCodecV2(grpctool.RawCodecWithProtoFallback{}),
		grpc.MaxRecvMsgSize(api.GRPCMaxMessageSize),
	}
	server := grpc.NewServer(append(sharedOpts, credsOpt...)...)
	inMemServer := grpc.NewServer(sharedOpts...)
	return server, inMemServer, nil
}

func newKASPool(log *slog.Logger, ot *obsTools, errRep errz.ErrReporter, jwtSecret []byte,
	ownPrivateAPIURL nettool.MultiURL, caCertificateFile string, v protovalidate.Validator,
	dialer func(context.Context, string) (net.Conn, error)) (grpctool.PoolInterface[nettool.MultiURL], error) {

	sharedPoolOpts := []grpc.DialOption{
		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.WithPerRPCCredentials(&grpctool.JWTCredentials{
			SigningMethod: jwt.SigningMethodHS256,
			SigningKey:    jwtSecret,
			Audience:      api.JWTKAS,
			Issuer:        api.JWTKAS,
			Insecure:      true, // We may or may not have TLS setup, so always say creds don't need TLS.
		}),
		grpc.WithChainStreamInterceptor(
			ot.streamClientProm,
			grpctool.StreamClientValidatingInterceptor(v),
		),
		grpc.WithChainUnaryInterceptor(
			ot.unaryClientProm,
			grpctool.UnaryClientValidatingInterceptor(v),
		),
		grpc.WithNoProxy(),
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(api.GRPCMaxMessageSize)),
	}

	// Construct in-memory connection to private API gRPC server
	inMemConn, err := grpc.NewClient("passthrough:private-server",
		append(sharedPoolOpts,
			grpc.WithContextDialer(dialer),
			grpc.WithTransportCredentials(insecure.NewCredentials()),
		)...,
	)
	if err != nil {
		return nil, err
	}
	tlsConfig, err := tlstool.ClientConfigWithCACert(caCertificateFile)
	if err != nil {
		return nil, err
	}
	noServerNameCreds := grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
	insecureCredsOpt := grpc.WithTransportCredentials(insecure.NewCredentials())
	kasPool := grpctool.NewPool(log, errRep, func(targetURL nettool.MultiURL) (string, []grpc.DialOption, error) {
		var credsOpt grpc.DialOption
		switch targetURL.Scheme() {
		case "grpc":
			credsOpt = insecureCredsOpt
		case "grpcs":
			tlsHost := targetURL.TLSHost()
			if tlsHost != "" {
				tlsCopy := tlsConfig.Clone()
				tlsCopy.ServerName = tlsHost
				credsOpt = grpc.WithTransportCredentials(credentials.NewTLS(tlsCopy))
			} else {
				credsOpt = noServerNameCreds
			}
		default:
			return "", nil, fmt.Errorf("unsupported URL scheme in %s", targetURL)
		}
		l := len(sharedPoolOpts)
		muOpts := append(
			sharedPoolOpts[:l:l], // set cap=len to ensure append() will copy the slice
			credsOpt,
		)
		if targetURL.Host() != "" {
			return "dns:" + targetURL.HostPort(), muOpts, nil
		} else {
			muOpts = append(muOpts,
				grpc.WithResolvers(grpctool.NewFixedAddressResolver(targetURL.Addresses(), targetURL.Port())),
			)
			return grpctool.FixedAddressResolverTarget, muOpts, nil
		}
	})
	return grpctool.NewPoolSelf(log, errRep, kasPool, ownPrivateAPIURL, inMemConn), nil
}
