handler.go

package handler

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/canpacis/payload"
)

type ctxKeyTyp string

const (
	responseWriterKey = ctxKeyTyp("response-writer")
	requestKey        = ctxKeyTyp("request")
	payloadKey        = ctxKeyTyp("payload")
	responseKey       = ctxKeyTyp("response")
	errorKey          = ctxKeyTyp("error")
)

// Of wraps a typed handler function into an [http.Handler]. It automatically
// decodes the incoming request into a value of type Payload using
// [payload.UnmarshalRequest], runs it through the configured validator, calls
// the provided method, and JSON-encodes the response. If the method returns a
// nil response, a 204 No Content status is written instead.
//
// Any error at each stage is forwarded to the configured [ErrorWriter]. Panics
// are caught by the configured [PanicHandler].
func Of[Payload, Response any](method func(context.Context, *Payload) (*Response, error), middlewares ...Middleware) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if p := recover(); p != nil {
				handlePanic(w, r, p)
			}
		}()

		ctx := r.Context()
		ctx = context.WithValue(ctx, responseWriterKey, w)
		ctx = context.WithValue(ctx, requestKey, r)

		var p Payload
		if err := payload.UnmarshalRequest(r, &p); err != nil {
			writeErr(w, r, wrapErr(DecodePayloadErrorKind, err))
			return
		}
		if err := validate(p); err != nil {
			writeErr(w, r, wrapErr(ValidatePayloadErrorKind, err))
			return
		}
		ctx = context.WithValue(ctx, payloadKey, p)
		var response *Response
		var err error

		next := func(ctx context.Context) context.Context {
			response, err = method(ctx, &p)
			ctx = context.WithValue(ctx, responseKey, p)
			ctx = context.WithValue(ctx, errorKey, err)
			return ctx
		}

		for _, middleware := range middlewares {
			next = middleware(next)
		}

		ctx = next(ctx)
		if err != nil {
			writeErr(w, r, wrapErr(CallMethodErrorKind, err))
			return
		}

		if w.Header().Get("Content-Type") == "" {
			w.Header().Set("Content-Type", "application/json")
		}
		if response == nil {
			w.WriteHeader(http.StatusNoContent)
			return
		}
		if err := json.NewEncoder(w).Encode(response); err != nil {
			writeErr(w, r, wrapErr(EncodeResponseErrorKind, err))
			return
		}
	})
}

// OfNoPayload is a convenience wrapper around [Of] for handler functions that
// do not require a request payload. The request body is ignored and the method
// receives only a [context.Context].
func OfNoPayload[Response any](method func(context.Context) (*Response, error)) http.Handler {
	return Of(func(ctx context.Context, p *struct{}) (*Response, error) {
		return method(ctx)
	})
}

// OfNoResponse is a convenience wrapper around [Of] for handler functions that
// do not produce a response body. A successful call results in a 204 No Content
// response.
func OfNoResponse[Payload any](method func(context.Context, *Payload) error) http.Handler {
	return Of(func(ctx context.Context, p *Payload) (*struct{}, error) {
		return nil, method(ctx, p)
	})
}

// OfAction is a convenience wrapper around [Of] for handler functions that
// require neither a request payload nor a response body — i.e. simple actions.
// A successful call results in a 204 No Content response.
func OfAction(method func(context.Context) error) http.Handler {
	return Of(func(ctx context.Context, p *struct{}) (*struct{}, error) {
		return nil, method(ctx)
	})
}

// OfStream wraps a streaming handler function into an [http.Handler]. Unlike
// [Of], the response is not JSON-encoded automatically; instead, the method
// receives the [http.ResponseWriter] as an [io.Writer] and is responsible for
// writing its own output. This is useful for server-sent events, chunked
// transfers, or any other streaming response format.
//
// Request payload decoding follows the same path as [Of]. Panics are caught by
// the configured [PanicHandler].
func OfStream[Payload any](method func(context.Context, *Payload, io.Writer) error) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if p := recover(); p != nil {
				handlePanic(w, r, p)
			}
		}()

		ctx := r.Context()
		ctx = context.WithValue(ctx, responseWriterKey, w)

		var p Payload
		if err := payload.UnmarshalRequest(r, &p); err != nil {
			writeErr(w, r, wrapErr(DecodePayloadErrorKind, err))
			return
		}
		if err := validate(p); err != nil {
			writeErr(w, r, wrapErr(ValidatePayloadErrorKind, err))
			return
		}

		if err := method(ctx, &p, w); err != nil {
			writeErr(w, r, wrapErr(CallMethodErrorKind, err))
			return
		}
	})
}

// OfActionStream is a convenience wrapper around [OfStream] for streaming
// handler functions that do not require a request payload.
func OfActionStream(method func(context.Context, io.Writer) error) http.Handler {
	return OfStream(func(ctx context.Context, p *struct{}, w io.Writer) error {
		return method(ctx, w)
	})
}

// Validator is a function type used to validate a decoded request payload.
// It should return a non-nil error if the payload is invalid.
type Validator func(any) error

// validate is the package-level validator, defaulting to a no-op.
var validate Validator = func(any) error {
	return nil
}

// SetDefaultValidator replaces the package-level [Validator] used by [Of] and
// its variants. This is typically called once during application startup to
// wire in a struct validation library such as go-playground/validator.
func SetDefaultValidator(v Validator) {
	validate = v
}

// ErrorWriter is a function type responsible for writing an error response to
// the client. Implementations may inspect the error (e.g. via [errors.As] on
// [*Error]) to choose an appropriate HTTP status code or response body format.
type ErrorWriter func(http.ResponseWriter, *http.Request, error) error

// HTTPErrorWriter is the default [ErrorWriter]. It responds with a 500 Internal
// Server Error status and writes the error message as plain text.
var HTTPErrorWriter ErrorWriter = func(w http.ResponseWriter, r *http.Request, err error) error {
	w.WriteHeader(http.StatusInternalServerError)
	_, e := w.Write([]byte(err.Error()))
	return e
}

var writeErr ErrorWriter = HTTPErrorWriter

// SetDefaultErrorWriter replaces the package-level [ErrorWriter]. Use this to
// customise error responses — for example to return JSON error payloads or map
// specific [ErrorKind] values to different HTTP status codes.
func SetDefaultErrorWriter(ew ErrorWriter) {
	writeErr = ew
}

// ErrorKind classifies the stage at which an error occurred during request
// handling, allowing [ErrorWriter] implementations to respond differently
// depending on the source of the error.
type ErrorKind int

const (
	// UnknownErrorKind indicates an error of unclassified origin.
	UnknownErrorKind = ErrorKind(iota)
	// DecodePayloadErrorKind indicates a failure to decode the request payload.
	DecodePayloadErrorKind
	// ValidatePayloadErrorKind indicates that the decoded payload failed validation.
	ValidatePayloadErrorKind
	// CallMethodErrorKind indicates that the handler method itself returned an error.
	CallMethodErrorKind
	// EncodeResponseErrorKind indicates a failure to JSON-encode the response.
	EncodeResponseErrorKind
)

var kindMap = map[ErrorKind]string{
	UnknownErrorKind:         "unknown error",
	DecodePayloadErrorKind:   "decode payload error",
	ValidatePayloadErrorKind: "validate payload error",
	CallMethodErrorKind:      "call method error",
	EncodeResponseErrorKind:  "encode response error",
}

// Error is the error type produced by this package. It wraps an underlying
// error with an [ErrorKind] that identifies which stage of request handling
// failed. It implements the standard error and unwrap interfaces.
type Error struct {
	Kind ErrorKind
	Err  error
}

// Error returns a human-readable description of the error, prefixed with the
// name of the error kind.
func (e *Error) Error() string {
	msg, ok := kindMap[e.Kind]
	if !ok {
		msg = kindMap[UnknownErrorKind]
	}
	return fmt.Sprintf("%s: %s", msg, e.Err.Error())
}

// Unwrap returns the underlying error, enabling use with [errors.Is] and
// [errors.As].
func (e *Error) Unwrap() error {
	return e.Err
}

func wrapErr(kind ErrorKind, err error) *Error {
	return &Error{Kind: kind, Err: err}
}

// PanicHandler is a function type called when a panic is recovered inside a
// handler. It receives the [http.ResponseWriter], the original [*http.Request],
// and the value passed to panic.
type PanicHandler func(http.ResponseWriter, *http.Request, any)

// handlePanic is the package-level panic handler. By default it re-panics,
// preserving standard Go panic behaviour. Replace it with [SetDefaultPanicHandler]
// to add logging, metrics, or a structured error response.
var handlePanic PanicHandler = func(w http.ResponseWriter, r *http.Request, p any) {
	panic(p)
}

// SetDefaultPanicHandler replaces the package-level [PanicHandler]. This is
// useful for recovering from panics gracefully — for example, logging the stack
// trace and returning a 500 response rather than crashing the server.
func SetDefaultPanicHandler(ph PanicHandler) {
	handlePanic = ph
}

// ResponseWriter retrieves the [http.ResponseWriter] stored in the context by
// the handler wrappers. This allows methods that only receive a
// [context.Context] to access the underlying response writer when they need to
// set headers or perform other low-level HTTP operations outside of the normal
// response encoding path. Returns nil if no writer is found in the context.
func ResponseWriter(ctx context.Context) http.ResponseWriter {
	w, ok := ctx.Value(responseWriterKey).(http.ResponseWriter)
	if !ok {
		return nil
	}
	return w
}

func Request(ctx context.Context) *http.Request {
	r, ok := ctx.Value(requestKey).(*http.Request)
	if !ok {
		return nil
	}
	return r
}

func Payload(ctx context.Context) any {
	return ctx.Value(payloadKey)
}

func Response(ctx context.Context) any {
	return ctx.Value(responseKey)
}

func GetError(ctx context.Context) error {
	err, ok := ctx.Value(errorKey).(error)
	if !ok {
		return nil
	}
	return err
}

type Handler = func(context.Context) context.Context
type Middleware = func(Handler) Handler