package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	graphql "github.com/hasura/go-graphql-client"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	interval int
	timeout  int
	port     int
	endpoint string

	graphqlClient *graphql.Client
)

var statuses = []string{
	"NEVER_RUN",
	"RUNNING",
	"PAUSED",
	"COMPLETED",
	"CANCELLED",
	"FAILED",
}

func init() {
	flag.IntVar(&interval, "interval", 30, "Duration, in seconds, between checks")
	flag.IntVar(&timeout, "timeout", 5000, "Timeout in ms when connecting to Unraid")
	flag.IntVar(&port, "port", 9091, "The port to listen on")
	flag.StringVar(&endpoint, "endpoint", "", "The Unraid GraphQL endpoint")
	flag.Parse()

	envEndpoint := os.Getenv("UNRAID_ENDPOINT")
	if envEndpoint != "" {
		endpoint = envEndpoint
	}
	if endpoint == "" {
		panic("Endpoint must be provided")
	}

	token := os.Getenv("UNRAID_TOKEN")
	if token == "" {
		panic("UNRAID_TOKEN env var must be provided")
	}

	graphqlClient = graphql.NewClient(
		endpoint,
		&http.Client{
			Timeout: time.Millisecond * time.Duration(timeout),
		},
	).
		WithRequestModifier(func(r *http.Request) {
			r.Header.Set("content-type", "application/json")
			r.Header.Set("x-api-key", token)
		})
}

func main() {
	fmt.Println("Unraid-exporter")
	collector := newUnraidCollector()
	prometheus.MustRegister(collector)

	go func() {
		if err := fetchUnraidStatus(collector); err != nil {
			fmt.Printf("failed to get status from Unraid: %v\n", err)
		}
		for range time.Tick(time.Second * time.Duration(interval)) {
			if err := fetchUnraidStatus(collector); err != nil {
				fmt.Printf("failed to get status from Unraid: %v\n", err)
			}
		}
	}()

	http.Handle("/metrics", promhttp.Handler())
	fmt.Printf("Starting server on %d\n", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))
}

func fetchUnraidStatus(collector *UnraidCollector) error {
	var query struct {
		Array struct {
			ParityCheckStatus struct {
				Progress int64
				Status   string
				Date     string
				Duration int64
				Errors   int64
				Speed    string
			}
		}
	}

	err := graphqlClient.Query(context.Background(), &query, nil)
	if err != nil {
		return err
	}

	collector.metrics = []prometheus.Metric{
		prometheus.MustNewConstMetric(
			collector.parityProgressMetric, prometheus.CounterValue,
			float64(query.Array.ParityCheckStatus.Progress),
		),
		prometheus.MustNewConstMetric(
			collector.parityErrorCountMetric, prometheus.CounterValue,
			float64(query.Array.ParityCheckStatus.Errors),
		),
	}
	for _, status := range statuses {
		val := 0
		if status == query.Array.ParityCheckStatus.Status {
			val = 1
		}
		collector.metrics = append(collector.metrics, prometheus.MustNewConstMetric(
			collector.parityStatusMetric, prometheus.CounterValue,
			float64(val), status,
		))
	}

	return nil
}