package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"go.uber.org/zap/zapcore"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/cluster-prober/config"
	"github.com/tidbcloud/infra-provider/cluster-prober/prober"
	dbprober "github.com/tidbcloud/infra-provider/cluster-prober/prober/database"
	infraprober "github.com/tidbcloud/infra-provider/cluster-prober/prober/infra"
	platformprober "github.com/tidbcloud/infra-provider/cluster-prober/prober/platform"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
)

var (
	scheme   = runtime.NewScheme()
	setupLog = log.Log.WithName("Setup")
)

func main() {
	rand.Seed(time.Now().UnixNano())

	var (
		rootConfigPatch string
	)
	flag.StringVar(&rootConfigPatch, "root-config", "", "The path to the root config.")

	opts := zap.Options{
		Development: true,
		TimeEncoder: zapcore.RFC3339TimeEncoder,
	}
	opts.BindFlags(flag.CommandLine)
	flag.Parse()
	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts)))

	rootConfig, err := config.LoadRootConfig(rootConfigPatch)
	exitIfErr(err, "unable to load root config")

	setupLog.Info("loading scheme")
	loadScheme()

	setupLog.Info("building dependencies")
	ctx := ctrl.SetupSignalHandler()
	deps := &prober.Dependency{
		Logger: log.Log,
	}
	predicates := rootConfig.Predicate
	selector, err := metav1.LabelSelectorAsSelector(&predicates)
	exitIfErr(err, "failed to build selector")

	if rootConfig.InfraKubeConfig != nil {
		cfg, err := k8s.GetKubeConfig(*rootConfig.InfraKubeConfig)
		exitIfErr(err, "failed to load kubeconfig of infra")
		cache, err := cache.New(cfg, cache.Options{
			Scheme:   scheme,
			ByObject: buildCacheSelector(selector),
		})
		exitIfErr(err, "failed to init cache")
		go func() {
			exitIfErr(cache.Start(ctx), "failed to start cache")
		}()
		ok := cache.WaitForCacheSync(ctx)
		if !ok {
			exitIfErr(fmt.Errorf("failed to sync cache"), "")
		}
		deps.InfraClient = cache
	}

	setupLog.Info("registering probers")
	promRegistry := prometheus.NewRegistry()
	promRegistry.MustRegister(buildProbers(rootConfig, deps)...)

	mux := http.DefaultServeMux
	mux.Handle(rootConfig.MetricPath, promhttp.HandlerFor(promRegistry, promhttp.HandlerOpts{}))
	exitIfErr(err, "failed to handle metrics")
	startHTTPServer(ctx, &http.Server{Addr: rootConfig.ListenAddr, Handler: mux})
}

func exitIfErr(err error, msg string) {
	if err != nil {
		setupLog.Error(err, msg)
		os.Exit(1)
	}
}

func loadScheme() {
	utilruntime.Must(cloudv1.SchemeBuilder.AddToScheme(scheme))
}

func buildCacheSelector(selector labels.Selector) map[client.Object]cache.ByObject {
	return map[client.Object]cache.ByObject{
		&cloudv1.Cluster{}: {
			Label: selector,
		},
	}
}

func startHTTPServer(ctx context.Context, server *http.Server) {
	closed := make(chan struct{})
	go func() {
		<-ctx.Done()
		if err := server.Shutdown(context.Background()); err != nil {
			setupLog.Error(err, "failed to shutdown http server")
		}
		close(closed)
	}()

	setupLog.Info("starting http server", "addr", server.Addr)
	err := server.ListenAndServe()
	if err != http.ErrServerClosed {
		exitIfErr(err, "failed to start metrics server")
	}

	<-closed
	setupLog.Info("closed http server")
}

func buildProbers(cfg *config.RootConfig, deps *prober.Dependency) []prometheus.Collector {
	probers := []prometheus.Collector{}

	if cfg.Probers.Database != nil {
		err := cfg.Probers.Database.Complete(cfg)
		exitIfErr(err, "failed to complete database prober config")
		p, err := dbprober.New(cfg.Probers.Database, deps)
		exitIfErr(err, "failed to register the database prober")

		probers = append(probers, p)
		setupLog.Info("registered the database prober")
	}
	if cfg.Probers.Infra != nil {
		err := cfg.Probers.Infra.Complete(cfg)
		exitIfErr(err, "failed to complete infra prober config")
		p, err := infraprober.New(cfg.Probers.Infra, deps)
		exitIfErr(err, "failed to register the infra prober")

		probers = append(probers, p)
		setupLog.Info("registered the infra prober")
	}
	if cfg.Probers.Platform != nil {
		err := cfg.Probers.Platform.Complete(cfg)
		exitIfErr(err, "failed to complete platform prober config")
		p, err := platformprober.New(cfg.Probers.Platform, deps)
		exitIfErr(err, "failed to register the platform prober")

		probers = append(probers, p)
		setupLog.Info("registered the platform prober")
	}

	return probers
}