package middleware

import (
	"context"
	"time"

	"git.echol.cn/loser/ai_proxy/server/model/common/response"
	"github.com/gin-gonic/gin"
)

// Timeout 超时中间件
func Timeout(timeout time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
		defer cancel()

		c.Request = c.Request.WithContext(ctx)

		finished := make(chan struct{})
		go func() {
			c.Next()
			finished <- struct{}{}
		}()

		select {
		case <-ctx.Done():
			response.FailWithMessage("请求超时", c)
			c.Abort()
		case <-finished:
		}
	}
}