package dnsstat

import (
	"context"

	"github.com/google/uuid"

	"github.com/coredns/caddy"
	"github.com/coredns/coredns/core/dnsserver"
	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/plugin/metrics"
	"github.com/coredns/coredns/request"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"

	"github.com/miekg/dns"
)

// Dump implement the plugin interface.
type DNSStat struct {
	Next plugin.Handler
}

// Session ID for current run for future better log partitioning in prometheus
var sessionID string

func init() {
	sessionID = uuid.New().String()
	plugin.Register("dnsstat", setup)
}

// Register
func setup(c *caddy.Controller) error {
	for c.Next() {
		if c.NextArg() {
			return plugin.Error("dnsstat", c.ArgErr())
		}
	}

	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
		return DNSStat{Next: next}
	})

	return nil
}

// Metrics counter
var requestCount = promauto.NewCounterVec(prometheus.CounterOpts{
	Namespace: plugin.Namespace,
	Subsystem: "dnsstat",
	Name:      "request_count_total",
	Help:      "Counter of requests made.",
}, []string{
	"server",
	"zone",
	"class",
	"type",
	"name",
	"client_ip",
	"session_id",
})

// ServeDNS implements the plugin.Handler interface.
func (d DNSStat) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
	state := request.Request{W: w, Req: r}
	requestCount.WithLabelValues(
		metrics.WithServer(ctx),
		state.Zone,
		state.Class(),
		state.Type(),
		state.Name(),
		state.IP(),
		sessionID,
	).Inc()
	return plugin.NextOrFailure(d.Name(), d.Next, ctx, w, r)
}

// Name implements the Handler interface.
func (d DNSStat) Name() string { return "dnsstat" }