package middleware

// v3.2：上传/下载带宽限速（字节/秒，0=不限速）。
//
// 设计要点：
//   - 令牌桶（token bucket）：每 Read 计算自上次起累计可消费字节，
//     未达上限则阻塞至令牌补足；保证长期平均速率，瞬时按桶容量突发。
//   - 桶容量 = rate（1 秒突发量）：简单且足够友好；后续若要更细粒度可参数化。
//   - 0 速率直接透传，不引入任何开销（性能基线保持）。
//   - 对 gin 中间件"上传限速"：包裹 c.Request.Body，使 multipart 解析/单分片读取
//     都被天然节流；下载侧由 api 层在 storage.Open 之后包裹 ReadCloser。

import (
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"

	"fileshare/internal/response"
)

// rateLimitedReader 按 bytes/sec 限速的 io.Reader（基于时间窗的精确调度）。
// 实现：累计已读字节，按 rate 推算"应到达时间"；超前则 sleep 补齐。
//   - rate <= 0 时退化为直通；
//   - 长期速率严格 ≤ rate，瞬时按调用方 Read 块大小自然突发；
//   - 并发安全（每次请求持独立实例即可，但保留锁便于复用）。
type rateLimitedReader struct {
	r     io.Reader
	rate  int       // 字节/秒
	start time.Time // 起始时间（首 Read 时记）
	bytes int64     // 累计已读字节
	mu    sync.Mutex
}

// NewRateLimitedReader 构造限速 reader；rate<=0 返回原始 reader。
func NewRateLimitedReader(r io.Reader, rate int) io.Reader {
	if r == nil || rate <= 0 {
		return r
	}
	return &rateLimitedReader{r: r, rate: rate}
}

// Read 限速：底层 Read 后若已超时应到达时间，sleep 补齐；错误/EOF 透传。
func (rl *rateLimitedReader) Read(p []byte) (int, error) {
	if rl.rate <= 0 {
		return rl.r.Read(p)
	}
	n, err := rl.r.Read(p)
	if n <= 0 {
		return n, err
	}
	rl.mu.Lock()
	if rl.start.IsZero() {
		rl.start = time.Now()
	}
	rl.bytes += int64(n)
	// 已读 n 字节的"应到达时间" = start + bytes/rate
	expected := rl.start.Add(time.Duration(float64(time.Second) * float64(rl.bytes) / float64(rl.rate)))
	now := time.Now()
	rl.mu.Unlock()
	if now.Before(expected) {
		time.Sleep(expected.Sub(now))
	}
	return n, err
}

// UploadBandwidthMiddleware 限速上传请求体（包裹 c.Request.Body）。
// rateProvider 每次请求动态读取最新配置（管理端改 KV 后立即生效）。
func UploadBandwidthMiddleware(rateProvider func() int) gin.HandlerFunc {
	return func(c *gin.Context) {
		rate := 0
		if rateProvider != nil {
			rate = rateProvider()
		}
		if rate > 0 && c.Request != nil && c.Request.Body != nil {
			c.Request.Body = WrapReadCloser(c.Request.Body, rate)
		}
		c.Next()
	}
}

// WrapReadCloser 包裹任意 io.ReadCloser 为限速版本（io.Copy 等下游使用）。
func WrapReadCloser(rc io.ReadCloser, rate int) io.ReadCloser {
	if rc == nil || rate <= 0 {
		return rc
	}
	return &rateLimitedReadCloser{rc: rc, rl: &rateLimitedReader{r: rc, rate: rate}}
}

type rateLimitedReadCloser struct {
	rc io.ReadCloser
	rl *rateLimitedReader
}

func (r *rateLimitedReadCloser) Read(p []byte) (int, error) { return r.rl.Read(p) }
func (r *rateLimitedReadCloser) Close() error               { return r.rc.Close() }

// DownloadRateReporter 限速命中提示：超长下载时可选上报。
// 当前仅占位——真正接入可在 audit 中按"下载耗时/字节比"标记慢速，但用户暂不需要。
var _ = response.Fail // 保留 response 包以备后续扩展
var _ = http.StatusOK