package bridge

import (
	"net"
	"sync"
	"time"
)

// rateLimitedConn limits read/write throughput to bytesPerSec (simple token bucket).
type rateLimitedConn struct {
	net.Conn
	bytesPerSec int64
	mu          sync.Mutex
	allowance   float64
	lastCheck   time.Time
}

func newRateLimitedConn(conn net.Conn, bytesPerSec int64) net.Conn {
	if bytesPerSec <= 0 {
		return conn
	}
	return &rateLimitedConn{
		Conn:        conn,
		bytesPerSec: bytesPerSec,
		lastCheck:   time.Now(),
		allowance:   float64(bytesPerSec),
	}
}

func (c *rateLimitedConn) wait(n int) {
	if c.bytesPerSec <= 0 || n <= 0 {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	now := time.Now()
	elapsed := now.Sub(c.lastCheck).Seconds()
	c.lastCheck = now
	c.allowance += elapsed * float64(c.bytesPerSec)
	if c.allowance > float64(c.bytesPerSec)*2 {
		c.allowance = float64(c.bytesPerSec) * 2
	}
	for float64(n) > c.allowance {
		deficit := float64(n) - c.allowance
		sleep := time.Duration(deficit/float64(c.bytesPerSec)*1e9) * time.Nanosecond
		if sleep < time.Millisecond {
			sleep = time.Millisecond
		}
		c.mu.Unlock()
		time.Sleep(sleep)
		c.mu.Lock()
		now = time.Now()
		elapsed = now.Sub(c.lastCheck).Seconds()
		c.lastCheck = now
		c.allowance += elapsed * float64(c.bytesPerSec)
		if c.allowance > float64(c.bytesPerSec)*2 {
			c.allowance = float64(c.bytesPerSec) * 2
		}
	}
	c.allowance -= float64(n)
}

func (c *rateLimitedConn) Read(b []byte) (int, error) {
	n, err := c.Conn.Read(b)
	if n > 0 {
		c.wait(n)
	}
	return n, err
}

func (c *rateLimitedConn) Write(b []byte) (int, error) {
	c.wait(len(b))
	return c.Conn.Write(b)
}