package metrics

import (
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	"sneak.berlin/go/smallwebwaf/internal/requestlog"
)

// other is the label under which the countries outside the busiest are
// counted.
const other = "other"

// countries are the metrics by the client's country, for requests whose
// client's country is known. The topN busiest countries, by their requests
// since the start, have series of their own, and the others are counted
// under other, so that there are never more than topN + 1 series. A
// country that drops out of the busiest loses its series, and its next
// requests are counted under other; one that becomes one of them gets a
// series that counts from then on. Each series therefore only ever goes
// up.
type countries struct {
	topN int

	requests      *prometheus.CounterVec
	requestBytes  *prometheus.CounterVec
	responseBytes *prometheus.CounterVec
	// refused are the requests the country lists refused.
	refused *prometheus.CounterVec

	mu sync.Mutex
	// seen is each country's requests since the start, by which the
	// countries are ranked. GeoJS gives two-letter codes, so it holds at
	// most a few hundred.
	seen map[string]int64
	// top are the countries with series of their own.
	top map[string]bool
}

// newCountries returns the metrics by country, with series of their own
// for the topN busiest countries.
func newCountries(topN int) *countries {
	byCountry := []string{"country"}

	return &countries{
		topN: topN,
		requests: counterVec("smallwebwaf_country_requests_total",
			"Requests, by the client's country.", byCountry),
		requestBytes: counterVec("smallwebwaf_country_request_bytes_total",
			"Request body bytes, by the client's country.", byCountry),
		responseBytes: counterVec("smallwebwaf_country_response_bytes_total",
			"Response body bytes, by the client's country.", byCountry),
		refused: counterVec("smallwebwaf_country_list_refusals_total",
			"Requests the country lists refused, by the client's country.",
			byCountry),
		seen: map[string]int64{},
		top:  map[string]bool{},
	}
}

// add counts a request from its log line, whose country is known.
func (c *countries) add(line *requestlog.Line) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.seen[line.Country]++

	label := c.label(line.Country)
	c.requests.WithLabelValues(label).Inc()
	c.requestBytes.WithLabelValues(label).Add(float64(line.RequestBytes))
	c.responseBytes.WithLabelValues(label).Add(float64(line.ResponseBytes))

	if line.Action == requestlog.ActionCountryDenied {
		c.refused.WithLabelValues(label).Inc()
	}
}

// label returns the label a request from country is counted under: the
// country while it is one of the busiest, other while it is not. A
// country busier than the least busy of them takes its place, and that
// country's series are dropped.
func (c *countries) label(country string) string {
	if c.top[country] {
		return country
	}

	if len(c.top) < c.topN {
		c.top[country] = true

		return country
	}

	least := ""

	for top := range c.top {
		if least == "" || c.seen[top] < c.seen[least] {
			least = top
		}
	}

	if c.seen[country] <= c.seen[least] {
		return other
	}

	delete(c.top, least)

	for _, vec := range []*prometheus.CounterVec{
		c.requests, c.requestBytes, c.responseBytes, c.refused,
	} {
		vec.DeleteLabelValues(least)
	}

	c.top[country] = true

	return country
}