// Copyright 2022 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"net/http"
	"os"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/promlog/flag"
	"github.com/prometheus/common/version"
	"github.com/prometheus/exporter-toolkit/web"
	webflag "github.com/prometheus/exporter-toolkit/web/kingpinflag"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

// SMARTctlManagerCollector implements the Collector interface.
type SMARTctlManagerCollector struct {
	CollectPeriod         string
	CollectPeriodDuration time.Duration
	Devices               []string

	logger log.Logger
}

// Describe sends the super-set of all possible descriptors of metrics
func (i SMARTctlManagerCollector) Describe(ch chan<- *prometheus.Desc) {
	prometheus.DescribeByCollect(i, ch)
}

// Collect is called by the Prometheus registry when collecting metrics.
func (i SMARTctlManagerCollector) Collect(ch chan<- prometheus.Metric) {
	info := NewSMARTctlInfo(ch)
	for _, device := range i.Devices {
		if json, err := readData(i.logger, device); err == nil {
			info.SetJSON(json)
			smart := NewSMARTctl(i.logger, json, ch)
			smart.Collect()
		} else {
			level.Error(i.logger).Log("msg", "Error collecting SMART data", "err", err.Error())
		}
	}
	info.Collect()
}

var (
	smartctlPath = kingpin.Flag("smartctl.path",
		"The path to the smartctl binary",
	).Default("/usr/sbin/smartctl").String()
	smartctlInterval = kingpin.Flag("smartctl.interval",
		"The interval between smarctl polls",
	).Default("60s").Duration()
	smartctlDevices = kingpin.Flag("smartctl.device",
		"The device to monitor (repeatable)",
	).Strings()
	smartctlFakeData = kingpin.Flag("smartctl.fake-data",
		"The device to monitor (repeatable)",
	).Default("false").Hidden().Bool()
)

func main() {
	listenAddress := kingpin.Flag("web.listen-address",
		"Address to listen on for web interface and telemetry",
	).Default(":9633").String()
	metricsPath := kingpin.Flag(
		"web.telemetry-path", "Path under which to expose metrics",
	).Default("/metrics").String()
	webConfig := webflag.AddFlags(kingpin.CommandLine)

	promlogConfig := &promlog.Config{}
	flag.AddFlags(kingpin.CommandLine, promlogConfig)
	kingpin.Version(version.Print("smartctl_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()
	logger := promlog.New(promlogConfig)

	level.Info(logger).Log("msg", "Starting systemd_exporter", "version", version.Info())
	level.Info(logger).Log("msg", "Build context", "build_context", version.BuildContext())

	// Scan the host devices
	json := readSMARTctlDevices(logger)
	scanDevices := json.Get("devices").Array()
	scanDevicesSet := make(map[string]bool)
	var scanDeviceNames []string
	for _, d := range scanDevices {
		deviceName := d.Get("name").String()
		level.Debug(logger).Log("msg", "Found device", "name", deviceName)
		scanDevicesSet[deviceName] = true
		scanDeviceNames = append(scanDeviceNames, deviceName)
	}

	// Read the configuration and verify that it is available
	devices := *smartctlDevices
	var readDeviceNames []string
	for _, device := range devices {
		if _, ok := scanDevicesSet[device]; ok {
			readDeviceNames = append(readDeviceNames, device)
		} else {
			level.Warn(logger).Log("msg", "Device unavailable", "name", device)
		}
	}

	if len(readDeviceNames) > 0 {
		devices = readDeviceNames
	} else {
		level.Info(logger).Log("msg", "No devices specified, trying to load them automatically")
		devices = scanDeviceNames
	}

	if len(devices) == 0 {
		level.Error(logger).Log("msg", "No devices found")
		os.Exit(1)
	}

	collector := SMARTctlManagerCollector{
		Devices: devices,
		logger:  logger,
	}

	reg := prometheus.NewPedanticRegistry()
	reg.MustRegister(
		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
		collectors.NewGoCollector(),
	)

	prometheus.WrapRegistererWithPrefix("", reg).MustRegister(collector)

	level.Info(logger).Log("msg", "Listening on", "address", *listenAddress)
	http.Handle(*metricsPath, promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		_, err := w.Write([]byte(`<html>
      <head><title>Smartctl Exporter</title></head>
      <body>
      <h1>Smartctl Exporter</h1>
      <p><a href="` + *metricsPath + `">Metrics</a></p>
      </body>
      </html>`))
		if err != nil {
			level.Error(logger).Log("msg", "Couldn't write response", "err", err)
		}
	})

	srv := &http.Server{Addr: *listenAddress}
	if err := web.ListenAndServe(srv, *webConfig, logger); err != nil {
		level.Error(logger).Log("err", err)
		os.Exit(1)
	}
}