package main

import (
	"context"
	"flag"
	"fmt"
	"os"

	"github.com/nberlee/go-netstat/netstat"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/expfmt"
)

var (
	maybe_port = flag.Uint("port", 8500, "Port to monitor")
)

func main() {
	flag.Parse()

	// Check port number before attempting to cast it to a uint16
	if *maybe_port < 1 || *maybe_port > 65535 {
		fmt.Fprintln(os.Stderr, "Invalid port number")
		os.Exit(1)
	}

	port := uint16(*maybe_port)

	reg := prometheus.NewRegistry()

	open_consul_connections := prometheus.NewGauge(prometheus.GaugeOpts{
		Subsystem: "consul",
		Name:      "open_tcp_connections",
		Help:      "Number of open consul tcp connections",
		ConstLabels: prometheus.Labels{
			"port": fmt.Sprintf("%d", port),
		},
	})

	reg.MustRegister(open_consul_connections)

	open_consul_connections.Set(float64(ConsulConnectionCount(port)))

	metricFamilies, err := reg.Gather()
	if err != nil {
		panic(err)
	}

	for _, mf := range metricFamilies {
		expfmt.MetricFamilyToText(os.Stdout, mf)
	}

}

// Filters
func PortFilter(port uint16) func(s *netstat.SockTabEntry) bool {
	return func(s *netstat.SockTabEntry) bool {
		return s.LocalEndpoint.Port == port
	}
}

func StateFilter(s *netstat.SockTabEntry) bool {
	return s.State == netstat.Established
}

func ConsulConnectionCount(port uint16) int {
	ctx := context.Background()
	features := netstat.EnableFeatures{
		TCP:  true,
		TCP6: true,
		UDP:  true,
		UDP6: true,
		PID:  true,
	}

	targetFilterFn := PortFilter(port)

	resp, err := netstat.Netstat(ctx, features, func(ste *netstat.SockTabEntry) bool {
		return targetFilterFn(ste) && StateFilter(ste)
	})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		return 0
	}

	return len(resp)
}