// Package metrics collects counters for the dashboard and renders them in the
// Prometheus text exposition format.
//
// The exposition format is small and stable, so it is written directly rather
// than pulling in the Prometheus client library and its dependency tree. The
// counters are plain atomics on the query path.
package metrics

import (
	"fmt"
	"io"
	"math"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// labelCounter is a set of counters keyed by a single label value.
type labelCounter struct {
	mu     sync.RWMutex
	values map[string]*atomic.Int64
}

func newLabelCounter() *labelCounter {
	return &labelCounter{values: map[string]*atomic.Int64{}}
}

func (c *labelCounter) Inc(label string) {
	c.mu.RLock()
	v, ok := c.values[label]
	c.mu.RUnlock()
	if ok {
		v.Add(1)
		return
	}
	c.mu.Lock()
	if v, ok = c.values[label]; !ok {
		v = &atomic.Int64{}
		c.values[label] = v
	}
	c.mu.Unlock()
	v.Add(1)
}

// Snapshot returns the counters sorted by label.
func (c *labelCounter) Snapshot() []LabelValue {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]LabelValue, 0, len(c.values))
	for k, v := range c.values {
		out = append(out, LabelValue{Label: k, Value: v.Load()})
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Value != out[j].Value {
			return out[i].Value > out[j].Value
		}
		return out[i].Label < out[j].Label
	})
	return out
}

// LabelValue is one labelled counter reading.
type LabelValue struct {
	Label string `json:"label"`
	Value int64  `json:"value"`
}

// histogram is a fixed-bucket latency histogram in seconds.
type histogram struct {
	bounds []float64
	counts []atomic.Int64
	sum    atomic.Uint64 // float64 bits
	total  atomic.Int64
}

func newHistogram(bounds []float64) *histogram {
	return &histogram{bounds: bounds, counts: make([]atomic.Int64, len(bounds)+1)}
}

func (h *histogram) Observe(d time.Duration) {
	v := d.Seconds()
	i := sort.SearchFloat64s(h.bounds, v)
	h.counts[i].Add(1)
	h.total.Add(1)
	for {
		old := h.sum.Load()
		nv := float64FromBits(old) + v
		if h.sum.CompareAndSwap(old, float64ToBits(nv)) {
			return
		}
	}
}

func (h *histogram) write(w io.Writer, name, help string) {
	fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s histogram\n", name, help, name)
	var cumulative int64
	for i, b := range h.bounds {
		cumulative += h.counts[i].Load()
		fmt.Fprintf(w, "%s_bucket{le=\"%s\"} %d\n", name, strconv.FormatFloat(b, 'g', -1, 64), cumulative)
	}
	cumulative += h.counts[len(h.bounds)].Load()
	fmt.Fprintf(w, "%s_bucket{le=\"+Inf\"} %d\n", name, cumulative)
	fmt.Fprintf(w, "%s_sum %s\n", name, strconv.FormatFloat(float64FromBits(h.sum.Load()), 'g', -1, 64))
	fmt.Fprintf(w, "%s_count %d\n", name, h.total.Load())
}

// Mean returns the average observation in milliseconds.
func (h *histogram) Mean() float64 {
	n := h.total.Load()
	if n == 0 {
		return 0
	}
	return float64FromBits(h.sum.Load()) / float64(n) * 1000
}

// The sum is kept as float64 bits inside an atomic so that Observe stays
// lock-free on the query path.
func float64ToBits(f float64) uint64   { return math.Float64bits(f) }
func float64FromBits(u uint64) float64 { return math.Float64frombits(u) }

// Gauges are values sampled at scrape time rather than counted incrementally.
type Gauges struct {
	CacheEntries     int64
	CacheBytes       int64
	Zones            int64
	Records          int64
	BlacklistDomains int64
	AllowlistDomains int64
	Networks         int64
	UpstreamsHealthy int64
	UpstreamsTotal   int64
	QueryLogRows     int64
}

// Metrics holds every counter the application exports.
type Metrics struct {
	start time.Time

	QueriesTotal    atomic.Int64
	Authoritative   atomic.Int64
	Recursive       atomic.Int64
	CacheHits       atomic.Int64
	CacheMisses     atomic.Int64
	StaleServed     atomic.Int64
	Blocked         atomic.Int64
	Refused         atomic.Int64
	RateLimited     atomic.Int64
	Errors          atomic.Int64
	TruncatedResp   atomic.Int64
	UDPQueries      atomic.Int64
	TCPQueries      atomic.Int64
	ResolverQueries atomic.Int64
	ResolverErrors  atomic.Int64

	byType   *labelCounter
	byRcode  *labelCounter
	bySource *labelCounter

	queryDuration   *histogram
	resolverLatency *histogram

	gaugeFn atomic.Value // func() Gauges

	buildVersion string
}

// New creates a metrics registry.
func New(version string) *Metrics {
	return &Metrics{
		start:        time.Now(),
		byType:       newLabelCounter(),
		byRcode:      newLabelCounter(),
		bySource:     newLabelCounter(),
		buildVersion: version,
		// Buckets chosen around the latencies a DNS server actually sees:
		// sub-millisecond for cache and authoritative hits, tens of
		// milliseconds for upstream queries.
		queryDuration: newHistogram([]float64{
			0.0001, 0.00025, 0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5,
		}),
		resolverLatency: newHistogram([]float64{
			0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5,
		}),
	}
}

// SetGaugeSource registers the callback used to sample gauges at scrape time.
func (m *Metrics) SetGaugeSource(fn func() Gauges) { m.gaugeFn.Store(fn) }

// Uptime returns how long the process has been serving.
func (m *Metrics) Uptime() time.Duration { return time.Since(m.start) }

// StartedAt returns the process start time.
func (m *Metrics) StartedAt() time.Time { return m.start }

// ObserveQuery records one completed query.
func (m *Metrics) ObserveQuery(qtype, rcode, source, protocol string, d time.Duration) {
	m.QueriesTotal.Add(1)
	m.byType.Inc(qtype)
	m.byRcode.Inc(rcode)
	m.bySource.Inc(source)
	m.queryDuration.Observe(d)
	if protocol == "tcp" {
		m.TCPQueries.Add(1)
	} else {
		m.UDPQueries.Add(1)
	}
}

// ObserveResolver records one upstream exchange.
func (m *Metrics) ObserveResolver(d time.Duration, err bool) {
	m.ResolverQueries.Add(1)
	if err {
		m.ResolverErrors.Add(1)
		return
	}
	m.resolverLatency.Observe(d)
}

// QueriesPerSecond returns the average query rate since start.
func (m *Metrics) QueriesPerSecond() float64 {
	secs := time.Since(m.start).Seconds()
	if secs < 1 {
		secs = 1
	}
	return float64(m.QueriesTotal.Load()) / secs
}

// CacheHitRate returns the cache hit percentage.
func (m *Metrics) CacheHitRate() float64 {
	h, ms := m.CacheHits.Load(), m.CacheMisses.Load()
	if h+ms == 0 {
		return 0
	}
	return float64(h) / float64(h+ms) * 100
}

// AvgQueryMS returns the mean query duration in milliseconds.
func (m *Metrics) AvgQueryMS() float64 { return m.queryDuration.Mean() }

// AvgResolverMS returns the mean upstream latency in milliseconds.
func (m *Metrics) AvgResolverMS() float64 { return m.resolverLatency.Mean() }

// ByType returns query counts per record type.
func (m *Metrics) ByType() []LabelValue { return m.byType.Snapshot() }

// ByRcode returns response counts per rcode.
func (m *Metrics) ByRcode() []LabelValue { return m.byRcode.Snapshot() }

// BySource returns response counts per answer source.
func (m *Metrics) BySource() []LabelValue { return m.bySource.Snapshot() }

// Reset zeroes every counter. Used by the "reset statistics" action.
func (m *Metrics) Reset() {
	m.start = time.Now()
	for _, c := range []*atomic.Int64{
		&m.QueriesTotal, &m.Authoritative, &m.Recursive, &m.CacheHits, &m.CacheMisses,
		&m.StaleServed, &m.Blocked, &m.Refused, &m.RateLimited, &m.Errors,
		&m.TruncatedResp, &m.UDPQueries, &m.TCPQueries, &m.ResolverQueries, &m.ResolverErrors,
	} {
		c.Store(0)
	}
	m.byType = newLabelCounter()
	m.byRcode = newLabelCounter()
	m.bySource = newLabelCounter()
}

// WritePrometheus renders every metric in the Prometheus text format.
func (m *Metrics) WritePrometheus(w io.Writer) {
	counter := func(name, help string, v int64) {
		fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s counter\n%s %d\n", name, help, name, name, v)
	}
	gauge := func(name, help string, v any) {
		fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s gauge\n%s %v\n", name, help, name, name, v)
	}

	fmt.Fprintf(w, "# HELP vibedns_build_info Build information.\n# TYPE vibedns_build_info gauge\n")
	fmt.Fprintf(w, "vibedns_build_info{version=%q} 1\n", escapeLabel(m.buildVersion))

	gauge("vibedns_uptime_seconds", "Seconds since the DNS server started.",
		strconv.FormatFloat(time.Since(m.start).Seconds(), 'f', 3, 64))

	counter("vibedns_dns_queries_total", "Total DNS queries received.", m.QueriesTotal.Load())
	counter("vibedns_dns_queries_authoritative_total", "Queries answered from a local authoritative zone.", m.Authoritative.Load())
	counter("vibedns_dns_queries_recursive_total", "Queries answered by forwarding upstream.", m.Recursive.Load())
	counter("vibedns_dns_queries_blocked_total", "Queries blocked by policy.", m.Blocked.Load())
	counter("vibedns_dns_queries_refused_total", "Queries refused, mostly recursion denied by ACL.", m.Refused.Load())
	counter("vibedns_dns_queries_ratelimited_total", "Queries dropped by the per-client rate limiter.", m.RateLimited.Load())
	counter("vibedns_dns_errors_total", "Queries that ended in a server error.", m.Errors.Load())
	counter("vibedns_dns_responses_truncated_total", "Responses truncated, prompting TCP retry.", m.TruncatedResp.Load())
	counter("vibedns_dns_queries_udp_total", "Queries received over UDP.", m.UDPQueries.Load())
	counter("vibedns_dns_queries_tcp_total", "Queries received over TCP.", m.TCPQueries.Load())
	counter("vibedns_cache_hits_total", "Cache lookups that were served from cache.", m.CacheHits.Load())
	counter("vibedns_cache_misses_total", "Cache lookups that missed.", m.CacheMisses.Load())
	counter("vibedns_cache_stale_served_total", "Responses served from the stale cache window.", m.StaleServed.Load())
	counter("vibedns_resolver_queries_total", "Queries forwarded to an upstream resolver.", m.ResolverQueries.Load())
	counter("vibedns_resolver_errors_total", "Upstream resolver failures.", m.ResolverErrors.Load())

	writeLabelled(w, "vibedns_dns_queries_by_type_total", "Queries by record type.", "type", m.byType.Snapshot())
	writeLabelled(w, "vibedns_dns_responses_by_rcode_total", "Responses by rcode.", "rcode", m.byRcode.Snapshot())
	writeLabelled(w, "vibedns_dns_responses_by_source_total", "Responses by answer source.", "source", m.bySource.Snapshot())

	m.queryDuration.write(w, "vibedns_dns_query_duration_seconds", "End-to-end query handling time.")
	m.resolverLatency.write(w, "vibedns_resolver_latency_seconds", "Upstream resolver round-trip time.")

	if fn, ok := m.gaugeFn.Load().(func() Gauges); ok && fn != nil {
		g := fn()
		gauge("vibedns_cache_entries", "Entries currently held in the resolver cache.", g.CacheEntries)
		gauge("vibedns_cache_bytes", "Estimated memory used by the resolver cache.", g.CacheBytes)
		gauge("vibedns_zones", "Configured authoritative zones.", g.Zones)
		gauge("vibedns_records", "Configured resource records.", g.Records)
		gauge("vibedns_blacklist_domains", "Domains across all blacklists.", g.BlacklistDomains)
		gauge("vibedns_allowlist_domains", "Domains across all allowlists.", g.AllowlistDomains)
		gauge("vibedns_client_networks", "Configured client networks.", g.Networks)
		gauge("vibedns_resolver_upstreams", "Configured upstream resolvers.", g.UpstreamsTotal)
		gauge("vibedns_resolver_upstreams_healthy", "Upstream resolvers currently considered healthy.", g.UpstreamsHealthy)
		gauge("vibedns_query_log_rows", "Rows currently stored in the query log.", g.QueryLogRows)
	}
}

func writeLabelled(w io.Writer, name, help, label string, values []LabelValue) {
	fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s counter\n", name, help, name)
	for _, v := range values {
		fmt.Fprintf(w, "%s{%s=%q} %d\n", name, label, escapeLabel(v.Label), v.Value)
	}
}

func escapeLabel(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `"`, `\"`)
	s = strings.ReplaceAll(s, "\n", `\n`)
	return s
}