package blacktip

import (
	"net"
	"net/netip"
	"sync"

	"git.veillax.com/Veillax/reef/model"
)

// maxTrackedSources bounds the per-port table of client addresses between two reports, so a flood
// from spoofed or scanning addresses can't grow it without limit. Addresses past it still count
// in the port's totals; they just aren't named.
const maxTrackedSources = 5000

// sourceTable counts a port's traffic per client address until the next report takes it.
type sourceTable struct {
	mu sync.Mutex
	by map[string]*model.SourceStat
}

func (t *sourceTable) add(ip string, conns, in, out int64) {
	if ip == "" {
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	s := t.by[ip]
	if s == nil {
		if len(t.by) >= maxTrackedSources {
			return
		}
		if t.by == nil {
			t.by = map[string]*model.SourceStat{}
		}
		s = &model.SourceStat{IP: ip}
		t.by[ip] = s
	}
	s.Conns += conns
	s.BytesIn += in
	s.BytesOut += out
}

func (t *sourceTable) take() map[string]*model.SourceStat {
	t.mu.Lock()
	defer t.mu.Unlock()
	by := t.by
	t.by = nil
	return by
}

// clientIP reduces a "host:port" remote address to the client's IP (IPv4-mapped IPv6 as plain IPv4).
// Anything that isn't an address gives "".
func clientIP(remote string) string {
	host, _, err := net.SplitHostPort(remote)
	if err != nil {
		host = remote
	}
	a, err := netip.ParseAddr(host)
	if err != nil {
		return ""
	}
	return a.WithZone("").Unmap().String()
}

// record counts traffic on a port, and names the client it came from when that is known.
func (ms *mapState) record(remote string, conns, in, out int64) {
	ms.conns.Add(conns)
	ms.bytesIn.Add(in)
	ms.bytesOut.Add(out)
	ms.sources.add(clientIP(remote), conns, in, out)
}

// takeStats hands over (and resets) the traffic counted since the last call.
func (ms *mapState) takeStats() (model.Stat, bool) {
	st := model.Stat{Conns: ms.conns.Swap(0), BytesIn: ms.bytesIn.Swap(0), BytesOut: ms.bytesOut.Swap(0)}
	st.Sources = model.TopSources(ms.sources.take(), model.MaxSources)
	return st, st.Conns|st.BytesIn|st.BytesOut != 0
}

// restoreStats puts taken traffic back after a report that didn't arrive, so none of it is lost.
func (ms *mapState) restoreStats(st model.Stat) {
	ms.conns.Add(st.Conns)
	ms.bytesIn.Add(st.BytesIn)
	ms.bytesOut.Add(st.BytesOut)
	for _, s := range st.Sources {
		ms.sources.add(s.IP, s.Conns, s.BytesIn, s.BytesOut)
	}
}