// Package bans is the ban ledger: the bans smallwebwaf makes on the
// netblocks of clients that break a rate limit, with their notes, as the
// "Bans" section of SPEC.md describes. The bans are kept in memory only.
package bans

import (
	"net/netip"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/golang-lru/v2/simplelru"
)

// repeatFactor is how many times as long as the netblock's last ban a ban
// for a limit broken again within the repeat window lasts.
const repeatFactor = 3

// maxTextBytes is how much of each text in a ban's notes is kept.
const maxTextBytes = 256

// Rules are how long a ban for a broken limit lasts, and how many bans
// are held.
type Rules struct {
	// LimitBanDuration is how long a first ban lasts.
	LimitBanDuration time.Duration
	// LimitBanRepeatWindow is how soon after the netblock's last ban
	// ended a broken limit counts as a repeat, which bans for
	// repeatFactor times as long as that ban.
	LimitBanRepeatWindow time.Duration
	// MaxBanDuration is the longest ban; a ban that would be longer is
	// permanent instead.
	MaxBanDuration time.Duration
	// MaxBans is the most bans held, at least one. Past it, the earliest
	// ban of the netblock that has gone longest without a request is
	// dropped.
	MaxBans int
}

// Ban is a ban on a netblock for a broken limit, the only kind of ban
// smallwebwaf makes so far.
type Ban struct {
	Netblock netip.Prefix
	Start    time.Time
	// Expires is when the ban ends, zero for a permanent ban.
	Expires time.Time
	Notes   Notes
}

// Permanent reports whether the ban never runs out.
func (b Ban) Permanent() bool {
	return b.Expires.IsZero()
}

// ActiveAt reports whether the ban refuses requests at now.
func (b Ban) ActiveAt(now time.Time) bool {
	return b.Permanent() || now.Before(b.Expires)
}

// Notes are what an admin needs to decide whether to lift a ban.
type Notes struct {
	// Country is the client's country, when it was looked up.
	Country string
	// Limit, Window and Count are the limit that was broken, its window,
	// "minute", "hour" or "day", and the count reached: the client's
	// requests in the window, the one that broke the limit included.
	// These are the requests that counted toward the ban, and the window
	// is the time over which they came.
	Limit  int64
	Window string
	Count  float64
	// Request is the request that broke the limit.
	Request Request
	// Refused is how many requests the ban has refused so far.
	Refused int64
	// EarlierBans is how many bans the netblock had before this one.
	EarlierBans int
}

// Request is a request in a ban's notes. Each text is cut to 256 bytes.
type Request struct {
	Time   time.Time
	Method string
	Host   string
	// Path is the path with its query string.
	Path string
	// Status is what the client was sent, 0 if nothing was.
	Status    int
	UserAgent string
}

// Ledger holds the bans. It is safe for concurrent use.
type Ledger struct {
	rules Rules

	mu sync.Mutex
	// netblocks holds each banned netblock's bans, oldest first. Each
	// request from a netblock makes it the most recently seen.
	netblocks *simplelru.LRU[netip.Prefix, *[]Ban]
	// held is how many bans netblocks holds, at most rules.MaxBans.
	held int
}

// New returns a Ledger with no ban yet.
func New(rules Rules) *Ledger {
	// Every netblock held has a ban, so there are never more netblocks
	// than rules.MaxBans, and the LRU never drops one itself.
	netblocks, err := simplelru.NewLRU[netip.Prefix, *[]Ban](rules.MaxBans, nil)
	if err != nil {
		panic(err) // NewLRU fails only for a size below one
	}

	return &Ledger{rules: rules, netblocks: netblocks}
}

// Check is called for each request from netblock, at now. It reports
// whether a ban on netblock is active, and returns that ban, with the
// request counted among those it refused.
func (l *Ledger) Check(netblock netip.Prefix, now time.Time) (Ban, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	bans, found := l.netblocks.Get(netblock)
	if !found {
		return Ban{}, false
	}

	// A ban is made only once the one before has ended, so only the last
	// can be active.
	last := &(*bans)[len(*bans)-1]
	if !last.ActiveAt(now) {
		return Ban{}, false
	}

	last.Notes.Refused++

	return *last, true
}

// BanForLimit bans netblock at now for a broken limit, with notes, and
// returns the ban. A first ban lasts LimitBanDuration. A ban made within
// LimitBanRepeatWindow after the netblock's last ban ended lasts
// repeatFactor times as long as that one. A ban that would be longer
// than MaxBanDuration is permanent instead. If a ban on netblock is still
// active, as when two of its requests break a limit at once, that ban is
// returned and no other is made. The ledger fills in the notes' Refused
// and EarlierBans itself.
func (l *Ledger) BanForLimit(netblock netip.Prefix, now time.Time, notes Notes) Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	var last *Ban

	bans, found := l.netblocks.Get(netblock)
	if found {
		last = &(*bans)[len(*bans)-1]
		if last.ActiveAt(now) {
			return *last
		}

		notes.EarlierBans = last.Notes.EarlierBans + 1
	}

	notes.Request = notes.Request.cut()
	ban := Ban{
		Netblock: netblock,
		Start:    now,
		Expires:  l.expiry(last, now),
		Notes:    notes,
	}

	if l.held == l.rules.MaxBans {
		l.dropOne()
	}

	// dropOne can have dropped netblock's last ban, and netblock with it.
	bans, found = l.netblocks.Peek(netblock)
	if !found {
		bans = &[]Ban{}
		l.netblocks.Add(netblock, bans)
	}

	*bans = append(*bans, ban)
	l.held++

	return ban
}

// Bans returns the bans held on netblock, oldest first. It is not a
// request from netblock, and leaves when it was last seen unchanged.
func (l *Ledger) Bans(netblock netip.Prefix) []Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	bans, found := l.netblocks.Peek(netblock)
	if !found {
		return nil
	}

	return slices.Clone(*bans)
}

// expiry returns when a ban for a broken limit made at now ends, or zero
// when it is permanent. last is the netblock's last ban, which has ended,
// or nil when it has none.
func (l *Ledger) expiry(last *Ban, now time.Time) time.Time {
	length := l.rules.LimitBanDuration

	if last != nil && now.Sub(last.Expires) <= l.rules.LimitBanRepeatWindow {
		lastLength := last.Expires.Sub(last.Start)
		// This is repeatFactor * lastLength > MaxBanDuration, written so
		// that it cannot overflow.
		if lastLength > l.rules.MaxBanDuration/repeatFactor {
			return time.Time{}
		}

		length = repeatFactor * lastLength
	}

	if length > l.rules.MaxBanDuration {
		return time.Time{}
	}

	return now.Add(length)
}

// dropOne drops the earliest ban of the netblock that has gone longest
// without a request, and the netblock with it if that was its only ban.
func (l *Ledger) dropOne() {
	netblock, bans, _ := l.netblocks.GetOldest()
	if len(*bans) == 1 {
		l.netblocks.Remove(netblock)
	} else {
		*bans = slices.Delete(*bans, 0, 1)
	}

	l.held--
}

// cut returns r with each text cut to maxTextBytes and copied, so that
// the notes do not keep the rest of the request in memory.
func (r Request) cut() Request {
	r.Method = cutText(r.Method)
	r.Host = cutText(r.Host)
	r.Path = cutText(r.Path)
	r.UserAgent = cutText(r.UserAgent)

	return r
}

// cutText returns a copy of the first maxTextBytes of text.
func cutText(text string) string {
	return strings.Clone(text[:min(len(text), maxTextBytes)])
}