// Copyright The Monogon Project Authors.
// SPDX-License-Identifier: Apache-2.0

// Package component implements reusable bits for cloud service components. Each
// component is currently defined as being a standalone Go binary with its own
// internal gRPC listener. Subsequent listeners (eg. public gRPC or HTTP) can be
// defined by users of this library.
package component

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"flag"
	"net"
	"net/http"
	"os"
	"path/filepath"

	"github.com/adrg/xdg"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"k8s.io/klog/v2"
)

// ComponentConfig is the common configuration of a component. It's
// supposed to be instantiated within a Configuration struct of a component.
//
// It can be configured by flags (via RegisterFlags) or manually (eg. in tests).
type ComponentConfig struct {
	// GRPCKeyPath is the filesystem path of the x509 key used to serve internal
	// gRPC traffic.
	GRPCKeyPath string
	// GRPCCertificatePath is the filesystem path of the x509 certificate used to
	// serve internal gRPC traffic.
	GRPCCertificatePath string
	// GRPCCAPath is the filesystem path of of the x509 CA certificate used to
	// verify incoming connections on internal gRPC traffic.
	GRPCCAPath string
	// GRPCListenAddress is the address on which the component should server
	// internal gRPC traffic.
	GRPCListenAddress string

	// DevCerts, if enabled, automatically generates development CA and component
	// certificates/keys at DevCertsPath, uses these to serve traffic.
	DevCerts bool
	// DevCertsPath sets the prefix in which DevCerts are generated. All components
	// should have the same path set so that they reuse the CA certificate.
	DevCertsPath string

	// ComponentName is the name of this component, which should be [a-z0-9+]. It's
	// used to prefix all flags set by the Configuration.
	ComponentName string

	// PrometheusListenAddress is the address on which the component should serve
	// Prometheus metrics.
	PrometheusListenAddress string
	// PrometheusInsecure enables serving Prometheus metrics without any TLS, running
	// a plain HTTP listener. If disabled, Prometheus metrics are served using the
	// same PKI setup as the components' gRPC server.
	PrometheusInsecure bool

	prometheusRegistry *prometheus.Registry
}

// RegisterFlags registers the component configuration to be provided by flags.
// This must be called exactly once before then calling flags.Parse().
func (c *ComponentConfig) RegisterFlags(componentName string) {
	flag.StringVar(&c.GRPCKeyPath, componentName+"_grpc_key_path", "", "Path to gRPC server/client key for "+componentName)
	flag.StringVar(&c.GRPCCertificatePath, componentName+"_grpc_certificate_path", "", "Path to gRPC server/client certificate for "+componentName)
	flag.StringVar(&c.GRPCCAPath, componentName+"_grpc_ca_certificate_path", "", "Path to gRPC CA certificate for "+componentName)
	flag.StringVar(&c.GRPCListenAddress, componentName+"_grpc_listen_address", ":4242", "Address to listen at for gRPC connections for "+componentName)
	flag.StringVar(&c.PrometheusListenAddress, componentName+"_prometheus_listen_address", ":4243", "Address to listen at for Prometheus connections for "+componentName)
	flag.BoolVar(&c.PrometheusInsecure, componentName+"_prometheus_insecure", false, "Serve plain HTTP prometheus without mTLS. If not set, main gRPC TLS credentials/certificates are used")

	flag.BoolVar(&c.DevCerts, componentName+"_dev_certs", false, "Use developer certificates (autogenerated) for "+componentName)
	flag.StringVar(&c.DevCertsPath, componentName+"_dev_certs_path", filepath.Join(xdg.ConfigHome, "monogon-dev-certs"), "Path for storing developer certificates")

	c.ComponentName = componentName
}

func (c *ComponentConfig) getTLSConfig() *tls.Config {
	var certPath, keyPath, caPath string
	if c.DevCerts {
		// Use devcerts if requested.
		certPath, keyPath, caPath = c.GetDevCerts()
	} else {
		// Otherwise, use data from flags.
		if c.GRPCKeyPath == "" {
			klog.Exitf("-grpc_key_path must be set")
		}
		if c.GRPCCertificatePath == "" {
			klog.Exitf("-grpc_certificate_path must be set")
		}
		if c.GRPCCAPath == "" {
			klog.Exitf("-grpc_ca_certificate_path must be set")
		}
		keyPath = c.GRPCKeyPath
		certPath = c.GRPCCertificatePath
		caPath = c.GRPCCAPath
	}

	ca, err := os.ReadFile(caPath)
	if err != nil {
		klog.Exitf("Could not read GRPC CA: %v", err)
	}
	certPool := x509.NewCertPool()
	if !certPool.AppendCertsFromPEM(ca) {
		klog.Exitf("Could not load GRPC CA: %v", err)
	}

	pair, err := tls.LoadX509KeyPair(certPath, keyPath)
	if err != nil {
		klog.Exitf("Could not load GRPC TLS keypair: %v", err)
	}
	return &tls.Config{
		Certificates: []tls.Certificate{pair},
		ClientAuth:   tls.RequireAndVerifyClientCert,
		ClientCAs:    certPool,
	}
}

// PrometheusRegistry returns this component's singleton Prometheus registry,
// creating it as needed. This method is not goroutine-safe, and should only be
// called during the setup process of the Component.
func (c *ComponentConfig) PrometheusRegistry() *prometheus.Registry {
	if c.prometheusRegistry == nil {
		c.prometheusRegistry = prometheus.NewRegistry()
		c.prometheusRegistry.Register(collectors.NewGoCollector())
		c.prometheusRegistry.Register(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
	}
	return c.prometheusRegistry
}

// StartPrometheus starts a Prometheus metrics server in a goroutine. It will
// serve any metrics that have been registered with the registry returned by
// PrometheusRegistry.
func (c *ComponentConfig) StartPrometheus(ctx context.Context) {
	reg := c.PrometheusRegistry()

	mux := http.NewServeMux()
	mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))

	var lis net.Listener
	var err error

	if c.PrometheusInsecure {
		lis, err = net.Listen("tcp", c.PrometheusListenAddress)
	} else {
		lis, err = tls.Listen("tcp", c.PrometheusListenAddress, c.getTLSConfig())
	}
	if err != nil {
		klog.Exitf("Could not listen on prometheus address: %v", err)
	}

	srv := http.Server{
		Handler: mux,
	}
	go func() {
		klog.Infof("Prometheus listening on %s", lis.Addr())
		if err := srv.Serve(lis); err != nil && ctx.Err() == nil {
			klog.Exitf("Prometheus serve failed: %v", err)
		}
	}()
	go func() {
		<-ctx.Done()
		srv.Close()
	}()
}

// GRPCServerOptions returns pre-built grpc.ServerOptions that this component
// should use to serve internal gRPC.
func (c *ComponentConfig) GRPCServerOptions() []grpc.ServerOption {
	return []grpc.ServerOption{
		grpc.Creds(credentials.NewTLS(c.getTLSConfig())),
	}
}

// GRPCServerOptionsPublic returns pre-built grpc.ServerOptions that this
// component should use to serve public gRPC. Any client will be allowed to
// connect, and it's up to the server implementation to authenticate incoming
// requests.
func (c *ComponentConfig) GRPCServerOptionsPublic() []grpc.ServerOption {
	var certPath, keyPath string
	if c.DevCerts {
		// Use devcerts if requested.
		certPath, keyPath, _ = c.GetDevCerts()
	} else {
		// Otherwise, use data from flags.
		if c.GRPCKeyPath == "" {
			klog.Exitf("-grpc_key_path must be set")
		}
		if c.GRPCCertificatePath == "" {
			klog.Exitf("-grpc_certificate_path must be set")
		}
		keyPath = c.GRPCKeyPath
		certPath = c.GRPCCertificatePath
	}

	pair, err := tls.LoadX509KeyPair(certPath, keyPath)
	if err != nil {
		klog.Exitf("Could not load GRPC TLS keypair: %v", err)
	}
	tlsConf := &tls.Config{
		Certificates: []tls.Certificate{pair},
		ClientAuth:   tls.RequestClientCert,
	}
	return []grpc.ServerOption{
		grpc.Creds(credentials.NewTLS(tlsConf)),
	}
}