package main

import (
	"context"
	"flag"
	"log"
	"os/signal"
	"syscall"
	"time"

	"git.govekar.net/LovroG05/nvllama-watchdog/config"
	"git.govekar.net/LovroG05/nvllama-watchdog/core"
	"git.govekar.net/LovroG05/nvllama-watchdog/exporters"
	"git.govekar.net/LovroG05/nvllama-watchdog/readers"
)

func main() {
	cfgPath := flag.String("config", "config/config.yaml", "path to the YAML config file")
	flag.Parse()

	cfg, err := config.Load(*cfgPath)
	if err != nil {
		log.Fatalf("config: %v", err)
	}

	// Build readers
	var rs []core.Reader
	for _, rc := range cfg.Readers {
		switch rc.Type {
		case "llamacpp":
			rs = append(rs, readers.NewLLamaCPPReader(rc.Endpoint, rc.Model))
		case "nvidia_smi":
			rs = append(rs, readers.NewNvidiaSmiReader())
		case "system":
			rs = append(rs, readers.NewSystemReader(rc.Enabled))
		}
	}

	// Build exporters
	var es []core.Exporter
	var prom *exporters.PrometheusExporter
	for _, ec := range cfg.Exporters {
		switch ec.Type {
		case "prometheus":
			prom = exporters.NewPrometheusExporter(ec.Addr)
			es = append(es, prom)
		case "dev":
			es = append(es, exporters.NewDevExporter(ec.JSONFile))
		case "influxdb":
			es = append(es, exporters.NewInfluxDBExporter(ec.URL, ec.Token, ec.User, ec.Pass))
		}
	}

	// Start Prometheus HTTP listener in background
	if prom != nil {
		go func() {
			log.Printf("prometheus exporter listening on %s", prom.Addr())
			if err := prom.Start(); err != nil {
				log.Fatalf("prometheus: %v", err)
			}
		}()
	}

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

	ticker := time.NewTicker(cfg.ScrapeInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			var all []core.Metric
			for _, r := range rs {
				ms, err := r.Read(ctx)
				if err != nil {
					log.Printf("reader %s: %v", r.Name(), err)
					continue
				}
				all = append(all, ms...)
			}
			for _, e := range es {
				if err := e.Write(ctx, all); err != nil {
					log.Printf("exporter %s: %v", e.Name(), err)
				}
			}
		}
	}
}