package middleware

import (
	"context"
	"time"

	"github.com/go-kratos/kratos/v2/errors"
	"github.com/go-kratos/kratos/v2/middleware"
)

// Timeout 超时中间件
func Timeout(timeout time.Duration) middleware.Middleware {
	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, req interface{}) (interface{}, error) {
			ctx, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()

			// 使用 buffered channel 避免 goroutine 泄漏
			type result struct {
				reply interface{}
				err   error
			}
			done := make(chan result, 1)
			panicChan := make(chan interface{}, 1)

			go func() {
				defer func() {
					if p := recover(); p != nil {
						select {
						case panicChan <- p:
						default:
						}
					}
				}()

				reply, err := handler(ctx, req)
				select {
				case done <- result{reply: reply, err: err}:
				default:
				}
			}()

			select {
			case p := <-panicChan:
				panic(p)
			case r := <-done:
				return r.reply, r.err
			case <-ctx.Done():
				return nil, errors.GatewayTimeout("TIMEOUT", "请求超时")
			}
		}
	}
}