package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

const (
	metricsPath = "/metrics"
	namespace   = "app"
)

var (
	listenAddress        = flag.String("listen", ":8080", "The address to listen on for HTTP requests.")
	endpointApp          = flag.String("endpoint", "http://localhost:8050/stats", "HTTP API address of the application")
	timeoutApp           = flag.Int("timeout", 500, "Connection timeout in ms")
	environment          = flag.String("environment", "", "Optional environment which will be added to the exported metrics")
	prometheusConstLabel = parseConstLabel()
)

// run before creating the descriptor
func parseConstLabel() prometheus.Labels {
	// parse flags in an early state, so we can retrieve the instance id
	flag.Parse()

	// generate constant label if environment is present
	if *environment != "" {
		return prometheus.Labels{"environment": *environment}
	}
	return prometheus.Labels{}
}

var (
	// Create the prometheus descriptors
	descUp = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, "", "up"),
		"Was the last query successful.",
		nil, prometheusConstLabel,
	)
	descRequestCount = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, "", "request_count"),
		"How many requests processed, partitioned by status code.",
		[]string{"code"}, prometheusConstLabel,
	)
	descRequestCountTotal = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, "", "request_count_total"),
		"How many requests processed of all status codes.",
		nil, prometheusConstLabel,
	)
	descRequestRates = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, "", "request_rates"),
		"How many requests processed in the last second, partitioned by status code.",
		[]string{"code"}, prometheusConstLabel,
	)
	descRequestRatesTotal = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, "", "request_rates_total"),
		"How many requests processed in the last second.",
		nil, prometheusConstLabel,
	)
	descDurationSum = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, "", "duration_sum"),
		"How much time consumed the requests in summary.",
		nil, prometheusConstLabel,
	)
	descDurationAvg = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, "", "duration_avg"),
		"How much time consumed the requests in average.",
		nil, prometheusConstLabel,
	)
)

// AppStats represent the schema of the returned json
type AppStats struct {
	// RequestCounters are the served status codes during app lifetime
	RequestCounters map[string]int `json:"requestCounters"`
	// RequestRates are the served status codes for the last second (QPS)
	RequestRates map[string]int `json:"requestRates"`
	// Duration represent some request stats during the lifetime
	Duration *AppDuration `json:"duration"`
}

// AppDuration the schema of the returned duration part of the json
type AppDuration struct {
	// Count is the total served request in the lifetime
	Count int `json:"count"`
	// Sum is the total time of taken time the requests have taken in seconds
	Sum float64 `json:"sum"`
	// Average time of usage a request has taken.
	Average float64 `json:"average"`
}

// appScraper is a helper to retrieve stats in a generic way
type appScraper struct {
	endpoint string
	client   *http.Client
}

// stats returns the fetched and parsed json
func (s *appScraper) stats() (*AppStats, error) {
	var stats AppStats
	response, err := s.client.Get(s.endpoint)
	if err != nil {
		return nil, err
	}

	buf, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(buf, &stats)
	if err != nil {
		return nil, err
	}

	// validate returned json is complete
	// requestCounter and requestRates should be empty maps, so we are fine here
	if stats.Duration == nil {
		return nil, errors.New("Invalid JSON returned, could not retrieve duration")
	}

	return &stats, err
}

// Exporter implements prometheus.Collector
type Exporter struct {
	scraper *appScraper
}

// Describe implements prometheus.Describe
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- descUp
	ch <- descRequestCount
	ch <- descRequestCountTotal
	ch <- descRequestRates
	ch <- descDurationSum
	ch <- descDurationAvg

}

// Collect implements prometheus.Describe.Collect
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {

	stats, err := e.scraper.stats()
	if err != nil {
		ch <- prometheus.MustNewConstMetric(
			descUp, prometheus.GaugeValue, 0,
		)
		log.Error("Failed to scrape app stats: ", err)
		return
	}
	ch <- prometheus.MustNewConstMetric(
		descUp, prometheus.GaugeValue, 1,
	)

	// Add counter per code on the fly (no need to update code if there are additional codes)
	for code, count := range stats.RequestCounters {
		ch <- prometheus.MustNewConstMetric(
			descRequestCount, prometheus.CounterValue, float64(count), code,
		)
	}

	// Add total of all requests
	ch <- prometheus.MustNewConstMetric(
		descRequestCountTotal, prometheus.CounterValue, float64(stats.Duration.Count),
	)

	// Add rates per code on the fly (no need to update code if there are additional codes)
	ratesSum := 0
	for code, count := range stats.RequestRates {
		ch <- prometheus.MustNewConstMetric(
			descRequestRates, prometheus.GaugeValue, float64(count), code,
		)
		ratesSum += count
	}
	// Additional sum of the rates, such like the requestCounter
	ch <- prometheus.MustNewConstMetric(
		descRequestRatesTotal, prometheus.CounterValue, float64(ratesSum),
	)

	ch <- prometheus.MustNewConstMetric(
		descDurationSum, prometheus.CounterValue, stats.Duration.Sum,
	)
	ch <- prometheus.MustNewConstMetric(
		descDurationAvg, prometheus.GaugeValue, stats.Duration.Average,
	)

}

// NewExporter creates a new prometheus exporter and app scraper
func NewExporter(endpoint string) (*Exporter, error) {

	if !strings.Contains(endpoint, "://") {
		endpoint = "http://" + endpoint
	}
	u, err := url.Parse(endpoint)
	if err != nil {
		return nil, fmt.Errorf("invalid endpoint URL: %s", err)
	}
	if u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") {
		return nil, fmt.Errorf("invalid endpoint URL: %s", endpoint)
	}

	// use custom http client with specific timeout
	client := &http.Client{
		Timeout: time.Duration(*timeoutApp) * time.Millisecond,
	}

	// create api client
	appScraper := &appScraper{
		client:   client,
		endpoint: endpoint,
	}

	return &Exporter{
		scraper: appScraper,
	}, nil
}

func init() {
	prometheus.MustRegister(version.NewCollector(fmt.Sprintf("%s_exporter", namespace)))
}

func main() {

	exporter, err := NewExporter(*endpointApp)
	if err != nil {
		log.Fatal(err)
	}
	prometheus.MustRegister(exporter)

	http.Handle(metricsPath, promhttp.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`<html>
             <head><title>App Exporter</title></head>
             <body>
             <h1>App Exporter</h1>
             <p><a href='` + metricsPath + `'>Metrics</a></p>
             </body>
             </html>`))
	})

	log.Infoln("Listening on", *listenAddress)
	log.Fatal(http.ListenAndServe(*listenAddress, nil))
}