package main

import (
	"context"
	"flag"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"

	pb "git.ipng.ch/ipng/nginx-logtail/proto/logtailpb"
	"google.golang.org/grpc"
)

func main() {
	listen := flag.String("listen", envOr("AGGREGATOR_LISTEN", ":9091"), "gRPC listen address (env: AGGREGATOR_LISTEN)")
	collectors := flag.String("collectors", envOr("AGGREGATOR_COLLECTORS", ""), "comma-separated collector host:port addresses (env: AGGREGATOR_COLLECTORS)")
	source := flag.String("source", envOr("AGGREGATOR_SOURCE", hostname()), "name for this aggregator in responses (env: AGGREGATOR_SOURCE, default: hostname)")
	flag.Parse()

	if *collectors == "" {
		log.Fatal("aggregator: --collectors / AGGREGATOR_COLLECTORS is required")
	}

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	var collectorAddrs []string
	for _, addr := range strings.Split(*collectors, ",") {
		addr = strings.TrimSpace(addr)
		if addr != "" {
			collectorAddrs = append(collectorAddrs, addr)
		}
	}

	merger := NewMerger()
	cache := NewCache(merger, *source)
	registry := NewTargetRegistry(collectorAddrs)
	go cache.Run(ctx)

	for _, addr := range collectorAddrs {
		sub := NewCollectorSub(addr, merger, registry)
		go sub.Run(ctx)
		log.Printf("aggregator: subscribing to collector %s", addr)
	}

	lis, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatalf("aggregator: failed to listen on %s: %v", *listen, err)
	}
	grpcServer := grpc.NewServer()
	pb.RegisterLogtailServiceServer(grpcServer, NewServer(cache, *source, registry))

	go func() {
		log.Printf("aggregator: gRPC listening on %s (source=%s)", *listen, *source)
		if err := grpcServer.Serve(lis); err != nil {
			log.Printf("aggregator: gRPC server stopped: %v", err)
		}
	}()

	<-ctx.Done()
	log.Printf("aggregator: shutting down")
	grpcServer.GracefulStop()
}

func hostname() string {
	h, err := os.Hostname()
	if err != nil {
		return "unknown"
	}
	return h
}

func envOr(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}