package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/prometheus/client_golang/prometheus"
)

type bindMetric struct {
	Name    string
	Help    string
	Labels  []string
	Extract func(*prometheus.Desc, *BindStats) []prometheus.Metric
	desc    *prometheus.Desc
}

type bindViewMetric struct {
	Name    string
	Help    string
	Labels  []string
	Extract func(*prometheus.Desc, string, *View) []prometheus.Metric
	desc    *prometheus.Desc
}

type bindZoneMetric struct {
	Name    string
	Help    string
	Labels  []string
	Extract func(*prometheus.Desc, string, *View, *Zone) []prometheus.Metric
	desc    *prometheus.Desc
}

type bindSocketMetric struct {
	Name    string
	Help    string
	Labels  []string
	Suffix  string
	Extract func(*prometheus.Desc, string, map[string]int) []prometheus.Metric
	desc    *prometheus.Desc
}

type bindCollector struct {
	ctx           context.Context
	target        string
	metrics       []*bindMetric
	viewMetrics   []*bindViewMetric
	zoneMetrics   []*bindZoneMetric
	socketMetrics []*bindSocketMetric
}

var _ prometheus.Collector = (*bindCollector)(nil)

func socketMetricExtractor(d *prometheus.Desc, suffix string, m map[string]int) []prometheus.Metric {
	out := []prometheus.Metric{}

	for _, p := range []string{"UDP", "TCP"} {
		for _, v := range []string{"4", "6"} {
			if value, ok := m[p+v+suffix]; ok {
				out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(value), p, v))
			}
		}
	}

	if value, ok := m["Raw"+suffix]; ok {
		out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(value), "Raw", ""))
	}

	return out
}

func kvExtractor(d *prometheus.Desc, m map[string]int) []prometheus.Metric {
	out := make([]prometheus.Metric, 0, len(m))

	for k, v := range m {
		out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v), k))
	}

	return out
}

func singleGaugeIntMetric(d *prometheus.Desc, v int) []prometheus.Metric {
	return singleGaugeInt64Metric(d, int64(v))
}

func singleGaugeInt64Metric(d *prometheus.Desc, v int64) []prometheus.Metric {
	return []prometheus.Metric{prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v))}
}

func NewBindCollector(ctx context.Context, target string) prometheus.Collector {
	c := &bindCollector{
		ctx:    ctx,
		target: target,
		metrics: []*bindMetric{
			&bindMetric{
				Name: "bind_boot_time",
				Help: "Unix timestamp of when the server was last booted",
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return singleGaugeInt64Metric(d, s.BootTime.Unix())
				},
			},
			&bindMetric{
				Name: "bind_config_time",
				Help: "Unix timestamp of when the server was last configured",
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return singleGaugeInt64Metric(d, s.ConfigTime.Unix())
				},
			},
			&bindMetric{
				Name: "bind_memory_total_use",
				Help: "Total memory bytes in used by the server",
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return singleGaugeIntMetric(d, s.MemoryStats.TotalUse)
				},
			},
			&bindMetric{
				Name: "bind_memory_in_use",
				Help: "Total memory bytes currently in use by the server",
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return singleGaugeIntMetric(d, s.MemoryStats.InUse)
				},
			},
			&bindMetric{
				Name: "bind_memory_malloced",
				Help: "Total memory bytes currently malloced by the server",
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return singleGaugeIntMetric(d, s.MemoryStats.Malloced)
				},
			},
			&bindMetric{
				Name:   "bind_server_op_codes",
				Help:   "Operation codes requested from the server as a whole",
				Labels: []string{"opcode"},
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return kvExtractor(d, s.OperationCodeCount)
				},
			},
			&bindMetric{
				Name:   "bind_notify_out",
				Help:   "Count of notifies sent by IP version",
				Labels: []string{"ipv"},
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					out := []prometheus.Metric{}
					if v4out, ok := s.ZoneStats["NotifyOutv4"]; ok {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v4out), "4"))
					}
					if v6out, ok := s.ZoneStats["NotifyOutv6"]; ok {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v6out), "6"))
					}
					return out
				},
			},
			&bindMetric{
				Name:   "bind_server_response_codes",
				Help:   "Response codes sent from the server as a whole",
				Labels: []string{"rcode"},
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return kvExtractor(d, s.ResponseCodeCount)
				},
			},
			&bindMetric{
				Name:   "bind_server_query_types",
				Help:   "Query types sent to the server as a whole",
				Labels: []string{"qtype"},
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return kvExtractor(d, s.QueryTypeCount)
				},
			},
			&bindMetric{
				Name:   "bind_server_stats",
				Help:   "Name server statistics for the server as a whole",
				Labels: []string{"stat"},
				Extract: func(d *prometheus.Desc, s *BindStats) []prometheus.Metric {
					return kvExtractor(d, s.NameserverStats)
				},
			},
		},
		viewMetrics: []*bindViewMetric{
			&bindViewMetric{
				Name:   "bind_view_resolver_stats",
				Help:   "Resolver stats for a view",
				Labels: []string{"view", "stat"},
				Extract: func(d *prometheus.Desc, name string, v *View) []prometheus.Metric {
					out := make([]prometheus.Metric, 0, len(v.Resolver.Stats))
					for k, v := range v.Resolver.Stats {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v), name, k))
					}
					return out
				},
			},
			&bindViewMetric{
				Name:   "bind_view_cache_stats",
				Help:   "Cache stats for a view",
				Labels: []string{"view", "stat"},
				Extract: func(d *prometheus.Desc, name string, v *View) []prometheus.Metric {
					out := make([]prometheus.Metric, 0, len(v.Resolver.CacheStats))
					for k, v := range v.Resolver.CacheStats {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v), name, k))
					}
					return out
				},
			},
			&bindViewMetric{
				Name:   "bind_view_query_count",
				Help:   "Count of queries by qtype in a view",
				Labels: []string{"view", "qtype"},
				Extract: func(d *prometheus.Desc, name string, v *View) []prometheus.Metric {
					out := make([]prometheus.Metric, 0, len(v.Resolver.QueryTypes))
					for k, v := range v.Resolver.QueryTypes {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v), name, k))
					}
					return out
				},
			},
			&bindViewMetric{
				Name:   "bind_view_cache_entry_count",
				Help:   "Cache entry count for a view by RRtype",
				Labels: []string{"view", "rrtype"},
				Extract: func(d *prometheus.Desc, name string, v *View) []prometheus.Metric {
					out := make([]prometheus.Metric, 0, len(v.Resolver.CacheRecordTypes))
					for k, v := range v.Resolver.CacheRecordTypes {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v), name, k))
					}
					return out
				},
			},
		},
		zoneMetrics: []*bindZoneMetric{
			&bindZoneMetric{
				Name:   "bind_zone_serial",
				Help:   "Serial number of zone",
				Labels: []string{"view", "zone", "type"},
				Extract: func(d *prometheus.Desc, viewName string, v *View, z *Zone) []prometheus.Metric {
					return []prometheus.Metric{
						prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(z.Serial), viewName, z.Name, z.Type),
					}
				},
			},
			&bindZoneMetric{
				Name:   "bind_zone_load_time",
				Help:   "Zone load time of a zone",
				Labels: []string{"view", "zone", "type"},
				Extract: func(d *prometheus.Desc, viewName string, v *View, z *Zone) []prometheus.Metric {
					return []prometheus.Metric{
						prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(z.LoadTime.Unix()), viewName, z.Name, z.Type),
					}
				},
			},
			&bindZoneMetric{
				Name:   "bind_zone_response_codes",
				Help:   "Response codes sent from the server for a zone in a view",
				Labels: []string{"view", "zone", "type", "rcode"},
				Extract: func(d *prometheus.Desc, viewName string, v *View, z *Zone) []prometheus.Metric {
					out := make([]prometheus.Metric, 0, len(z.ResponseCodes))
					for k, v := range z.ResponseCodes {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v), viewName, z.Name, z.Type, k))
					}
					return out
				},
			},
			&bindZoneMetric{
				Name:   "bind_zone_query_types",
				Help:   "Query types sent to the server for a zone in a view",
				Labels: []string{"view", "zone", "type", "qtype"},
				Extract: func(d *prometheus.Desc, viewName string, v *View, z *Zone) []prometheus.Metric {
					out := make([]prometheus.Metric, 0, len(z.QueryTypes))
					for k, v := range z.QueryTypes {
						out = append(out, prometheus.MustNewConstMetric(d, prometheus.GaugeValue, float64(v), viewName, z.Name, z.Type, k))
					}
					return out
				},
			},
		},
		socketMetrics: []*bindSocketMetric{
			&bindSocketMetric{
				Name:    "bind_socket_open_count",
				Help:    "Count of sockets opened by protocol and version",
				Suffix:  "Open",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_socket_close_count",
				Help:    "Count of sockets closed by protocol and version",
				Suffix:  "Close",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_socket_bind_fail_count",
				Help:    "Count of sockets bind failures by protocol and version",
				Suffix:  "BindFail",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_socket_connect_count",
				Help:    "Count of sockets connections by protocol and version",
				Suffix:  "Conn",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_tcp_accept_fail_count",
				Help:    "Count of TCP sockets accept failures by version",
				Suffix:  "AcceptFail",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_tcp_connect_fail_count",
				Help:    "Count of sockets connect failures by version",
				Suffix:  "ConnFail",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_accept_count",
				Help:    "Count of sockets accept success by version",
				Suffix:  "Accept",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_socket_recv_error_count",
				Help:    "Count of sockets receive error by protocol and version",
				Suffix:  "RecvErr",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
			&bindSocketMetric{
				Name:    "bind_socket_active_count",
				Help:    "Count of active sockets by protocol and version",
				Suffix:  "Active",
				Labels:  []string{"proto", "version"},
				Extract: socketMetricExtractor,
			},
		},
	}

	for _, m := range c.metrics {
		m.desc = prometheus.NewDesc(m.Name, m.Help, m.Labels, nil)
	}

	for _, m := range c.viewMetrics {
		m.desc = prometheus.NewDesc(m.Name, m.Help, m.Labels, nil)
	}

	for _, m := range c.zoneMetrics {
		m.desc = prometheus.NewDesc(m.Name, m.Help, m.Labels, nil)
	}

	for _, m := range c.socketMetrics {
		m.desc = prometheus.NewDesc(m.Name, m.Help, m.Labels, nil)
	}

	return c
}

func (c *bindCollector) Describe(ch chan<- *prometheus.Desc) {
	for _, m := range c.metrics {
		ch <- m.desc
	}

	for _, m := range c.viewMetrics {
		ch <- m.desc
	}

	for _, m := range c.zoneMetrics {
		ch <- m.desc
	}

	for _, m := range c.socketMetrics {
		ch <- m.desc
	}
}

func (c *bindCollector) Collect(ch chan<- prometheus.Metric) {
	var s BindStats

	req, err := http.NewRequestWithContext(c.ctx, http.MethodGet, c.target, nil)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(
			prometheus.NewDesc("bind_error", "Error building http request", nil, nil),
			fmt.Errorf("Error building http request: %w", err))
		return
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(
			prometheus.NewDesc("bind_error", "Error retrieving bind metrics", nil, nil),
			fmt.Errorf("HTTP error while retrieving bind metrics: %w", err))
		return
	}
	defer res.Body.Close()

	if err := json.NewDecoder(res.Body).Decode(&s); err != nil {
		ch <- prometheus.NewInvalidMetric(
			prometheus.NewDesc("bind_error", "Error decoding bind metrics", nil, nil),
			fmt.Errorf("Error JSON decoding bind metrics: %w", err))
		return
	}

	for _, m := range c.metrics {
		for _, v := range m.Extract(m.desc, &s) {
			ch <- v
		}
	}

	for vn, v := range s.Views {
		for _, m := range c.viewMetrics {
			for _, mv := range m.Extract(m.desc, vn, &v) {
				ch <- mv
			}
		}

		for _, z := range v.Zones {
			for _, m := range c.zoneMetrics {
				for _, mv := range m.Extract(m.desc, vn, &v, &z) {
					ch <- mv
				}
			}
		}
	}

	for _, m := range c.socketMetrics {
		for _, mv := range m.Extract(m.desc, m.Suffix, s.SocketStats) {
			ch <- mv
		}
	}
}