package stats

import (
	"cmp"
	"context"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	"housecarl/internal/api"
	"housecarl/internal/nft"
)

// Resolutions and retention, and the sources the figures come from.
type Options struct {
	TopN         int
	RetainMinute time.Duration
	RetainHour   time.Duration
	RetainDay    time.Duration
	Location     *time.Location // for days (default: local time)
	// Interval is how often the collector samples (default: a minute).
	Interval time.Duration
}

// Sources are what the collector reads. Nil functions are left out.
type Sources struct {
	// Table reads Housecarl's table: the counters by mark (totals across
	// reloads), the quotas, the per-device sets.
	Table func(context.Context) (*nft.Table, error)
	// Zones maps the interfaces to their zones.
	Zones func() map[string]string
	// Classify says which chain a mark's counter is in and whether what
	// it counts is dropped ("r12" → forward, true).
	Classify func(mark string) (chain string, drop bool)
	// SysNet and Proc are /sys/class/net and /proc (tests point them
	// elsewhere).
	SysNet, Proc string
}

// Counter keys. Traffic: rx:IF, tx:IF (bytes), rxp:IF, txp:IF (packets),
// zrx:ZONE, ztx:ZONE (bytes received from and sent to a zone's
// interfaces). Hits: hits:MARK (packets), hitb:MARK (bytes). Drops:
// drop:CHAIN and drops (packets). Quotas: quota:r12 (bytes). Gauges
// (g:KEY, see gauge) are below.
const (
	GaugeConntrack    = "conntrack"
	GaugeConntrackMax = "conntrack_max"
	GaugeRSS          = "rss"  // housecarld's resident memory, bytes
	GaugeCPU          = "cpu"  // housecarld's CPU, tenths of a percent of one CPU
	GaugeMem          = "mem"  // the host's memory in use, bytes
	GaugeLoad         = "load" // the host's 1-minute load average, times 100
)

// Top-list kinds: per hour and per day, the TopN largest.
const (
	KindBlockedSource = "blocked_src"  // logged drops by source address
	KindBlockedPort   = "blocked_port" // logged drops by what they tried: tcp/22, icmp
	KindDeviceUp      = "device_up"    // bytes sent to the internet, by inside address
	KindDeviceDown    = "device_down"  // bytes received from it
	// A per-address quota's use by address is "quota_r12".
)

// Kinds are the top lists Top accepts (and quota_rN).
var Kinds = []string{KindBlockedSource, KindBlockedPort, KindDeviceUp, KindDeviceDown}

// Status describes the collector's health.
type Status struct {
	Enabled     bool       `json:"enabled"`
	Database    string     `json:"database,omitempty"`
	Size        int64      `json:"size"`
	LastSample  *time.Time `json:"last_sample,omitempty"`
	SampleError string     `json:"sample_error,omitempty"`
	FirstData   *time.Time `json:"first_data,omitempty"`
	// PerDevice: the firewall counts traffic per device (a setting).
	PerDevice bool `json:"per_device"`
	Retention struct {
		Minute string `json:"minute"`
		Hour   string `json:"hour"`
		Day    string `json:"day"`
	} `json:"retention"`
}

// Latest are the newest gauge values.
type Latest struct {
	At           time.Time `json:"at"`
	Conntrack    int64     `json:"conntrack"`
	ConntrackMax int64     `json:"conntrack_max"`
	RSS          int64     `json:"rss"`
	CPUPermille  int64     `json:"cpu_permille"`
	MemUsed      int64     `json:"mem_used"`
	Load         float64   `json:"load"`
}

// Collector gathers statistics into a Store.
type Collector struct {
	opts  Options
	src   Sources
	store *Store
	log   *slog.Logger
	now   func() time.Time

	// Each source's last totals: the next sample stores the difference.
	ifs, table deltas
	cpuTicks   int64
	cpuAt      time.Time
	lastPrune  time.Time

	mu     sync.Mutex
	hours  map[int64]map[string]map[string]int64 // live hours: hour -> kind -> key -> count
	status Status
	latest *Latest
}

// maxKeys bounds each kind of an hour's top list in memory.
const maxKeys = 100_000

// NewCollector creates a collector. Run starts it.
func NewCollector(store *Store, opts Options, src Sources, log *slog.Logger) *Collector {
	if opts.Location == nil {
		opts.Location = time.Local
	}
	if log == nil {
		log = slog.Default()
	}
	if opts.TopN <= 0 {
		opts.TopN = 100
	}
	if opts.Interval <= 0 {
		opts.Interval = time.Minute
	}
	if src.SysNet == "" {
		src.SysNet = "/sys/class/net"
	}
	if src.Proc == "" {
		src.Proc = "/proc"
	}
	c := &Collector{opts: opts, src: src, store: store, log: log, now: time.Now, hours: map[int64]map[string]map[string]int64{}}
	c.status.Enabled = true
	c.status.Database = store.Path()
	c.status.Retention.Minute = humanDuration(opts.RetainMinute)
	c.status.Retention.Hour = humanDuration(opts.RetainHour)
	c.status.Retention.Day = humanDuration(opts.RetainDay)
	return c
}

// SetClock replaces the clock (tests).
func (c *Collector) SetClock(now func() time.Time) { c.now = now }

// Store returns the database.
func (c *Collector) Store() *Store { return c.store }

// Options returns the collector's options.
func (c *Collector) Options() Options { return c.opts }

// Status returns the collector's health.
func (c *Collector) Status() Status {
	c.mu.Lock()
	defer c.mu.Unlock()
	s := c.status
	s.Size = c.store.Size()
	return s
}

// Latest returns the newest gauges (nil before the first sample).
func (c *Collector) Latest() *Latest {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.latest == nil {
		return nil
	}
	l := *c.latest
	return &l
}

func (c *Collector) minuteOf(t time.Time) int64 { return t.Unix() - t.Unix()%60 }
func (c *Collector) hourOf(t time.Time) int64   { return t.Unix() - t.Unix()%3600 }

// dayOf returns the start of the local day containing t.
func (c *Collector) dayOf(t time.Time) int64 {
	t = t.In(c.opts.Location)
	return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, c.opts.Location).Unix()
}

// Run collects until ctx is done.
func (c *Collector) Run(ctx context.Context) {
	c.Start(ctx)
	for {
		t := time.NewTimer(time.Until(c.nextSample()))
		select {
		case <-ctx.Done():
			t.Stop()
			c.write(context.Background(), &Batch{})
			return
		case <-t.C:
			c.Sample(ctx)
		}
	}
}

// nextSample: two seconds into the next minute, so the minute that just
// ended is complete (or after the interval, when it isn't a minute).
func (c *Collector) nextSample() time.Time {
	if c.opts.Interval != time.Minute {
		return c.now().Add(c.opts.Interval)
	}
	return time.Unix(c.minuteOf(c.now())+62, 0)
}

// Start takes the first sample (the baseline the next ones are counted
// from), continues the current hour's top lists, and adds hours that
// ended while housecarld wasn't running to their days.
func (c *Collector) Start(ctx context.Context) {
	if v, _ := c.store.Meta(ctx, "first_data"); v != "" {
		if n, err := strconv.ParseInt(v, 10, 64); err == nil {
			t := time.Unix(n, 0)
			c.status.FirstData = &t
		}
	}
	cur := c.hourOf(c.now())
	b := &Batch{}
	merged, _ := c.store.Meta(ctx, "merged_through")
	if last, err := strconv.ParseInt(merged, 10, 64); err == nil {
		for h := max(last+3600, cur-48*3600); h < cur; h += 3600 {
			kinds, _ := c.store.KindsAt(ctx, ResHour, h)
			for _, k := range kinds {
				if es, err := c.store.TopsAt(ctx, ResHour, h, k); err == nil && len(es) > 0 {
					b.AddTops(ResDay, c.dayOf(time.Unix(h, 0)), k, es, 3*c.opts.TopN)
				}
			}
		}
	}
	b.SetMeta("merged_through", strconv.FormatInt(cur-3600, 10))
	kinds, _ := c.store.KindsAt(ctx, ResHour, cur)
	c.mu.Lock()
	for _, k := range kinds {
		if es, err := c.store.TopsAt(ctx, ResHour, cur, k); err == nil {
			m := c.topsMap(cur, k)
			for _, e := range es {
				m[e.Key] = e.Count
			}
		}
	}
	c.mu.Unlock()
	if err := c.store.Commit(ctx, b); err != nil {
		c.log.Warn("statistics: catching up failed", "err", err)
	}
	c.sample(ctx, -1)
	c.lastPrune = c.now()
	c.prune(ctx)
	c.log.Info("statistics collector started", "database", c.store.Path())
}

// Sample reads every source and writes what changed since the last
// sample.
func (c *Collector) Sample(ctx context.Context) {
	now := c.now()
	// The interval that just ended: with a minute's interval, sampling two
	// seconds into a minute, the previous minute.
	minute := c.minuteOf(now.Add(-c.opts.Interval / 2))
	vals := c.sample(ctx, minute)
	b := &Batch{}
	if len(vals) > 0 {
		t := time.Unix(minute, 0)
		b.AddCounters(ResMinute, minute, vals)
		b.AddCounters(ResHour, c.hourOf(t), vals)
		b.AddCounters(ResDay, c.dayOf(t), vals)
		c.mu.Lock()
		if c.status.FirstData == nil {
			c.status.FirstData = &t
			b.SetMeta("first_data", strconv.FormatInt(minute, 10))
		}
		c.mu.Unlock()
	}
	c.write(ctx, b)
	if now.Sub(c.lastPrune) >= time.Hour {
		c.lastPrune = now
		c.prune(ctx)
	}
}

// write adds the live hours' top lists to b, ends the hours that are
// over, and commits.
func (c *Collector) write(ctx context.Context, b *Batch) {
	c.mu.Lock()
	cur := c.hourOf(c.now())
	for h, kinds := range c.hours {
		for k, m := range kinds {
			b.SetTops(ResHour, h, k, top(m, c.opts.TopN))
		}
		if h < cur {
			// The hour is over: add it to its day.
			for k, m := range kinds {
				b.AddTops(ResDay, c.dayOf(time.Unix(h, 0)), k, top(m, c.opts.TopN), 3*c.opts.TopN)
			}
			delete(c.hours, h)
			b.SetMeta("merged_through", strconv.FormatInt(h, 10))
		}
	}
	c.mu.Unlock()
	if err := c.store.Commit(ctx, b); err != nil {
		c.log.Error("statistics: writing failed", "err", err)
	}
}

// deltas turn a source's totals into what changed since its last read.
type deltas struct {
	prev, next map[string]int64
	based      bool
}

// of returns how much key grew. The first read is the baseline (nothing
// counts); after it, a new key counts whole (a new interface, rule or
// device starts from zero), and a total that went down started again (an
// interface came back, a quota was reset, a device's entry expired).
func (d *deltas) of(key string, cur int64) int64 {
	if d.next == nil {
		d.next = map[string]int64{}
	}
	d.next[key] = cur
	if !d.based {
		return 0
	}
	p, ok := d.prev[key]
	if !ok || cur < p {
		return cur
	}
	return cur - p
}

// done keeps this read's totals for the next (a failed read skips it, so
// its keys aren't new next time).
func (d *deltas) done() { d.prev, d.next, d.based = d.next, nil, true }

// sample reads the sources; minute is the bucket the changes go to (-1:
// only the baseline). It returns the counters.
func (c *Collector) sample(ctx context.Context, minute int64) map[string]int64 {
	now := c.now()
	vals := map[string]int64{}
	add := func(k string, v int64) {
		if v != 0 {
			vals[k] += v
		}
	}
	gauge := func(k string, v int64) {
		vals["g:"+k] += v
		vals["gn:"+k]++
	}
	tops := map[string]map[string]int64{}
	addTop := func(kind, key string, v int64) {
		if v <= 0 {
			return
		}
		if tops[kind] == nil {
			tops[kind] = map[string]int64{}
		}
		tops[kind][key] += v
	}
	var errs []string

	// Interfaces, and their zones.
	var zones map[string]string
	if c.src.Zones != nil {
		zones = c.src.Zones()
	}
	for name, ic := range readInterfaces(c.src.SysNet) {
		rx, tx := c.ifs.of("rx:"+name, ic.RxBytes), c.ifs.of("tx:"+name, ic.TxBytes)
		add("rx:"+name, rx)
		add("tx:"+name, tx)
		add("rxp:"+name, c.ifs.of("rxp:"+name, ic.RxPackets))
		add("txp:"+name, c.ifs.of("txp:"+name, ic.TxPackets))
		if z := zones[name]; z != "" {
			add("zrx:"+z, rx)
			add("ztx:"+z, tx)
		}
	}
	c.ifs.done()

	// Housecarl's table: hits, drops, quotas, devices.
	if c.src.Table != nil {
		t, err := c.src.Table(ctx)
		switch {
		case err == nil:
			for mark, ctr := range t.Counters {
				p := c.table.of("hits:"+mark, int64(ctr.Packets))
				add("hits:"+mark, p)
				add("hitb:"+mark, c.table.of("hitb:"+mark, int64(ctr.Bytes)))
				if c.src.Classify != nil && p > 0 {
					if chain, drop := c.src.Classify(mark); drop {
						add("drop:"+chain, p)
						add("drops", p)
					}
				}
			}
			for name, used := range t.Quotas {
				add("quota:"+strings.TrimPrefix(name, "quota_"), c.table.of("quota:"+name, used))
			}
			for set, es := range t.QuotaSets {
				kind := strings.TrimSuffix(strings.TrimSuffix(set, "_v4"), "_v6")
				for _, e := range es {
					addTop(kind, e.Addr, c.table.of(set+"/"+e.Addr, e.Used))
				}
			}
			for set, es := range t.DeviceSets {
				kind := map[bool]string{true: KindDeviceUp, false: KindDeviceDown}[strings.HasPrefix(set, "stats_up")]
				for _, e := range es {
					key := e.Addr
					if strings.HasSuffix(set, "_v6") {
						key += "/64"
					}
					addTop(kind, key, c.table.of(set+"/"+e.Addr, int64(e.Bytes)))
				}
			}
			c.table.done()
		case errors.Is(err, nft.ErrNoTable):
		default:
			errs = append(errs, "reading the table: "+err.Error())
		}
	}

	// Gauges.
	l := &Latest{At: now}
	if n, m, ok := readConntrack(c.src.Proc); ok {
		l.Conntrack, l.ConntrackMax = n, m
		gauge(GaugeConntrack, n)
		gauge(GaugeConntrackMax, m)
	}
	if rss, ticks, ok := readSelf(c.src.Proc); ok {
		l.RSS = rss
		gauge(GaugeRSS, rss)
		if !c.cpuAt.IsZero() && ticks >= c.cpuTicks {
			secs := max(int64(now.Sub(c.cpuAt)/time.Second), 1)
			l.CPUPermille = (ticks - c.cpuTicks) * 1000 / clockTicks / secs
			gauge(GaugeCPU, l.CPUPermille)
		}
		c.cpuTicks, c.cpuAt = ticks, now
	}
	if mem, load, ok := readHost(c.src.Proc); ok {
		l.MemUsed, l.Load = mem, load
		gauge(GaugeMem, mem)
		gauge(GaugeLoad, int64(load*100))
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	c.latest = l
	c.status.LastSample = &now
	c.status.SampleError = strings.Join(errs, "; ")
	if minute < 0 {
		return nil
	}
	h := c.hourOf(time.Unix(minute, 0))
	for kind, m := range tops {
		for k, v := range m {
			c.addTop(h, kind, k, v)
		}
	}
	return vals
}

// Packet counts a packet the firewall logged: drops by source and by
// what they tried.
func (c *Collector) Packet(e api.LogEntry) {
	if e.Verdict != "drop" && e.Verdict != "reject" {
		return
	}
	what := e.Protocol
	switch {
	case e.DestPort > 0:
		what = fmt.Sprintf("%s/%d", e.Protocol, e.DestPort)
	case e.ICMPType != "":
		what = e.Protocol + "/" + e.ICMPType
	}
	t := e.Time
	if t.IsZero() {
		t = c.now()
	}
	h := c.hourOf(t)
	c.mu.Lock()
	defer c.mu.Unlock()
	if e.Source != "" {
		c.addTop(h, KindBlockedSource, e.Source, 1)
	}
	if what != "" {
		c.addTop(h, KindBlockedPort, what, 1)
	}
}

// addTop counts into an hour's top list. Called with c.mu held.
func (c *Collector) addTop(hour int64, kind, key string, n int64) {
	m := c.topsMap(hour, kind)
	if _, ok := m[key]; !ok && len(m) >= maxKeys {
		return
	}
	m[key] += n
}

func (c *Collector) topsMap(hour int64, kind string) map[string]int64 {
	h := c.hours[hour]
	if h == nil {
		h = map[string]map[string]int64{}
		c.hours[hour] = h
	}
	m := h[kind]
	if m == nil {
		m = map[string]int64{}
		h[kind] = m
	}
	return m
}

// top returns the n largest entries of m.
func top(m map[string]int64, n int) []Entry {
	out := make([]Entry, 0, len(m))
	for k, v := range m {
		out = append(out, Entry{Key: k, 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)
	})
	if len(out) > n {
		out = out[:n]
	}
	return out
}

func (c *Collector) prune(ctx context.Context) {
	now := c.now().Unix()
	for res, keep := range map[string]time.Duration{ResMinute: c.opts.RetainMinute, ResHour: c.opts.RetainHour, ResDay: c.opts.RetainDay} {
		if keep <= 0 {
			continue
		}
		if n, err := c.store.Prune(ctx, res, now-int64(keep.Seconds())); err != nil {
			c.log.Warn("statistics: pruning failed", "err", err)
		} else if n > 0 {
			c.log.Debug("statistics: pruned", "resolution", res, "rows", n)
		}
	}
}