package main

import (
	"context"
	"encoding/json"
	"flag"
	"io"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

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

// original /api/telemetry type struct
type MetricsV1 struct {
	TempNozzle    int     `json:"temp_nozzle"`
	TempBed       int     `json:"temp_bed"`
	Material      string  `json:"material"`
	ZPosition     float64 `json:"pos_z_mm"`
	PrintingSpeed int     `json:"printing_speed"`
	FlowFactor    int     `json:"flow_factor"`
	Progress      int     `json:"progress"`
	PrintDur      string  `json:"print_dur"`
	TimeEst       string  `json:"time_est"`
	TimeZone      string  `json:"time_zone"`
	ProjectName   string  `json:"project_name"`
}

// new 4.4.0+ /api/printer struct
type MetricsV2 struct {
	Telemetry   Telemetry   `json:"telemetry"`
	Temperature Temperature `json:"temperature"`
	State       State       `json:"state"`
}

type Telemetry struct {
	TempBed    float64 `json:"temp-bed"`
	TempNozzle float64 `json:"temp-nozzle"`
	PrintSpeed int     `json:"print-speed"`
	ZHeight    float64 `json:"z-height"`
	Material   string  `json:"material"`
}

type Temperature struct {
	Tool0 Tempatures `json:"tool0"`
	Bed   Tempatures `json:"bed"`
}

type Tempatures struct {
	Actual  float64 `json:"actual"`
	Target  float64 `json:"target"`
	Display float64 `json:"display"`
	Offset  float64 `json:"offset"`
}

type State struct {
	Text  string `json:"text"`
	Flags Flags  `json:"flags"`
}

type Flags struct {
	Operational   bool `json:"operational"`
	Paused        bool `json:"paused"`
	Printing      bool `json:"printing"`
	Cancelling    bool `json:"cancelling"`
	Pausing       bool `json:"pausing"`
	SdReady       bool `json:"sdReady"`
	Error         bool `json:"error"`
	ClosedOnError bool `json:"closedOnError"`
	Ready         bool `json:"ready"`
	Busy          bool `json:"busy"`
}

// /api/job struct
type JobV2 struct {
	State    string   `json:"state"`
	Job      Job      `json:"job"`
	Progress Progress `json:"progress"`
}

type File struct {
	Name    string `json:"name"`
	Path    string `json:"path"`
	Display string `json:"display"`
}

type Job struct {
	EstimatedPrintTime int  `json:"estimatedPrintTime"`
	File               File `json:"file"`
}

type Progress struct {
	Completion    float64 `json:"completion"`
	PrintTime     int     `json:"printTime"`
	PrintTimeLeft int     `json:"printTimeLeft"`
}

type Config struct {
	QueryHostname string
	QueryInterval int
	Port          string
	Path          string
	APIKey        string
}

var (
	errCount = 0
	gauges   = map[string]prometheus.Gauge{
		"flow_factor": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "flow_factor",
			Help:      "Current flow factor, as a unitless number",
		}),
		"printing_speed": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "printing_speed",
			Help:      "Current print speed, as a percentage",
		}),
		"z_position": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "z_position",
			Help:      "Vertical depth, in mm, of the Z head",
		}),
		"temp_bed": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "temp_bed",
			Help:      "Temperature, in celsius, of the print bed",
		}),
		"temp_nozzle": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "temp_nozzle",
			Help:      "Temperature, in celsius, of the print nozzle",
		}),
		"progress": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "progress",
			Help:      "Current print completeness, as a percentage",
		}),
		"duration": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "duration",
			Help:      "Duration of current print job, as seconds since start",
		}),
		"remaining": prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "prusa_connect",
			Name:      "remaining",
			Help:      "Estimated remaining time of current print job, as seconds",
		}),
	}
)

func errLog(err error) {
	errCount++
	// reset metrics
	for _, g := range gauges {
		g.Set(0)
	}
	if errCount == 5 {
		log.Printf("suppressing further error logging")
		return
	} else if errCount > 5 {
		return
	}
	log.Printf("error retrieving telemetry: %v", err)
}

func recordMetrics(ctx context.Context, config Config) {
	go func() {
		for {
			select {
			case <-ctx.Done():
				log.Print("exiting!")
				return
			default:
				time.Sleep(time.Second * time.Duration(config.QueryInterval))
				if config.APIKey == "" {
					parseMetricsV1(ctx, config)
					break
				}
				parseMetricsV2(ctx, config)
			}
		}
	}()
}

func parseMetricsV1(ctx context.Context, config Config) {
	t, err := getTelemetry(config.QueryHostname)
	if err != nil {
		errLog(err)
		return
	}
	if errCount != 0 {
		errCount = 0
		log.Printf("connection established")
	}
	gauges["flow_factor"].Set(float64(t.FlowFactor))
	gauges["printing_speed"].Set(float64(t.PrintingSpeed))
	gauges["z_position"].Set(t.ZPosition)
	gauges["temp_bed"].Set(float64(t.TempBed))
	gauges["temp_nozzle"].Set(float64(t.TempNozzle))
	gauges["progress"].Set(float64(t.Progress))
	gauges["duration"].Set(float64(parseDuration(t.PrintDur) / time.Second))
	estimatedTimeRemaining, err := strconv.ParseFloat(t.TimeEst, 64)
	if err == nil {
		gauges["remaining"].Set(float64(estimatedTimeRemaining))
	}
}

func parseMetricsV2(ctx context.Context, config Config) {
	t, err := getPrinter(config.Port, config.APIKey)
	if err != nil {
		errLog(err)
		return
	}
	j, err := getJob(config.Port, config.APIKey)
	if err != nil {
		errLog(err)
		return
	}
	if errCount != 0 {
		errCount = 0
		log.Printf("connection established")
	}
	gauges["printing_speed"].Set(float64(t.Telemetry.PrintSpeed))
	gauges["z_position"].Set(t.Telemetry.ZHeight)
	gauges["temp_bed"].Set(float64(t.Telemetry.TempBed))
	gauges["temp_nozzle"].Set(float64(t.Telemetry.TempNozzle))
	gauges["progress"].Set(float64(j.Progress.Completion) * 100)
	gauges["duration"].Set(float64(j.Progress.PrintTime))
	gauges["remaining"].Set(float64(j.Progress.PrintTimeLeft))
}

func getJob(hostname, apiKey string) (*JobV2, error) {
	r, err := http.NewRequest(http.MethodGet, hostname+"/api/job", nil)
	if err != nil {
		return &JobV2{}, err
	}
	r.Header.Set("X-Api-Key", apiKey)
	res, err := http.DefaultClient.Do(r)
	if err != nil {
		return &JobV2{}, err
	}
	b, err := io.ReadAll(res.Body)
	if err != nil {
		return &JobV2{}, err
	}
	j := &JobV2{}
	err = json.Unmarshal(b, j)
	if err != nil {
		return &JobV2{}, err
	}
	return j, nil
}

func getPrinter(hostname, apiKey string) (*MetricsV2, error) {
	r, err := http.NewRequest(http.MethodGet, hostname+"/api/printer", nil)
	if err != nil {
		return &MetricsV2{}, err
	}
	r.Header.Set("X-Api-Key", apiKey)
	res, err := http.DefaultClient.Do(r)
	if err != nil {
		return &MetricsV2{}, err
	}
	b, err := io.ReadAll(res.Body)
	if err != nil {
		return &MetricsV2{}, err
	}
	j := &MetricsV2{}
	err = json.Unmarshal(b, j)
	if err != nil {
		return &MetricsV2{}, err
	}
	return j, nil
}

func getTelemetry(hostname string) (*MetricsV1, error) {
	res, err := http.Get(hostname + "/api/telemetry")
	if err != nil {
		return &MetricsV1{}, err
	}
	b, err := io.ReadAll(res.Body)
	if err != nil {
		return &MetricsV1{}, err
	}
	t := &MetricsV1{}
	err = json.Unmarshal(b, t)
	if err != nil {
		return &MetricsV1{}, err
	}
	return t, nil
}

// taking a guess on proper parsing here...
func parseDuration(d string) time.Duration {
	dur := 0
	d = strings.TrimSpace(d)
	number := []rune{}
	for _, t := range d {
		switch t {
		case ' ':
			continue
		case '1', '2', '3', '4', '5', '6', '7', '8', '9', '0':
			number = append(number, t)
		case 's':
			i, _ := strconv.Atoi(string(number))
			dur += i * int(time.Second)
			number = []rune{}
		case 'm':
			i, _ := strconv.Atoi(string(number))
			dur += (i * int(time.Minute))
			number = []rune{}
		case 'h':
			i, _ := strconv.Atoi(string(number))
			dur += (i * int(time.Hour))
			number = []rune{}
		case 'd':
			i, _ := strconv.Atoi(string(number))
			dur += (i * int(time.Hour) * 24)
			number = []rune{}
		}

	}
	return time.Duration(dur)
}

func main() {
	var hostname string
	flag.StringVar(&hostname, "hostname", "localhost", "Hostname the Prusa Connect API is available at (assumes http)")
	var queryInterval int
	flag.IntVar(&queryInterval, "interval", 2, "How often, in seconds, to query the API")
	var port string
	flag.StringVar(&port, "port", "2112", "Local port to export metrics on")
	var path string
	flag.StringVar(&path, "path", "metrics", "Local path to export metrics on")
	var apiKey string
	flag.StringVar(&apiKey, "apikey", "", "Prusa Connect API key (see Main Menu -> Settings -> Network on the printer)")
	flag.Parse()

	if queryInterval < 1 {
		log.Fatalf("query interval must be greater than 0; %d received", queryInterval)
	}

	config := Config{
		QueryHostname: "http://" + hostname,
		QueryInterval: queryInterval,
		Port:          port,
		Path:          path,
		APIKey:        apiKey,
	}

	r := prometheus.NewRegistry()
	for _, g := range gauges {
		r.MustRegister(g)
	}
	recordMetrics(context.Background(), config)

	log.Printf("starting exporter on :%v", config.Port)

	http.Handle("/"+config.Path, promhttp.HandlerFor(r, promhttp.HandlerOpts{}))
	http.ListenAndServe(":"+config.Port, nil)
}