package main

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"sync"
	"time"
)

func proxy(target *url.URL, timeout time.Duration) *httputil.ReverseProxy {
	tr := &http.Transport{DialContext: (&net.Dialer{Timeout: time.Second}).DialContext, ResponseHeaderTimeout: timeout, MaxIdleConnsPerHost: 8, IdleConnTimeout: 30 * time.Second}
	return &httputil.ReverseProxy{
		Rewrite:   func(r *httputil.ProxyRequest) { r.SetURL(target); r.SetXForwarded() },
		Transport: tr,
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			http.Error(w, "bad gateway", http.StatusBadGateway)
		},
	}
}

type bucket struct {
	mu                     sync.Mutex
	tokens, capacity, rate float64
	last                   time.Time
}

func (b *bucket) allow(now time.Time) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if now.After(b.last) {
		b.tokens = min(b.capacity, b.tokens+now.Sub(b.last).Seconds()*b.rate)
		b.last = now
	}
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}
func main() {
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "xff=%s proto=%s forwarded=%q\n", r.Header.Get("X-Forwarded-For"), r.Header.Get("X-Forwarded-Proto"), r.Header.Get("Forwarded"))
	}))
	defer backend.Close()
	target, err := url.Parse(backend.URL)
	if err != nil {
		panic(err)
	}
	p := proxy(target, time.Second)
	defer p.Transport.(*http.Transport).CloseIdleConnections()
	req := httptest.NewRequest(http.MethodGet, "http://api.example.test/", nil)
	req.RemoteAddr = "192.0.2.1:1234"
	req.Header.Set("X-Forwarded-For", "198.51.100.99")
	req.Header.Set("Forwarded", "for=198.51.100.99;proto=https")
	w := httptest.NewRecorder()
	p.ServeHTTP(w, req)
	body, err := io.ReadAll(w.Result().Body)
	if err != nil {
		panic(err)
	}
	fmt.Print(string(body))
	start := time.Unix(0, 0)
	b := bucket{tokens: 2, capacity: 2, rate: 1, last: start}
	fmt.Println("bucket:", b.allow(start), b.allow(start), b.allow(start), b.allow(start.Add(time.Second)))
}