package main

import (
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

type limiter struct {
	mu     sync.Mutex
	hits   map[string][]time.Time
	max    int
	window time.Duration
}

func newLimiter(max int, window time.Duration) *limiter {
	return &limiter{hits: map[string][]time.Time{}, max: max, window: window}
}

func (l *limiter) prune(key string, now time.Time) []time.Time {
	cutoff := now.Add(-l.window)
	kept := l.hits[key][:0]
	for _, t := range l.hits[key] {
		if t.After(cutoff) {
			kept = append(kept, t)
		}
	}
	if len(kept) == 0 {
		delete(l.hits, key)
	} else {
		l.hits[key] = kept
	}
	if len(l.hits) > 10000 {
		for k, v := range l.hits {
			if !v[len(v)-1].After(cutoff) {
				delete(l.hits, k)
			}
		}
	}
	return kept
}

func (l *limiter) blocked(key string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	return len(l.prune(key, time.Now())) >= l.max
}

func (l *limiter) fail(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := time.Now()
	l.hits[key] = append(l.prune(key, now), now)
}

func clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	if ip := net.ParseIP(host); ip != nil && ip.IsLoopback() {
		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
			hops := strings.Split(xff, ",")
			return strings.TrimSpace(hops[len(hops)-1])
		}
	}
	return host
}

type bucket struct {
	tokens float64
	last   time.Time
}

type buckets struct {
	mu    sync.Mutex
	m     map[string]*bucket
	rate  float64
	burst float64
}

func newBuckets(rate, burst float64) *buckets {
	return &buckets{m: map[string]*bucket{}, rate: rate, burst: burst}
}

func (b *buckets) allow(key string) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	now := time.Now()
	k, ok := b.m[key]
	if !ok {
		if len(b.m) >= 20000 {
			for name, v := range b.m {
				if v.tokens+now.Sub(v.last).Seconds()*b.rate >= b.burst {
					delete(b.m, name)
				}
			}
			if len(b.m) >= 20000 {
				return false
			}
		}
		k = &bucket{tokens: b.burst, last: now}
		b.m[key] = k
	}
	k.tokens = min(b.burst, k.tokens+now.Sub(k.last).Seconds()*b.rate)
	k.last = now
	if k.tokens < 1 {
		return false
	}
	k.tokens--
	return true
}

func heavyRequest(r *http.Request) bool {
	p := r.URL.Path
	q := r.URL.Query()
	return p == "/search" || strings.Contains(p, "/archive/") || strings.Contains(p, "/blame/") ||
		strings.Contains(p, "/interdiff") || (strings.HasSuffix(p, "/search") && q.Get("q") != "") ||
		(strings.Contains(p, "/log") && q.Get("q") != "") || p == "/login" || p == "/register"
}

func (s *Server) throttle(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/static/") {
			next.ServeHTTP(w, r)
			return
		}
		ip := clientIP(r)
		if !s.requests.allow(ip) || (heavyRequest(r) && !s.heavy.allow(ip)) {
			w.Header().Set("Retry-After", "10")
			http.Error(w, "too many requests, slow down", http.StatusTooManyRequests)
			return
		}
		next.ServeHTTP(w, r)
	})
}