package server

import (
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"errors"
	"fmt"
	"log/slog"
	"math"
	"net"
	"net/http"
	"net/netip"
	"slices"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"jarl/internal/dnssec"
	"jarl/internal/engine"
	"jarl/internal/model"
	"jarl/internal/render"
	"jarl/internal/stats"
	"jarl/internal/validate"
)

// Metrics configures the Prometheus endpoint /metrics (docs/metrics.md).
type Metrics struct {
	Enabled bool
	// Listen are extra addresses serving only /metrics (plain HTTP).
	Listen []string
	// Allow may scrape without a token; Token, when set, lets anyone who
	// sends it scrape.
	Allow []netip.Prefix
	Token string
}

// metricsStateTTL is how long the Jarl state part of /metrics is reused:
// it runs rndc status, validation and the drift check, too much for every
// scrape when several Prometheus servers scrape often.
const metricsStateTTL = 15 * time.Second

// metricsCache holds the last Jarl state part of /metrics.
type metricsCache struct {
	mu   sync.Mutex
	at   time.Time
	body []byte
}

const metricsContentType = "text/plain; version=0.0.4; charset=utf-8"

// metrics serves GET /metrics in Prometheus's text format.
func (s *Server) metrics(w http.ResponseWriter, r *http.Request) {
	if !s.cfg.Metrics.Enabled {
		http.Error(w, "metrics are disabled; set metrics.enabled: true in /etc/jarl/jarl.yaml and restart jarl (see docs/metrics.md)", http.StatusNotFound)
		return
	}
	if code, msg := s.metricsAccess(r); code != 0 {
		if code == http.StatusUnauthorized {
			w.Header().Set("WWW-Authenticate", `Bearer realm="jarl metrics"`)
		}
		http.Error(w, msg, code)
		return
	}
	var buf bytes.Buffer
	buf.Write(s.stateMetrics(r.Context()))
	if c := s.cfg.Stats; c != nil {
		pw := &promWriter{}
		statsMetrics(pw, c.Totals(), c.Server(), c.Status())
		pw.writeTo(&buf)
	}
	w.Header().Set("Content-Type", metricsContentType)
	w.Header().Set("Cache-Control", "no-store")
	w.Write(buf.Bytes())
}

// metricsAccess decides whether a scrape is allowed: from an address in
// metrics.allow, with metrics.token, or on the local socket from root or
// jarld's own user. It returns 0 when allowed, or the status and message
// to refuse with.
func (s *Server) metricsAccess(r *http.Request) (int, string) {
	m := s.cfg.Metrics
	ip := s.clientIP(r)
	if h := r.Header.Get("Authorization"); h != "" {
		scheme, tok, _ := strings.Cut(h, " ")
		if m.Token == "" || !strings.EqualFold(scheme, "Bearer") {
			return http.StatusUnauthorized, "this server takes no token for /metrics (metrics.token is not set)"
		}
		if s.limiter.blocked("ip:"+ip, failsPerIP) {
			return http.StatusTooManyRequests, "too many failed attempts; try again later"
		}
		got, want := sha256.Sum256([]byte(strings.TrimSpace(tok))), sha256.Sum256([]byte(m.Token))
		if subtle.ConstantTimeCompare(got[:], want[:]) != 1 {
			s.limiter.fail("ip:" + ip)
			s.log.Warn("metrics: wrong token", "ip", ip)
			return http.StatusUnauthorized, "wrong token for /metrics"
		}
		return 0, ""
	}
	if pr := peerOf(r.Context()); pr != nil {
		if trustedUID(pr.UID) {
			return 0, ""
		}
		return http.StatusForbidden, "on the local socket, only root and jarld's own user may read /metrics without metrics.token"
	}
	if a, err := netip.ParseAddr(ip); err == nil {
		a = a.Unmap().WithZone("")
		for _, p := range m.Allow {
			if p.Contains(a) {
				return 0, ""
			}
		}
	}
	if m.Token != "" {
		return http.StatusUnauthorized, "send metrics.token as 'Authorization: Bearer TOKEN', or add " + ip + " to metrics.allow in jarl.yaml"
	}
	return http.StatusForbidden, ip + " may not read /metrics: add it to metrics.allow in jarl.yaml (or set metrics.token)"
}

// serveMetricsListeners serves /metrics alone on metrics.listen.
func (s *Server) serveMetricsListeners(errc chan error) error {
	if !s.cfg.Metrics.Enabled || len(s.cfg.Metrics.Listen) == 0 {
		return nil
	}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", s.metrics)
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "this address serves only /metrics", http.StatusNotFound)
	})
	h := s.recoverer(s.logRequests(mux))
	for _, addr := range s.cfg.Metrics.Listen {
		ln, err := net.Listen("tcp", addr)
		if err != nil {
			return fmt.Errorf("metrics.listen %s: %w", addr, err)
		}
		srv := &http.Server{Handler: h, ReadHeaderTimeout: 10 * time.Second, WriteTimeout: time.Minute, IdleTimeout: 2 * time.Minute,
			ErrorLog: slog.NewLogLogger(s.log.Handler(), slog.LevelWarn)}
		s.mu.Lock()
		s.srvs = append(s.srvs, srv)
		s.mu.Unlock()
		s.log.Info("serving metrics", "url", "http://"+ln.Addr().String()+"/metrics")
		go func() {
			if err := srv.Serve(ln); !errors.Is(err, http.ErrServerClosed) {
				errc <- err
			}
		}()
	}
	return nil
}

// stateMetrics renders Jarl's own state (cached for metricsStateTTL).
func (s *Server) stateMetrics(ctx context.Context) []byte {
	c := &s.metricsCache
	c.mu.Lock()
	defer c.mu.Unlock()
	now := time.Now()
	if c.body != nil && now.Sub(c.at) < metricsStateTTL {
		return c.body
	}
	pw := &promWriter{}
	s.collectState(ctx, pw, now)
	var buf bytes.Buffer
	pw.writeTo(&buf)
	c.body, c.at = buf.Bytes(), now
	return c.body
}

func (s *Server) collectState(ctx context.Context, pw *promWriter, now time.Time) {
	pw.gauge("jarl_build_info", "Jarl's version (always 1).", 1, "version", s.cfg.Version)
	pw.gauge("jarl_start_time_seconds", "When jarld started (Unix time).", unix(s.started))
	ok := 1.0
	defer func() {
		pw.gauge("jarl_metrics_state_ok", "1 when Jarl's state could be read for this scrape.", ok)
	}()

	st, err := s.eng.Status(ctx)
	if err != nil {
		s.log.Warn("metrics: reading the status failed", "err", err)
		ok = 0
		return
	}
	pw.gauge("jarl_adopted", "1 when named loads Jarl's configuration.", b2f(st.Adopted))
	if st.Named != nil || st.NamedError != "" {
		pw.gauge("jarl_named_up", "1 when named answers rndc status.", b2f(st.Named != nil && st.Named.Running))
	}
	if n := st.Named; n != nil {
		if v := fractions(n.RecursiveClients); len(v) == 3 {
			pw.gauge("jarl_named_recursive_clients", "Recursive clients named is serving now (rndc status).", v[0])
			pw.gauge("jarl_named_recursive_clients_soft_limit", "Recursive clients at which named starts dropping the oldest queries.", v[1])
			pw.gauge("jarl_named_recursive_clients_limit", "Most recursive clients named serves at once (recursive-clients).", v[2])
		}
		if v := fractions(n.TCPClients); len(v) == 2 {
			pw.gauge("jarl_named_tcp_clients", "TCP clients named is serving now (rndc status).", v[0])
			pw.gauge("jarl_named_tcp_clients_limit", "Most TCP clients named serves at once (tcp-clients).", v[1])
		}
	}
	pw.gauge("jarl_open_changesets", "Changesets with staged changes that are not applied yet.", float64(st.OpenChangesets))
	if st.LastApplied != nil && st.LastApplied.AppliedAt != nil {
		pw.gauge("jarl_last_apply_timestamp_seconds", "When the last changeset was applied (Unix time).", unix(*st.LastApplied.AppliedAt))
	}
	sev := map[validate.Severity]int{validate.Error: 0, validate.Warning: 0, validate.Info: 0}
	for _, p := range st.Problems {
		sev[p.Severity]++
	}
	for _, k := range []validate.Severity{validate.Error, validate.Warning, validate.Info} {
		pw.gauge("jarl_problems", "Validation problems in the current configuration, as on the dashboard.", float64(sev[k]), "severity", string(k))
	}
	pw.gauge("jarl_drifted_files", "Files Jarl manages that were changed or removed outside Jarl.", float64(len(st.Drift)))
	for _, k := range sortedKeys(st.Counts) {
		pw.gauge("jarl_objects", "Objects in Jarl's configuration.", float64(st.Counts[k]), "kind", k)
	}

	if s.cfg.BackupDir != "" {
		bs := s.backupStatus(ctx, now)
		if bs.LastAt != nil {
			pw.gauge("jarl_backup_last_success_timestamp_seconds", "When the newest backup was made by 'jarld backup' (Unix time).", unix(*bs.LastAt))
			pw.gauge("jarl_backup_last_size_bytes", "Size of the newest backup.", float64(bs.LastSize))
		}
		if bs.LastErrorAt != nil {
			pw.gauge("jarl_backup_last_failure_timestamp_seconds", "When a backup last failed, if after the newest backup (Unix time).", unix(*bs.LastErrorAt))
		}
		if s.cfg.BackupWarnAfter > 0 {
			pw.gauge("jarl_backup_stale", "1 when the newest backup is older than backup.warn_after.", b2f(bs.Stale))
		}
	}

	state, err := s.store.LoadState(ctx)
	if err != nil {
		ok = 0
		return
	}
	if len(state.Blocklists) > 0 || len(state.BlockRules) > 0 {
		pw.gauge("jarl_blocking_paused", "1 while blocking is paused.", b2f(state.Settings.BlockingPaused()))
		if t := state.Settings.BlockingPausedUntil; t != nil {
			pw.gauge("jarl_blocking_paused_until_timestamp_seconds", "When paused blocking resumes (Unix time).", unix(*t))
		}
	}
	statuses, _ := s.store.BlocklistStatuses(ctx)
	for _, b := range state.SortedBlocklists() {
		pw.gauge("jarl_blocklist_enabled", "1 for blocklists that are switched on.", b2f(!b.Disabled), "list", b.Name)
		bs := statuses[b.ID]
		if bs == nil {
			continue
		}
		pw.gauge("jarl_blocklist_entries", "Domains in the blocklist's last good download.", float64(bs.Entries), "list", b.Name)
		pw.gauge("jarl_blocklist_failing", "1 when the last download or check of an enabled blocklist failed.", b2f(bs.Error != "" && !b.Disabled), "list", b.Name)
		if bs.CheckedAt != nil {
			pw.gauge("jarl_blocklist_last_check_timestamp_seconds", "When the blocklist was last checked for changes (Unix time).", unix(*bs.CheckedAt), "list", b.Name)
		}
		if bs.UpdatedAt != nil {
			pw.gauge("jarl_blocklist_last_change_timestamp_seconds", "When the blocklist's content last changed (Unix time).", unix(*bs.UpdatedAt), "list", b.Name)
		}
	}

	zones := make([]*model.Zone, 0, len(state.Zones))
	for _, z := range state.Zones {
		if z.DNSSECPolicy != "" && !z.Disabled {
			zones = append(zones, z)
		}
	}
	slices.SortFunc(zones, func(a, b *model.Zone) int {
		return strings.Compare(a.Name+" "+state.ViewName(a), b.Name+" "+state.ViewName(b))
	})
	for _, z := range zones {
		keys, err := dnssec.Keys(render.KeyDir(state, z, st.Paths), z.Name)
		if err != nil {
			continue
		}
		ph := dnssec.Assess(z.DNSSECPolicy, keys)
		pw.gauge("jarl_dnssec_phase", "The DNSSEC phase of each signed zone (always 1; see docs/dnssec.md). ready: add the DS at the parent; removing_ds: remove it.", 1,
			"zone", engine.Display(z.Name), "view", state.ViewName(z), "phase", ph.Code)
	}

	if tokens, err := s.store.ListTokens(ctx, 0); err == nil {
		for _, t := range tokens {
			if t.ExpiresAt != nil && t.RevokedAt == nil {
				pw.gauge("jarl_token_expiry_timestamp_seconds", "When each API token with an expiry date expires (Unix time).", unix(*t.ExpiresAt),
					"id", strconv.FormatInt(t.ID, 10), "user", t.Username, "token", t.Name)
			}
		}
	}
}

// counterMetrics maps the statistics counters (docs/stats.md) to metrics:
// a key "PREFIX:VALUE" becomes metric{label="VALUE"}, a key without a
// colon a metric without labels.
var counterMetrics = map[string]struct{ name, label, help string }{
	"queries":   {"jarl_queries_total", "", "Queries named received (statistics channel)."},
	"recursion": {"jarl_recursion_queries_total", "", "Queries that caused recursion."},
	"dropped":   {"jarl_dropped_queries_total", "", "Queries named dropped."},
	"duplicate": {"jarl_duplicate_queries_total", "", "Duplicate queries named ignored."},
	"rcode":     {"jarl_responses_by_rcode_total", "rcode", "Responses by response code."},
	"qtype":     {"jarl_queries_by_type_total", "type", "Queries by record type."},
	"answer":    {"jarl_answers_total", "source", "Answers that were authoritative or not (recursive)."},
	"transport": {"jarl_queries_by_transport_total", "transport", "Queries over UDP and TCP."},
	"ratelimit": {"jarl_ratelimit_responses_total", "action", "Responses rate limiting dropped or slipped (truncated)."},
	"rtt":       {"jarl_upstream_responses_by_rtt_total", "rtt", "Responses from other servers during recursion, by round-trip time."},
	"dnssec":    {"jarl_dnssec_validations_total", "result", "DNSSEC validations of recursive answers by result."},
	"cache":     {"jarl_cache_lookups_total", "result", "Cache lookups that hit or missed."},
	"responses": {"jarl_logged_responses_total", "", "Responses read from named's response log (without stats.exclude_clients)."},
	"view":      {"jarl_responses_by_view_total", "view", "Logged responses by view."},
	"zone":      {"jarl_responses_by_zone_total", "zone", "Logged responses by local zone (\"-\": names in no local zone)."},
	"blocked":   {"jarl_blocked_queries_total", "", "Queries blocking answered."},
	"blocklist": {"jarl_blocked_by_list_total", "list", "Blocked queries by blocklist (\"rules\": Jarl's own rules)."},
}

// statsMetrics adds the query statistics: counters since jarld started,
// and named's process and server gauges.
func statsMetrics(pw *promWriter, totals map[string]int64, srv *stats.Server, st stats.Status) {
	if _, ok := totals["queries"]; !ok {
		totals["queries"] = 0
	}
	for _, k := range sortedKeys(totals) {
		prefix, value, labeled := strings.Cut(k, ":")
		m, ok := counterMetrics[prefix]
		if !ok || labeled != (m.label != "") {
			continue
		}
		if labeled {
			pw.counter(m.name, m.help, float64(totals[k]), m.label, value)
		} else {
			pw.counter(m.name, m.help, float64(totals[k]))
		}
	}
	pw.gauge("jarl_stats_poll_ok", "1 when the last poll of named's statistics channel worked.", b2f(st.PollError == ""))
	if st.LastPoll != nil {
		pw.gauge("jarl_stats_last_poll_timestamp_seconds", "When named's statistics channel was last read (Unix time).", unix(*st.LastPoll))
	}
	if st.QueryLog {
		pw.counter("jarl_stats_log_lines_total", "Lines read from named's response log.", float64(st.LinesRead))
		if st.LastLine != nil {
			pw.gauge("jarl_stats_last_log_line_timestamp_seconds", "Time of the last response log line read (Unix time).", unix(*st.LastLine))
		}
	}
	pw.gauge("jarl_stats_warnings", "Warnings about statistics collection, as on the Stats page.", float64(len(st.Warnings)))
	if srv == nil {
		return
	}
	if t, err := time.Parse(time.RFC3339Nano, srv.BootTime); err == nil {
		pw.gauge("jarl_named_start_time_seconds", "When named started (Unix time).", unix(t))
	}
	if t, err := time.Parse(time.RFC3339Nano, srv.ConfigTime); err == nil {
		pw.gauge("jarl_named_config_time_seconds", "When named last loaded its configuration (Unix time).", unix(t))
	}
	if srv.PID != 0 {
		pw.gauge("jarl_named_resident_memory_bytes", "named's resident memory (/proc).", float64(srv.RSS))
		pw.gauge("jarl_named_threads", "named's threads (/proc).", float64(srv.Threads))
		pw.counter("jarl_named_cpu_seconds_total", "CPU time named used since it started (/proc).", srv.CPUSeconds)
	}
	if srv.Version != "" {
		pw.gauge("jarl_named_memory_in_use_bytes", "Memory named counts as in use (statistics channel).", float64(srv.MemInUse))
		pw.gauge("jarl_named_cache_rrsets", "RRsets in named's caches, all views.", float64(srv.CacheRRsets))
		pw.gauge("jarl_named_cache_memory_bytes", "Memory named's caches use, all views.", float64(srv.CacheMem))
		pw.gauge("jarl_named_sockets", "Open sockets.", float64(srv.SocketsUDP), "proto", "udp")
		pw.gauge("jarl_named_sockets", "Open sockets.", float64(srv.SocketsTCP), "proto", "tcp")
		pw.gauge("jarl_named_recursive_clients_highwater", "Most recursive clients at once since named started.", float64(srv.RecursionHighwater))
		pw.gauge("jarl_named_tcp_clients_highwater", "Most TCP clients at once since named started.", float64(srv.TCPHighwater))
	}
}

// promWriter collects metrics and writes them in Prometheus's text format,
// each metric's samples together under one HELP and TYPE.
type promWriter struct {
	order []string
	fams  map[string]*promFamily
}

type promFamily struct {
	typ, help string
	lines     []string
}

func (w *promWriter) gauge(name, help string, v float64, labels ...string) {
	w.add(name, "gauge", help, v, labels)
}

func (w *promWriter) counter(name, help string, v float64, labels ...string) {
	w.add(name, "counter", help, v, labels)
}

func (w *promWriter) add(name, typ, help string, v float64, labels []string) {
	if w.fams == nil {
		w.fams = map[string]*promFamily{}
	}
	f := w.fams[name]
	if f == nil {
		f = &promFamily{typ: typ, help: help}
		w.fams[name] = f
		w.order = append(w.order, name)
	}
	var b strings.Builder
	b.WriteString(name)
	if len(labels) > 0 {
		b.WriteByte('{')
		for i := 0; i+1 < len(labels); i += 2 {
			if i > 0 {
				b.WriteByte(',')
			}
			b.WriteString(labels[i])
			b.WriteString(`="`)
			b.WriteString(labelEscaper.Replace(labels[i+1]))
			b.WriteByte('"')
		}
		b.WriteByte('}')
	}
	b.WriteByte(' ')
	b.WriteString(formatFloat(v))
	f.lines = append(f.lines, b.String())
}

func (w *promWriter) writeTo(buf *bytes.Buffer) {
	for _, name := range w.order {
		f := w.fams[name]
		fmt.Fprintf(buf, "# HELP %s %s\n# TYPE %s %s\n", name, helpEscaper.Replace(f.help), name, f.typ)
		for _, l := range f.lines {
			buf.WriteString(l)
			buf.WriteByte('\n')
		}
	}
}

var (
	labelEscaper = strings.NewReplacer(`\`, `\\`, `"`, `\"`, "\n", `\n`)
	helpEscaper  = strings.NewReplacer(`\`, `\\`, "\n", `\n`)
)

func formatFloat(v float64) string {
	switch {
	case math.IsInf(v, 1):
		return "+Inf"
	case math.IsInf(v, -1):
		return "-Inf"
	case math.IsNaN(v):
		return "NaN"
	}
	return strconv.FormatFloat(v, 'f', -1, 64)
}

func unix(t time.Time) float64 { return float64(t.UnixMilli()) / 1000 }

func b2f(b bool) float64 {
	if b {
		return 1
	}
	return 0
}

// fractions parses rndc status values such as "0/900/1000".
func fractions(s string) []float64 {
	if s == "" {
		return nil
	}
	var out []float64
	for _, p := range strings.Split(s, "/") {
		v, err := strconv.ParseFloat(strings.TrimSpace(p), 64)
		if err != nil {
			return nil
		}
		out = append(out, v)
	}
	return out
}

func sortedKeys[V any](m map[string]V) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}