package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

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

type BindStats struct {
	Version            string                 `json:"json-stats-version"`
	BindVersion        string                 `json:"version"`
	BootTime           time.Time              `json:"boot-time"`
	ConfigTime         time.Time              `json:"config-time"`
	OperationCodeCount map[string]int         `json:"opcodes"`
	ResponseCodeCount  map[string]int         `json:"rcodes"`
	QueryTypeCount     map[string]int         `json:"qtypes"`
	NameserverStats    map[string]int         `json:"nsstats"`
	ZoneStats          map[string]int         `json:"zonestats"`
	MemoryStats        MemoryStats            `json:"memory"`
	ResolverStats      struct{ Mismatch int } `json:"resstats"`
	SocketStats        map[string]int         `json:"sockstats"`
	TrafficHistograms  map[string]Histogram   `json:"traffic"`
	Views              map[string]View        `json:"views"`
}

type Histogram struct {
	Buckets []HistogramBucket
}

func (h *Histogram) UnmarshalJSON(v []byte) error {
	mv := map[string]int{}

	if err := json.Unmarshal(v, &mv); err != nil {
		return err
	}

	if h.Buckets == nil {
		h.Buckets = make([]HistogramBucket, 0, len(mv))
	}

	for k, v := range mv {
		var min, max int

		if strings.HasSuffix(k, "+") {
			min, _ = strconv.Atoi(k[:len(k)-1])
			max = 0
		} else {
			mm := strings.Split(k, "-")
			min, _ = strconv.Atoi(mm[0])
			max, _ = strconv.Atoi(mm[1])
		}
		h.Buckets = append(h.Buckets, HistogramBucket{
			Min:          min,
			Max:          max,
			Observations: v,
		})
	}

	return nil
}

type HistogramBucket struct {
	Min          int
	Max          int
	Observations int
}

type View struct {
	Zones    []Zone            `json:"zones"`
	Resolver ViewResolverStats `json:"resolver"`
}

type Zone struct {
	Name          string         `json:"name"`
	Class         string         `json:"class"`
	Serial        int            `json:"serial"`
	Type          string         `json:"type"`
	LoadTime      time.Time      `json:"loaded"`
	ResponseCodes map[string]int `json:"rcodes"`
	QueryTypes    map[string]int `json:"qtypes"`
}

type ViewResolverStats struct {
	Stats            map[string]int `json:"stats"`
	CacheStats       map[string]int `json:"cachestats"`
	QueryTypes       map[string]int `json:"qtypes"`
	CacheRecordTypes map[string]int `json:"cache"`
	ADB              struct {
		NumberOfEntries int `json:"nentries"`
		EntryCount      int `json:"entriescnt"`
		NumberOfNames   int `json:"nnames"`
		NamesCount      int `json:"namescnt"`
	} `json:"adb"`
}

type MemoryStats struct {
	TotalUse    int `json:"TotalUse"`
	InUse       int `json:"InUse"`
	Malloced    int `json:"Malloced"`
	BlockSize   int `json:"BlockSize"`
	ContextSize int `json:"ContextSize"`
	Lost        int `json:"Lost"`
}

func main() {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		target := r.URL.Query().Get("target")
		if target == "" {
			http.Error(w, "'target' must be specified exactly once", http.StatusBadRequest)
			return
		}

		u, err := url.Parse(fmt.Sprintf("http://%s", target))
		if err != nil {
			http.Error(w, "unable to parse target URL", http.StatusBadRequest)
			return
		}

		u.Scheme = "http"
		u.Path = "/json"

		ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
		defer cancel()

		registry := prometheus.NewRegistry()
		c := NewBindCollector(ctx, u.String())
		registry.MustRegister(c)
		h := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
		h.ServeHTTP(w, r)
	})

	http.ListenAndServe(":8080", nil)
}