// Adapted from https://github.com/periph/cmd/blob/main/ina219/main.go
//
// Copyright 2018 The Periph Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package main // import "drjosh.dev/ina219-exporter"

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"syscall"
	"time"

	"drjosh.dev/jmetrics"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"periph.io/x/conn/v3/i2c/i2creg"
	"periph.io/x/conn/v3/physic"
	"periph.io/x/devices/v3/ina219"
	"periph.io/x/host/v3"
)

func main() {
	httpAddr := flag.String("http-address", ":9455", "Listen addr for HTTP handler")
	i2cAddress := flag.Int("i2c-address", 0x45, "I²C address")
	i2cBus := flag.String("i2c-bus", "", "I²C bus (/dev/i2c-1)")
	senseResistor := flag.String("sense-resistor", "10m", "Resistance of the shunt resistor in ohms (optionally with suffix m = milli, {µ,mu} = micro, n = nano)")
	maxCurrent := flag.String("max-current", "8", "Maximum current through the device in amps (optional suffix as for --sense-resistor)")

	flag.Parse()

	resistor, err := parseNanoSI[physic.ElectricResistance](*senseResistor)
	if err != nil {
		log.Fatalf("parsing --sense-resistor: %w", err)
	}
	if resistor <= 0 {
		log.Fatalf("--sense-resistor must be positive")
	}
	current, err := parseNanoSI[physic.ElectricCurrent](*maxCurrent)
	if err != nil {
		log.Fatalf("parsing --max-current: %w", err)
	}
	if current <= 0 {
		log.Fatalf("--max-current must be positive")
	}

	if _, err := host.Init(); err != nil {
		log.Fatalf("host.Init() error: %v", err)
	}

	// Open default I²C bus.
	bus, err := i2creg.Open(*i2cBus)
	if err != nil {
		log.Fatalf("i2creg.Open(%v): %v", *i2cBus, err)
	}
	defer bus.Close()

	sensor, err := ina219.New(bus, &ina219.Opts{
		Address:       *i2cAddress,
		SenseResistor: resistor,
		MaxCurrent:    current,
	})
	if err != nil {
		log.Fatalf("ina219.New(...): %v", err)
	}

	constLabels := prometheus.Labels{"i2c_addr": fmt.Sprintf("0x%x", *i2cAddress)}

	busVoltageHist := promauto.NewHistogram(prometheus.HistogramOpts{
		Namespace:                    "ina219",
		Name:                         "bus_voltage",
		Help:                         "Bus voltage (V)",
		ConstLabels:                  constLabels,
		NativeHistogramBucketFactor:  1.001,
		NativeHistogramZeroThreshold: 0.001,
	})
	busVoltageSumm := jmetrics.NewLiteGaugeSummary(jmetrics.LiteGaugeSummaryOpts{
		Namespace:   "ina219",
		Subsystem:   "summ",
		Name:        "bus_voltage",
		Help:        "Bus voltage (V)",
		ConstLabels: constLabels,
	})
	prometheus.MustRegister(busVoltageSumm)

	busCurrentHist := promauto.NewHistogram(prometheus.HistogramOpts{
		Namespace:                    "ina219",
		Name:                         "bus_current",
		Help:                         "Bus current (A)",
		ConstLabels:                  constLabels,
		NativeHistogramBucketFactor:  1.001,
		NativeHistogramZeroThreshold: 0.001,
	})
	busCurrentSumm := jmetrics.NewLiteGaugeSummary(jmetrics.LiteGaugeSummaryOpts{
		Namespace:   "ina219",
		Subsystem:   "summ",
		Name:        "bus_current",
		Help:        "Bus current (A)",
		ConstLabels: constLabels,
	})
	prometheus.MustRegister(busCurrentSumm)

	busPowerHist := promauto.NewHistogram(prometheus.HistogramOpts{
		Namespace:                    "ina219",
		Name:                         "bus_power",
		Help:                         "Bus power (W)",
		ConstLabels:                  constLabels,
		NativeHistogramBucketFactor:  1.001,
		NativeHistogramZeroThreshold: 0.001,
	})
	busPowerSumm := jmetrics.NewLiteGaugeSummary(jmetrics.LiteGaugeSummaryOpts{
		Namespace:   "ina219",
		Subsystem:   "summ",
		Name:        "bus_power",
		Help:        "Bus power (W)",
		ConstLabels: constLabels,
	})
	prometheus.MustRegister(busPowerSumm)

	shuntVoltageHist := promauto.NewHistogram(prometheus.HistogramOpts{
		Namespace:                    "ina219",
		Name:                         "shunt_voltage",
		Help:                         "Shunt voltage (V)",
		ConstLabels:                  constLabels,
		NativeHistogramBucketFactor:  1.001,
		NativeHistogramZeroThreshold: 0.001,
	})
	shuntVoltageSumm := jmetrics.NewLiteGaugeSummary(jmetrics.LiteGaugeSummaryOpts{
		Namespace:   "ina219",
		Subsystem:   "summ",
		Name:        "shunt_voltage",
		Help:        "Shunt voltage (V)",
		ConstLabels: constLabels,
	})
	prometheus.MustRegister(shuntVoltageSumm)

	go func() {
		http.Handle("/metrics", promhttp.Handler())
		log.Fatal(http.ListenAndServe(*httpAddr, nil))
	}()

	// Read values from sensor every second.
	everySecond := time.NewTicker(time.Second).C
	halt := make(chan os.Signal, 1)
	signal.Notify(halt, syscall.SIGTERM)
	signal.Notify(halt, syscall.SIGINT)

	for {
		select {
		case <-everySecond:
			p, err := sensor.Sense()
			if err != nil {
				log.Fatalf("sensor.Sense() error: %v", err)
			}

			// sanity check sensor outputs:
			// current = shunt / resistor, but resistor is fixed,
			// so if current == 0 if and only if shunt == 0
			if (p.Current == 0) != (p.Shunt == 0) {
				log.Fatalf("Ohm's Law violation: current = %v but shunt = %v", p.Current, p.Shunt)
			}
			// power = current * voltage, similar logic
			if (p.Power == 0) != (p.Current == 0 || p.Voltage == 0) {
				log.Fatalf("Joule's First Law violation: power = %v but current = %v and voltage = %v", p.Power, p.Current, p.Voltage)
			}

			busVolts := float64(p.Voltage) / float64(physic.Volt)
			busVoltageHist.Observe(busVolts)
			busVoltageSumm.Observe(busVolts)

			busAmps := float64(p.Current) / float64(physic.Ampere)
			busCurrentHist.Observe(busAmps)
			busCurrentSumm.Observe(busAmps)

			busWatts := float64(p.Power) / float64(physic.Watt)
			busPowerHist.Observe(busWatts)
			busPowerSumm.Observe(busWatts)

			shuntVolts := float64(p.Shunt) / float64(physic.Volt)
			shuntVoltageHist.Observe(shuntVolts)
			shuntVoltageSumm.Observe(shuntVolts)

		case <-halt:
			return
		}
	}
}

var siSmallRE = regexp.MustCompile(`^(\d+)(m|mu|µ|n)?$`)

func parseNanoSI[U ~int64](s string) (U, error) {
	submatches := siSmallRE.FindStringSubmatch(s)
	if len(submatches) < 2 {
		return 0, fmt.Errorf("value %q did not match %q", s, siSmallRE)
	}
	valueInt, err := strconv.Atoi(submatches[1])
	if err != nil {
		return 0, err
	}
	value := U(valueInt)
	switch len(submatches) {
	case 2:
		value *= 1_000_000_000
	case 3:
		switch suffix := submatches[2]; suffix {
		case "":
			value *= 1_000_000_000
		case "m":
			value *= 1_000_000
		case "mu", "µ":
			value *= 1_000
		case "n":
			value *= 1
		default:
			return 0, fmt.Errorf("invalid suffix %q", suffix)
		}
	}
	return value, nil
}