package main

import (
	"flag"
	"fmt"
	"io/fs"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	httpHost  string
	httpPort  int
	interval  int
	namespace string
	path      string
	subsystem string

	extensionKeys map[string]bool
	gauges        map[string]*prometheus.GaugeVec
)

func main() {
	// Set options
	flag.StringVar(&httpHost, "host", "[::]", "The path to export metrics of")
	flag.IntVar(&httpPort, "port", 9091, "The path to export metrics of")
	flag.IntVar(&interval, "interval", 10, "The path to export metrics of")
	flag.StringVar(&namespace, "namespace", "fsdir_exporter", "Namespace to use for the Prometheus metrics")
	flag.StringVar(&path, "path", "", "The path to export metrics of")
	flag.StringVar(&subsystem, "subsystem", "", "Subsystem to use for the Prometheus metrics")
	flag.Parse()

	// Check required options
	if path == "" {
		slog.Error("No path set")
		os.Exit(1)
	}

	// Instantiate Prometheus metrics
	gauges = make(map[string]*prometheus.GaugeVec)

	gauges["count"] = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Namespace: namespace,
		Subsystem: subsystem,
		Name:      "count",
		Help:      "Number of files found in the given directory",
	}, []string{
		"extension",
	})

	gauges["size"] = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Namespace: namespace,
		Subsystem: subsystem,
		Name:      "size",
		Help:      "Size (in bytes) of files found in the given directory",
	}, []string{
		"extension",
	})

	// Register Prometheus metrics
	slog.Debug("Registering Prometheus metrics")
	prometheus.MustRegister(
		gauges["count"],
		gauges["size"],
	)

	// Start the update process
	extensionKeys = make(map[string]bool)

	update()
	go updateLoop()

	// Bind HTTP handlers
	slog.Debug("Binding HTTP handlers")
	http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "OK\n")
	})
	http.Handle("/metrics", promhttp.Handler())

	// Run HTTP service
	slog.Info("Serving HTTP requests", "address", httpHost, "port", httpPort)
	err := http.ListenAndServe(httpHost+":"+strconv.Itoa(httpPort), nil)

	slog.Error(err.Error())
}

func update() {
	sizes := make(map[string]float64)
	counts := make(map[string]float64)

	slog.Debug("Scanning filesystem for updates")
	err := filepath.Walk(path, func(path string, info fs.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		ext := filepath.Ext(info.Name())

		// Remove the leading . if needed
		if ext != "" {
			ext = ext[1:]
		}

		extensionKeys[ext] = true

		counts[ext]++
		sizes[ext] += float64(info.Size())

		return nil
	})

	if err != nil {
		slog.Error(fmt.Sprintf("error walking the path %v\n", err))
		return
	}

	// Loop over all keys and set the respective gauge values
	slog.Debug("Writing new values to metrics")
	for key, _ := range extensionKeys {
		if _, ok := counts[key]; ok {
			gauges["count"].With(prometheus.Labels{
				"extension": key,
			}).Set(counts[key])
			gauges["size"].With(prometheus.Labels{
				"extension": key,
			}).Set(sizes[key])

			continue
		}

		// Clean up extension which no longer exists in the directory
		gauges["count"].Delete(prometheus.Labels{"extension": key})
		gauges["size"].Delete(prometheus.Labels{"extension": key})
		delete(extensionKeys, key)
	}
}

func updateLoop() {
	for range time.Tick(time.Duration(interval) * time.Second) {
		update()
	}
}