package stats

import (
	"cmp"
	"context"
	"fmt"
	"slices"
	"strconv"
	"strings"
	"time"
)

// Querier answers questions about stored statistics.
type Querier struct {
	Store *Store
	Opts  Options
	Now   func() time.Time
}

// NewQuerier returns a querier for a collector's data.
func NewQuerier(c *Collector) *Querier {
	return &Querier{Store: c.store, Opts: c.opts, Now: c.now}
}

func (q *Querier) loc() *time.Location {
	if q.Opts.Location != nil {
		return q.Opts.Location
	}
	return time.Local
}

// Resolution picks the finest stored resolution that covers a time range
// starting since ago.
func (q *Querier) Resolution(since time.Duration) string {
	switch {
	case since <= 48*time.Hour && since <= q.Opts.RetainMinute:
		return ResMinute
	case since <= q.Opts.RetainHour:
		return ResHour
	}
	return ResDay
}

// topsResolution picks the resolution for top lists, which exist per hour
// and per day.
func (q *Querier) topsResolution(since time.Duration) string {
	if since <= 7*24*time.Hour && since <= q.Opts.RetainHour {
		return ResHour
	}
	return ResDay
}

// bucketStart returns the start of the bucket containing t at a resolution.
func (q *Querier) bucketStart(res string, t time.Time) int64 {
	switch res {
	case ResMinute:
		return t.Unix() - t.Unix()%60
	case ResHour:
		return t.Unix() - t.Unix()%3600
	}
	l := t.In(q.loc())
	return time.Date(l.Year(), l.Month(), l.Day(), 0, 0, 0, 0, q.loc()).Unix()
}

// sums returns the totals of counters (optionally with a key prefix) in
// [from, to).
func (s *Store) sums(ctx context.Context, res string, from, to int64, prefix string) (map[string]int64, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT key, SUM(value) FROM counters WHERE res = ? AND ts >= ? AND ts < ? AND key LIKE ? ESCAPE '\' GROUP BY key`,
		res, from, to, likePrefix(prefix))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[string]int64{}
	for rows.Next() {
		var k string
		var v int64
		if err := rows.Scan(&k, &v); err != nil {
			return nil, err
		}
		out[k] = v
	}
	return out, rows.Err()
}

func likePrefix(p string) string {
	r := strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
	return r.Replace(p) + "%"
}

// series returns the values of a counter per bucket in [from, to).
func (s *Store) series(ctx context.Context, res string, from, to int64, key string) (map[int64]int64, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT ts, value FROM counters WHERE res = ? AND key = ? AND ts >= ? AND ts < ?`, res, key, from, to)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[int64]int64{}
	for rows.Next() {
		var ts, v int64
		if err := rows.Scan(&ts, &v); err != nil {
			return nil, err
		}
		out[ts] = v
	}
	return out, rows.Err()
}

// tops sums top lists (of one or more kinds) in [from, to).
func (s *Store) tops(ctx context.Context, res string, kinds []string, from, to int64, limit int) ([]Entry, error) {
	args := []any{res}
	for _, k := range kinds {
		args = append(args, k)
	}
	args = append(args, from, to, limit)
	rows, err := s.db.QueryContext(ctx, `SELECT key, SUM(count) AS n FROM tops WHERE res = ? AND kind IN (`+strings.TrimSuffix(strings.Repeat("?,", len(kinds)), ",")+`) AND ts >= ? AND ts < ?
		GROUP BY key ORDER BY n DESC, key LIMIT ?`, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []Entry{}
	for rows.Next() {
		var e Entry
		if err := rows.Scan(&e.Key, &e.Count); err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

func sorted(m map[string]int64, prefix string) []Entry {
	out := []Entry{}
	for k, v := range m {
		if strings.HasPrefix(k, prefix) && v != 0 {
			out = append(out, Entry{Key: strings.TrimPrefix(k, prefix), Count: v})
		}
	}
	slices.SortFunc(out, func(a, b Entry) int {
		if a.Count != b.Count {
			return cmp.Compare(b.Count, a.Count)
		}
		return cmp.Compare(a.Key, b.Key)
	})
	return out
}

// Top returns the largest entries of a top list over the last since.
func (q *Querier) Top(ctx context.Context, kind string, since time.Duration, limit int) ([]Entry, error) {
	if !slices.Contains(Kinds, kind) && !quotaKind(kind) {
		return nil, fmt.Errorf("unknown list %q (use %s, or quota_rN)", kind, strings.Join(Kinds, ", "))
	}
	now := q.Now()
	res := q.topsResolution(since)
	return q.Store.tops(ctx, res, []string{kind}, q.bucketStart(res, now.Add(-since)), now.Unix()+1, limit)
}

func quotaKind(k string) bool {
	n, ok := strings.CutPrefix(k, "quota_r")
	_, err := strconv.ParseInt(n, 10, 64)
	return ok && err == nil
}

// Sums returns the totals of the counters with a prefix over the last
// since, largest first, the prefix taken off ("hits:" → "r12").
func (q *Querier) Sums(ctx context.Context, prefix string, since time.Duration) ([]Entry, error) {
	now := q.Now()
	res := q.Resolution(since)
	m, err := q.Store.sums(ctx, res, q.bucketStart(res, now.Add(-since)), now.Unix()+1, prefix)
	if err != nil {
		return nil, err
	}
	return sorted(m, prefix), nil
}

// Traffic is what an interface or zone received and sent, in bytes.
type Traffic struct {
	Name     string `json:"name"`
	Received int64  `json:"received"`
	Sent     int64  `json:"sent"`
}

// Summary is an overview of the last since.
type Summary struct {
	Since      string    `json:"since"`
	From       time.Time `json:"from"`
	Zones      []Traffic `json:"zones"`
	Interfaces []Traffic `json:"interfaces"`
	// Drops are the packets dropped (rules, built-ins, policies), and by
	// chain.
	Drops        int64   `json:"drops"`
	DropsByChain []Entry `json:"drops_by_chain"`
	// Rules are the busiest rules, built-ins and policies, by packets.
	Rules []Entry `json:"rules"`
	// Connections: the most tracked at once (per-bucket averages), and
	// the table's size.
	PeakConnections int64   `json:"peak_connections"`
	ConntrackMax    int64   `json:"conntrack_max"`
	BlockedSources  []Entry `json:"blocked_sources"`
	BlockedPorts    []Entry `json:"blocked_ports"`
	DevicesUp       []Entry `json:"devices_up"`
	DevicesDown     []Entry `json:"devices_down"`
}

// Summary sums up the last since, with top lists of up to topLimit.
func (q *Querier) Summary(ctx context.Context, since time.Duration, topLimit int) (*Summary, error) {
	now := q.Now()
	res := q.Resolution(since)
	from := q.bucketStart(res, now.Add(-since))
	all, err := q.Store.sums(ctx, res, from, now.Unix()+1, "")
	if err != nil {
		return nil, err
	}
	s := &Summary{Since: humanDuration(since), From: time.Unix(from, 0), Zones: []Traffic{}, Interfaces: []Traffic{}}
	traffic := func(rx, tx string) []Traffic {
		by := map[string]*Traffic{}
		for k, v := range all {
			var name string
			var in bool
			if n, ok := strings.CutPrefix(k, rx); ok {
				name, in = n, true
			} else if n, ok := strings.CutPrefix(k, tx); ok {
				name = n
			} else {
				continue
			}
			t := by[name]
			if t == nil {
				t = &Traffic{Name: name}
				by[name] = t
			}
			if in {
				t.Received += v
			} else {
				t.Sent += v
			}
		}
		out := []Traffic{}
		for _, t := range by {
			out = append(out, *t)
		}
		slices.SortFunc(out, func(a, b Traffic) int {
			if c := cmp.Compare(b.Received+b.Sent, a.Received+a.Sent); c != 0 {
				return c
			}
			return cmp.Compare(a.Name, b.Name)
		})
		return out
	}
	s.Zones, s.Interfaces = traffic("zrx:", "ztx:"), traffic("rx:", "tx:")
	s.Drops = all["drops"]
	s.DropsByChain = sorted(all, "drop:")
	s.Rules = sorted(all, "hits:")
	if len(s.Rules) > topLimit {
		s.Rules = s.Rules[:topLimit]
	}
	if tl, err := q.Timeline(ctx, since, "auto", []string{"g:" + GaugeConntrack, "g:" + GaugeConntrackMax}); err == nil {
		for _, v := range tl.Series[0].Values {
			s.PeakConnections = max(s.PeakConnections, v)
		}
		for _, v := range tl.Series[1].Values {
			s.ConntrackMax = max(s.ConntrackMax, v)
		}
	}
	for _, l := range []struct {
		kind string
		dst  *[]Entry
	}{{KindBlockedSource, &s.BlockedSources}, {KindBlockedPort, &s.BlockedPorts}, {KindDeviceUp, &s.DevicesUp}, {KindDeviceDown, &s.DevicesDown}} {
		es, err := q.Top(ctx, l.kind, since, topLimit)
		if err != nil {
			return nil, err
		}
		*l.dst = es
	}
	return s, nil
}

// Timeline is a set of series over time.
type Timeline struct {
	Step   string    `json:"step"`
	Times  []int64   `json:"times"` // bucket starts (Unix seconds)
	Series []*Series `json:"series"`
}

// Series is one counter over time. For a gauge (g:KEY) the values are
// averages, and -1 where there is no sample.
type Series struct {
	Key    string  `json:"key"`
	Values []int64 `json:"values"`
}

var steps = []struct {
	name string
	d    time.Duration
	res  string
}{
	{"1m", time.Minute, ResMinute}, {"5m", 5 * time.Minute, ResMinute}, {"15m", 15 * time.Minute, ResMinute},
	{"1h", time.Hour, ResHour}, {"6h", 6 * time.Hour, ResHour}, {"1d", 24 * time.Hour, ResDay},
}

// Timeline returns counters over the last since. keys are counter names;
// a key ending in ":*" means every counter of that group. step is "auto"
// or one of 1m, 5m, 15m, 1h, 6h, 1d.
func (q *Querier) Timeline(ctx context.Context, since time.Duration, step string, keys []string) (*Timeline, error) {
	idx := -1
	for i, st := range steps {
		ok := st.res == ResMinute && since <= q.Opts.RetainMinute || st.res == ResHour && since <= q.Opts.RetainHour || st.res == ResDay
		if step == st.name {
			if !ok {
				return nil, fmt.Errorf("step %s needs %s data, which is only kept for %s", st.name, map[string]string{ResMinute: "per-minute", ResHour: "hourly"}[st.res], humanDuration(map[string]time.Duration{ResMinute: q.Opts.RetainMinute, ResHour: q.Opts.RetainHour}[st.res]))
			}
			idx = i
			break
		}
		if step == "auto" || step == "" {
			if ok && since/st.d <= 300 {
				idx = i
				break
			}
		}
	}
	if idx < 0 {
		if step != "auto" && step != "" {
			return nil, fmt.Errorf("unknown step %q (use auto, 1m, 5m, 15m, 1h, 6h or 1d)", step)
		}
		idx = len(steps) - 1
	}
	st := steps[idx]
	now := q.Now()
	// Buckets aligned to the step (days to local midnight).
	var starts []int64
	if st.res == ResDay {
		for t := time.Unix(q.bucketStart(ResDay, now.Add(-since)), 0).In(q.loc()); !t.After(now); t = t.AddDate(0, 0, 1) {
			starts = append(starts, t.Unix())
		}
	} else {
		sd := int64(st.d.Seconds())
		first := now.Add(-since).Unix()
		first -= first % sd
		for t := first; t <= now.Unix(); t += sd {
			starts = append(starts, t)
		}
	}
	from, to := starts[0], now.Unix()+1
	// Expand group keys.
	var expanded []string
	for _, k := range keys {
		if g, ok := strings.CutSuffix(k, ":*"); ok {
			m, err := q.Store.sums(ctx, st.res, from, to, g+":")
			if err != nil {
				return nil, err
			}
			for _, e := range sorted(m, "") {
				expanded = append(expanded, e.Key)
			}
			continue
		}
		expanded = append(expanded, k)
	}
	tl := &Timeline{Step: st.name, Times: starts, Series: []*Series{}}
	bucketed := func(key string) ([]int64, error) {
		raw, err := q.Store.series(ctx, st.res, from, to, key)
		if err != nil {
			return nil, err
		}
		vals := make([]int64, len(starts))
		for ts, v := range raw {
			// The bucket is the last start <= ts.
			i, found := slices.BinarySearch(starts, ts)
			if !found {
				i--
			}
			if i >= 0 && i < len(vals) {
				vals[i] += v
			}
		}
		return vals, nil
	}
	for _, k := range expanded {
		vals, err := bucketed(k)
		if err != nil {
			return nil, err
		}
		// A gauge (g:KEY) is the average of its samples: its sum divided
		// by gn:KEY, their number. A bucket without samples is -1.
		if g, ok := strings.CutPrefix(k, "g:"); ok {
			n, err := bucketed("gn:" + g)
			if err != nil {
				return nil, err
			}
			for i := range vals {
				if n[i] > 0 {
					vals[i] /= n[i]
				} else {
					vals[i] = -1
				}
			}
		}
		tl.Series = append(tl.Series, &Series{Key: k, Values: vals})
	}
	return tl, nil
}

// Heatmap is average activity per weekday and hour of day (local time).
type Heatmap struct {
	Key string `json:"key"`
	// Avg[d][h] is the average per hour on weekday d (0 = Monday) at
	// hour h; -1 where there is no data.
	Avg [7][24]float64 `json:"avg"`
	// Samples[d][h] is how many hours the average is based on.
	Samples [7][24]int `json:"samples"`
	From    time.Time  `json:"from"`
}

// Heatmap averages a counter (default "drops") by weekday and hour over
// the last since (at most the hourly retention).
func (q *Querier) Heatmap(ctx context.Context, key string, since time.Duration) (*Heatmap, error) {
	if key == "" {
		key = "drops"
	}
	since = min(since, q.Opts.RetainHour)
	now := q.Now()
	from := q.bucketStart(ResHour, now.Add(-since))
	raw, err := q.Store.series(ctx, ResHour, from, now.Unix()+1, key)
	if err != nil {
		return nil, err
	}
	hm := &Heatmap{Key: key, From: time.Unix(from, 0)}
	var sum [7][24]int64
	for ts, v := range raw {
		t := time.Unix(ts, 0).In(q.loc())
		d := (int(t.Weekday()) + 6) % 7
		sum[d][t.Hour()] += v
		hm.Samples[d][t.Hour()]++
	}
	for d := range 7 {
		for h := range 24 {
			if n := hm.Samples[d][h]; n > 0 {
				hm.Avg[d][h] = float64(sum[d][h]) / float64(n)
			} else {
				hm.Avg[d][h] = -1
			}
		}
	}
	return hm, nil
}

// ParseSince parses a range such as 90m, 24h, 7d or 2w.
func ParseSince(s string) (time.Duration, error) {
	if s == "" {
		return 24 * time.Hour, nil
	}
	mult := time.Duration(0)
	switch {
	case strings.HasSuffix(s, "d"):
		mult = 24 * time.Hour
	case strings.HasSuffix(s, "w"):
		mult = 7 * 24 * time.Hour
	}
	if mult > 0 {
		var n int
		if _, err := fmt.Sscanf(s[:len(s)-1], "%d", &n); err != nil || n <= 0 || fmt.Sprint(n) != s[:len(s)-1] {
			return 0, fmt.Errorf("invalid range %q (use e.g. 90m, 24h, 7d, 2w)", s)
		}
		return time.Duration(n) * mult, nil
	}
	d, err := time.ParseDuration(s)
	if err != nil || d <= 0 {
		return 0, fmt.Errorf("invalid range %q (use e.g. 90m, 24h, 7d, 2w)", s)
	}
	return d, nil
}

func humanDuration(d time.Duration) string {
	if d%(24*time.Hour) == 0 && d > 0 {
		return strconv.Itoa(int(d/(24*time.Hour))) + "d"
	}
	return d.String()
}