package server

import (
	"context"
	"crypto/ed25519"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"strconv"
	"time"

	"buf.build/go/protovalidate"
	"github.com/golang-jwt/jwt/v5"
	"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/modserver"
	"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/syncz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/tlstool"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/info"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/rpc"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tool/retry"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tunclient"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	otelmetric "go.opentelemetry.io/otel/metric"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"
	"go.opentelemetry.io/otel/trace/noop"
	"google.golang.org/grpc"
	"google.golang.org/grpc/backoff"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/encoding/gzip"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/metadata"
)

var (
	_ syncz.Worker = (*connWorker)(nil)
)

type AgentConnPool interface {
	Add(agentKey api.AgentKey, conn grpc.ClientConnInterface) error
	Remove(agentKey api.AgentKey)
}

type connWorker struct {
	log                  *slog.Logger
	api                  modserver.API
	minIdleConnections   int32
	maxConnections       int32
	scaleUpStep          int32
	maxIdleTime          time.Duration
	userAgent            string
	jwtAudience          string
	jwtIssuer            string
	connectionPollConfig retry.PollConfig
	apiDescriptor        *info.APIDescriptor
	agent                *gapi.ReceptiveAgent
	agentConnPool        AgentConnPool
	agentServerConn      grpc.ClientConnInterface
	tp                   trace.TracerProvider
	mp                   otelmetric.MeterProvider
	p                    propagation.TextMapPropagator
	streamClientProm     grpc.StreamClientInterceptor
	unaryClientProm      grpc.UnaryClientInterceptor
	validator            protovalidate.Validator
	grpcClientTracing    bool
}

func (w *connWorker) Run(ctx context.Context) {
	conn, err := w.constructAgentConnection()
	// NOTE: Receptive agents are currently only agentk.
	//		 If we add other agents, then construct the agentKey based on the type of the agent and its ID.
	agentKey := api.AgentKey{ID: w.agent.Id, Type: api.AgentTypeKubernetes}
	if err != nil {
		w.api.HandleProcessingError(ctx, w.log, "Failed to construct agent connection", err, logz.AgentKey(agentKey))
		return
	}
	defer func() {
		err = conn.Close()
		if err != nil {
			w.api.HandleProcessingError(ctx, w.log, "Failed to close agent connection", err, logz.AgentKey(agentKey))
		}
	}()
	err = w.agentConnPool.Add(agentKey, conn)
	if err != nil {
		w.api.HandleProcessingError(ctx, w.log, "Failed to add agent connection to the pool", err, logz.AgentKey(agentKey))
		return
	}
	defer w.agentConnPool.Remove(agentKey)

	client := rpc.NewReverseTunnelClient(conn)
	cm := tunclient.NewConnectionManager(
		w.minIdleConnections,
		w.maxConnections,
		w.scaleUpStep,
		w.maxIdleTime,
		func(onActive, onIdle func(tunclient.ConnectionInterface)) tunclient.ConnectionInterface {
			return &tunclient.Connection{
				Log:             w.log,
				Descriptor:      w.apiDescriptor,
				Client:          client,
				OwnServerConn:   w.agentServerConn,
				PollConfig:      w.connectionPollConfig,
				OnActive:        onActive,
				OnIdle:          onIdle,
				PrepareMetadata: w.prepareMetadata,
			}
		},
	)
	cm.Run(ctx)
}

func (w *connWorker) prepareMetadata(md metadata.MD) metadata.MD {
	return prepareMetadata(api.AgentKey{ID: w.agent.Id, Type: api.AgentTypeKubernetes}, md)
}

func prepareMetadata(agentKey api.AgentKey, md metadata.MD) metadata.MD {
	if md == nil {
		md = metadata.MD{}
	}
	// MUST override, not append so that the value cannot be spoofed!
	md[agentIDMDKey] = []string{strconv.FormatInt(agentKey.ID, 10)}
	md[agentTypeMDKey] = []string{agentKey.Type.String()}
	return md
}

func (w *connWorker) constructAgentConnection() (*grpc.ClientConn, error) {
	u, err := url.Parse(w.agent.Url)
	if err != nil {
		return nil, fmt.Errorf("invalid agent URL: %w", err)
	}
	var creds credentials.TransportCredentials
	switch u.Scheme {
	case "grpc":
		creds = insecure.NewCredentials()
	case "grpcs":
		tlsConfig, err := w.tlsConfigForAgent() //nolint: govet
		if err != nil {
			return nil, err
		}
		creds = credentials.NewTLS(tlsConfig)
	default:
		return nil, fmt.Errorf("unsupported scheme in agent URL: %q", u.Scheme)
	}

	var tp trace.TracerProvider
	if w.grpcClientTracing {
		tp = w.tp
	} else {
		tp = noop.NewTracerProvider()
	}
	opts := []grpc.DialOption{
		grpc.WithTransportCredentials(creds),
		grpc.WithStatsHandler(otelgrpc.NewClientHandler(
			otelgrpc.WithTracerProvider(tp),
			otelgrpc.WithMeterProvider(w.mp),
			otelgrpc.WithPropagators(w.p),
			otelgrpc.WithMessageEvents(otelgrpc.ReceivedEvents, otelgrpc.SentEvents),
		)),
		// 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.WithSharedWriteBuffer(true),
		grpc.WithDefaultCallOptions(grpc.UseCompressor(gzip.Name), grpc.MaxCallRecvMsgSize(api.GRPCMaxMessageSize)),
		grpc.WithUserAgent(w.userAgent),
		// keepalive.ClientParameters must be specified at least as large as what is allowed by the
		// server-side grpc.KeepaliveEnforcementPolicy
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			// kas allows min 20 seconds, trying to stay below 60 seconds (typical load-balancer timeout) and
			// above kas' server keepalive Time so that kas pings the client sometimes. This helps mitigate
			// reverse-proxies' enforced server response timeout.
			Time:                55 * time.Second,
			PermitWithoutStream: true,
		}),
		grpc.WithChainStreamInterceptor(
			w.streamClientProm,
			grpctool.StreamClientValidatingInterceptor(w.validator),
		),
		grpc.WithChainUnaryInterceptor(
			w.unaryClientProm,
			grpctool.UnaryClientValidatingInterceptor(w.validator),
		),
		// See https://github.com/grpc/grpc/blob/master/doc/service_config.md.
		// See https://github.com/grpc/grpc/blob/master/doc/load-balancing.md.
		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig":[{"round_robin":{}}]}`),
	}
	if j := w.agent.GetJwt(); j != nil {
		jwtSecret, decodeErr := base64.StdEncoding.DecodeString(j.PrivateKey)
		if decodeErr != nil {
			return nil, fmt.Errorf("decode agent JWT private key: %w", decodeErr)
		}
		if len(jwtSecret) != ed25519.SeedSize {
			return nil, fmt.Errorf("invalid agent JWT private key length: %d", len(jwtSecret))
		}
		opts = append(opts, grpc.WithPerRPCCredentials(&grpctool.JWTCredentials{
			SigningMethod: jwt.SigningMethodEdDSA,
			SigningKey:    ed25519.NewKeyFromSeed(jwtSecret),
			Audience:      w.jwtAudience,
			Issuer:        w.jwtIssuer,
			Insecure:      true, // We may or may not have TLS setup, so always say creds don't need TLS.
		}))
	}
	// See https://github.com/grpc/grpc/blob/master/doc/naming.md.
	addressToDial := "dns:" + grpctool.HostWithPort(u)
	conn, err := grpc.NewClient(addressToDial, opts...)
	if err != nil {
		return nil, fmt.Errorf("gRPC.dial: %w", err)
	}
	return conn, nil
}

func (w *connWorker) tlsConfigForAgent() (*tls.Config, error) {
	tlsConfig := tlstool.ClientConfig()
	if w.agent.CaCert != "" {
		certPool, err := x509.SystemCertPool()
		if err != nil {
			return nil, fmt.Errorf("SystemCertPool: %w", err)
		}
		ok := certPool.AppendCertsFromPEM([]byte(w.agent.CaCert))
		if !ok {
			return nil, errors.New("AppendCertsFromPEM() failed")
		}
		tlsConfig.RootCAs = certPool
	}
	if m := w.agent.GetMtls(); m != nil {
		cert, err := tls.X509KeyPair([]byte(m.ClientCert), []byte(m.ClientKey))
		if err != nil {
			return nil, fmt.Errorf("X509KeyPair: %w", err)
		}
		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	tlsConfig.ServerName = w.agent.TlsHost
	return tlsConfig, nil
}
