package httpx

import (
	"errors"
	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"git.yuntuo.hk/k2safe/go-omni/response"
	"net/http"
)

// LimitBody is useful for handwritten routes. Built-in App.HandleWithBodyLimit
// additionally installs the reader cap before global middleware can read it.
func LimitBody(maxBytes int64) Middleware {
	if maxBytes <= 0 || maxBytes > 1<<40 {
		panic("httpx: body limit must be 1..1 TiB")
	}
	return func(next Handler) Handler {
		return func(c *Context) error {
			r := c.Request()
			if r.ContentLength > maxBytes {
				return bodyTooLarge()
			}
			if r.Body != nil {
				r.Body = http.MaxBytesReader(c.Raw(), r.Body, maxBytes)
			}
			defer CleanupMultipart(r)
			err := next(c)
			var limitErr *http.MaxBytesError
			if errors.As(err, &limitErr) {
				return bodyTooLarge()
			}
			return err
		}
	}
}
func bodyTooLarge() error { return apierrors.New(413, "BODY_TOO_LARGE", "request body exceeds limit") }

// CleanupMultipart removes disk spill files created by ParseMultipartForm.
// Callers still close individual multipart files. Streaming MultipartReader
// handlers own any files they create themselves.
func CleanupMultipart(r *http.Request) {
	if r != nil && r.MultipartForm != nil {
		_ = r.MultipartForm.RemoveAll()
	}
}

type BodyLimitRegistrar interface {
	HandleWithBodyLimit(string, string, int64, response.Encoder, Handler, ...Middleware)
}

func HandleWithBodyLimit(router interface {
	Handle(string, string, Handler, ...Middleware)
}, method, path string, maxBytes int64, encoder response.Encoder, handler Handler, middleware ...Middleware) {
	if r, ok := router.(BodyLimitRegistrar); ok {
		r.HandleWithBodyLimit(method, path, maxBytes, encoder, handler, middleware...)
		return
	}
	// Custom registrars without the early hook cap at route middleware entry.
	all := append([]Middleware{LimitBody(maxBytes)}, middleware...)
	HandleWithResponse(router, method, path, encoder, handler, all...)
}