package relay

import (
	"strconv"
	"sync"
	"time"
)

type Limits struct {
	AppLive, AppPerMinute          int
	AppBytesPerHour                int64
	AccountLive, AccountPerMinute  int
	AccountBytesPerHour            int64
	BoxLive, BoxPerMinute          int
	AllocBitsPerSec, BoxBitsPerSec int64
}

func DefaultLimits() Limits {
	return Limits{
		AppLive: 64, AppPerMinute: 30, AppBytesPerHour: 20 << 30,
		AccountLive: 256, AccountPerMinute: 60, AccountBytesPerHour: 60 << 30,
		BoxLive: 2048, BoxPerMinute: 500,
		AllocBitsPerSec: 8_000_000, BoxBitsPerSec: 200_000_000,
	}
}

func LimitsFromEnv(getenv func(string) string) Limits {
	l := DefaultLimits()
	i := func(name string, dst *int) {
		if v, err := strconv.Atoi(getenv("CARLOS_RELAY_LIMIT_" + name)); err == nil && v > 0 {
			*dst = v
		}
	}
	i64 := func(name string, dst *int64) {
		if v, err := strconv.ParseInt(getenv("CARLOS_RELAY_LIMIT_"+name), 10, 64); err == nil && v > 0 {
			*dst = v
		}
	}
	i("APP_LIVE", &l.AppLive)
	i("APP_PER_MINUTE", &l.AppPerMinute)
	i64("APP_BYTES_PER_HOUR", &l.AppBytesPerHour)
	i("ACCOUNT_LIVE", &l.AccountLive)
	i("ACCOUNT_PER_MINUTE", &l.AccountPerMinute)
	i64("ACCOUNT_BYTES_PER_HOUR", &l.AccountBytesPerHour)
	i("BOX_LIVE", &l.BoxLive)
	i("BOX_PER_MINUTE", &l.BoxPerMinute)
	i64("ALLOC_BPS", &l.AllocBitsPerSec)
	i64("BOX_BPS", &l.BoxBitsPerSec)
	return l
}

type Identity struct{ Account, App string }

type window struct { // sliding counters bucketed per minute, 60 buckets for an hour
	at [60]time.Time
	n  [60]int64
}

func (w *window) add(now time.Time, n int64) {
	m := now.Truncate(time.Minute)
	i := int(m.Unix()/60) % 60
	if !w.at[i].Equal(m) {
		w.at[i], w.n[i] = m, 0
	}
	w.n[i] += n
}

func (w *window) sum(now time.Time, span time.Duration) int64 {
	var s int64
	cut := now.Add(-span)
	for i := range w.at {
		if !w.at[i].IsZero() && w.at[i].After(cut) {
			s += w.n[i]
		}
	}
	return s
}

type counters struct {
	live   int
	starts window // attempts per minute
	bytes  window // bytes per rolling hour
}

type Admission struct {
	mu   sync.Mutex
	l    Limits
	now  func() time.Time
	app  map[Identity]*counters
	acct map[string]*counters
	box  counters
}

func NewAdmission(l Limits, now func() time.Time) *Admission {
	return &Admission{l: l, now: now, app: map[Identity]*counters{}, acct: map[string]*counters{}}
}

func (a *Admission) get(id Identity) (*counters, *counters) {
	c, ok := a.app[id]
	if !ok {
		c = &counters{}
		a.app[id] = c
	}
	ac, ok := a.acct[id.Account]
	if !ok {
		ac = &counters{}
		a.acct[id.Account] = ac
	}
	return c, ac
}

// Admit counts an attempt for the app, the account and the box whether or
// not it is admitted, then checks live caps, rates and budgets in that order.
func (a *Admission) Admit(id Identity) (bool, string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	now := a.now()
	c, ac := a.get(id)
	c.starts.add(now, 1)
	ac.starts.add(now, 1)
	a.box.starts.add(now, 1)
	switch {
	case c.live >= a.l.AppLive:
		return false, "app live"
	case ac.live >= a.l.AccountLive:
		return false, "account live"
	case a.box.live >= a.l.BoxLive:
		return false, "box live"
	case c.starts.sum(now, time.Minute) > int64(a.l.AppPerMinute):
		return false, "app rate"
	case ac.starts.sum(now, time.Minute) > int64(a.l.AccountPerMinute):
		return false, "account rate"
	case a.box.starts.sum(now, time.Minute) > int64(a.l.BoxPerMinute):
		return false, "box rate"
	case c.bytes.sum(now, time.Hour) > a.l.AppBytesPerHour:
		return false, "app budget"
	case ac.bytes.sum(now, time.Hour) > a.l.AccountBytesPerHour:
		return false, "account budget"
	}
	c.live++
	ac.live++
	a.box.live++
	return true, ""
}

func (a *Admission) Release(id Identity) {
	a.mu.Lock()
	defer a.mu.Unlock()
	c, ac := a.get(id)
	if c.live > 0 {
		c.live--
	}
	if ac.live > 0 {
		ac.live--
	}
	if a.box.live > 0 {
		a.box.live--
	}
}

// AddBytes records forwarded bytes and reports whether the app or account
// rolling-hour budget is now exceeded.
func (a *Admission) AddBytes(id Identity, n int64) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	now := a.now()
	c, ac := a.get(id)
	c.bytes.add(now, n)
	ac.bytes.add(now, n)
	return c.bytes.sum(now, time.Hour) > a.l.AppBytesPerHour || ac.bytes.sum(now, time.Hour) > a.l.AccountBytesPerHour
}

func (a *Admission) Seed(id Identity, at time.Time, n int64) {
	a.mu.Lock()
	defer a.mu.Unlock()
	c, ac := a.get(id)
	c.bytes.add(at, n)
	ac.bytes.add(at, n)
}

func (a *Admission) OverBudget(id Identity, now time.Time) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	c, ac := a.get(id)
	return c.bytes.sum(now, time.Hour) > a.l.AppBytesPerHour || ac.bytes.sum(now, time.Hour) > a.l.AccountBytesPerHour
}

// Bucket is a token bucket in bytes with a one-second burst.
type Bucket struct {
	mu     sync.Mutex
	rate   float64 // bytes per second
	tokens float64
	last   time.Time
	now    func() time.Time
}

func NewBucket(bitsPerSec int64, now func() time.Time) *Bucket {
	r := float64(bitsPerSec) / 8
	return &Bucket{rate: r, tokens: r, last: now(), now: now}
}

func (b *Bucket) Allow(nbytes int) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	t := b.now()
	b.tokens += t.Sub(b.last).Seconds() * b.rate
	if b.tokens > b.rate {
		b.tokens = b.rate
	}
	b.last = t
	if float64(nbytes) > b.tokens {
		return false
	}
	b.tokens -= float64(nbytes)
	return true
}

// Peek refills the bucket and reports whether nbytes would fit, without
// spending anything. Pair it with Spend so multiple buckets can be checked
// before any of them is debited.
func (b *Bucket) Peek(nbytes int) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	t := b.now()
	b.tokens += t.Sub(b.last).Seconds() * b.rate
	if b.tokens > b.rate {
		b.tokens = b.rate
	}
	b.last = t
	return float64(nbytes) <= b.tokens
}

// Spend debits nbytes unconditionally, clamping at zero.
func (b *Bucket) Spend(nbytes int) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.tokens -= float64(nbytes)
	if b.tokens < 0 {
		b.tokens = 0
	}
}