package ratelimit

import (
	"container/list"
	"errors"
	"net/netip"
	"sync"
	"time"
)

type Rate struct {
	Burst int
	Every time.Duration
}

type entry[K comparable] struct {
	key    K
	tokens float64
	at     time.Time
}

var errNoEntries = errors.New("ratelimit: a table needs at least one entry")

type Table[K comparable] struct {
	mu      sync.Mutex
	now     func() time.Time
	max     int
	entries map[K]*list.Element
	order   *list.List
	evicted uint64
}

func New[K comparable](max int, now func() time.Time) *Table[K] {
	if max < 1 {
		panic(errNoEntries)
	}
	return &Table[K]{now: now, max: max, entries: map[K]*list.Element{}, order: list.New()}
}

func (t *Table[K]) bucket(k K, rate Rate, now time.Time) *entry[K] {
	el, ok := t.entries[k]
	if !ok {
		el = t.order.PushFront(&entry[K]{key: k, tokens: float64(rate.Burst), at: now})
		t.entries[k] = el
		for t.order.Len() > t.max {
			oldest := t.order.Back()
			t.order.Remove(oldest)
			delete(t.entries, oldest.Value.(*entry[K]).key)
			t.evicted++
		}
	}
	t.order.MoveToFront(el)
	e := el.Value.(*entry[K])
	if elapsed := now.Sub(e.at); elapsed > 0 {
		e.tokens = min(float64(rate.Burst), e.tokens+float64(elapsed)/float64(rate.Every))
	}
	e.at = now
	return e
}

func wait(tokens float64, rate Rate) time.Duration {
	return time.Duration((1 - tokens) * float64(rate.Every))
}

func (t *Table[K]) Take(k K, rate Rate) (time.Duration, bool) {
	if rate.Burst <= 0 || rate.Every <= 0 {
		return 0, true
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	e := t.bucket(k, rate, t.now())
	if e.tokens >= 1 {
		e.tokens--
		return 0, true
	}
	return wait(e.tokens, rate), false
}

func (t *Table[K]) Allowed(k K, rate Rate) (time.Duration, bool) {
	if rate.Burst <= 0 || rate.Every <= 0 {
		return 0, true
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if _, ok := t.entries[k]; !ok {
		return 0, true
	}
	e := t.bucket(k, rate, t.now())
	if e.tokens >= 1 {
		return 0, true
	}
	return wait(e.tokens, rate), false
}

func (t *Table[K]) Len() int {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.order.Len()
}

func Prefix(a netip.Addr) netip.Prefix {
	if !a.IsValid() {
		return netip.Prefix{}
	}
	a = a.WithZone("").Unmap()
	bits := 32
	if a.Is6() {
		bits = 64
	}
	p, _ := a.Prefix(bits)
	return p
}