router.go

// ABOUTME: WebSocket method router with middleware support.
// ABOUTME: Provides group-based routing similar to Gin.

package websocket

import (
	"sync"

	"github.com/bingo-project/websocket/jsonrpc"
)

// Router routes JSON-RPC methods to handlers with middleware.
type Router struct {
	mu          sync.RWMutex
	middlewares []Middleware
	handlers    map[string]*handlerEntry
}

type handlerEntry struct {
	handler     Handler
	middlewares []Middleware
	compiled    Handler // cached compiled handler chain
}

// NewRouter creates a new Router.
func NewRouter() *Router {
	return &Router{
		handlers: make(map[string]*handlerEntry),
	}
}

// Use adds global middleware that applies to all handlers.
func (r *Router) Use(middlewares ...Middleware) *Router {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.middlewares = append(r.middlewares, middlewares...)
	// Invalidate compiled handlers
	for _, entry := range r.handlers {
		entry.compiled = nil
	}

	return r
}

// Handle registers a handler for a method with optional middlewares.
// Middlewares are applied in order: global -> group -> method-specific.
func (r *Router) Handle(method string, handler Handler, middlewares ...Middleware) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.handlers[method] = &handlerEntry{
		handler:     handler,
		middlewares: middlewares,
	}
}

// Dispatch routes a request to its handler.
func (r *Router) Dispatch(c *Context) *jsonrpc.Response {
	r.mu.RLock()
	entry, ok := r.handlers[c.Method]
	if !ok {
		r.mu.RUnlock()

		return jsonrpc.NewErrorResponse(c.Request.ID,
			NewError(404, "MethodNotFound", "Method not found: %s", c.Method))
	}

	// Compile handler chain if not cached
	if entry.compiled == nil {
		r.mu.RUnlock()
		r.mu.Lock()
		// Double-check after acquiring write lock
		if entry.compiled == nil {
			all := make([]Middleware, 0, len(r.middlewares)+len(entry.middlewares))
			all = append(all, r.middlewares...)
			all = append(all, entry.middlewares...)
			entry.compiled = Chain(all...)(entry.handler)
		}
		r.mu.Unlock()
		r.mu.RLock()
	}

	compiled := entry.compiled
	r.mu.RUnlock()

	return compiled(c)
}

// Methods returns all registered method names.
func (r *Router) Methods() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	methods := make([]string, 0, len(r.handlers))
	for method := range r.handlers {
		methods = append(methods, method)
	}

	return methods
}

// Group is a collection of handlers with shared middleware.
type Group struct {
	router      *Router
	middlewares []Middleware
}

// Group creates a new handler group with additional middleware.
func (r *Router) Group(middlewares ...Middleware) *Group {
	return &Group{
		router:      r,
		middlewares: middlewares,
	}
}

// Use adds middleware to this group.
func (g *Group) Use(middlewares ...Middleware) *Group {
	g.middlewares = append(g.middlewares, middlewares...)

	return g
}

// Handle registers a handler in this group.
func (g *Group) Handle(method string, handler Handler, middlewares ...Middleware) {
	// Combine group middlewares with handler-specific middlewares
	all := make([]Middleware, 0, len(g.middlewares)+len(middlewares))
	all = append(all, g.middlewares...)
	all = append(all, middlewares...)
	g.router.Handle(method, handler, all...)
}