Moonshark/http/router/router.go

package router

import (
	"errors"
	"strings"
	"sync"
)

// node represents a node in the radix trie
type node struct {
	segment    string
	handler    int // Lua function reference
	children   []*node
	isDynamic  bool // :param
	isWildcard bool // *param
	paramName  string
}

// Router is a string-based HTTP router with efficient lookup
type Router struct {
	get, post, put, patch, delete *node
	mu                            sync.RWMutex
}

// Params holds URL parameters
type Params struct {
	Keys   []string
	Values []string
}

// Get returns a parameter value by name
func (p *Params) Get(name string) string {
	for i, key := range p.Keys {
		if key == name && i < len(p.Values) {
			return p.Values[i]
		}
	}
	return ""
}

// New creates a new Router instance
func New() *Router {
	return &Router{
		get:    &node{},
		post:   &node{},
		put:    &node{},
		patch:  &node{},
		delete: &node{},
	}
}

// methodNode returns the root node for a method
func (r *Router) methodNode(method string) *node {
	switch method {
	case "GET":
		return r.get
	case "POST":
		return r.post
	case "PUT":
		return r.put
	case "PATCH":
		return r.patch
	case "DELETE":
		return r.delete
	default:
		return nil
	}
}

// AddRoute adds a new route with handler reference
func (r *Router) AddRoute(method, path string, handlerRef int) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	root := r.methodNode(method)
	if root == nil {
		return errors.New("unsupported HTTP method")
	}

	if path == "/" {
		root.handler = handlerRef
		return nil
	}

	return r.addRoute(root, path, handlerRef)
}

// addRoute adds a route to the trie
func (r *Router) addRoute(root *node, path string, handlerRef int) error {
	segments := r.parseSegments(path)
	current := root

	for _, seg := range segments {
		isDyn := strings.HasPrefix(seg, ":")
		isWC := strings.HasPrefix(seg, "*")

		if isWC && seg != segments[len(segments)-1] {
			return errors.New("wildcard must be the last segment")
		}

		paramName := ""
		if isDyn {
			paramName = seg[1:]
			seg = ":"
		} else if isWC {
			paramName = seg[1:]
			seg = "*"
		}

		// Find or create child
		var child *node
		for _, c := range current.children {
			if c.segment == seg {
				child = c
				break
			}
		}

		if child == nil {
			child = &node{
				segment:    seg,
				isDynamic:  isDyn,
				isWildcard: isWC,
				paramName:  paramName,
			}
			current.children = append(current.children, child)
		}

		current = child
	}

	current.handler = handlerRef
	return nil
}

// parseSegments splits path into segments
func (r *Router) parseSegments(path string) []string {
	segments := strings.Split(strings.Trim(path, "/"), "/")
	var result []string
	for _, seg := range segments {
		if seg != "" {
			result = append(result, seg)
		}
	}
	return result
}

// Lookup finds handler and parameters for a method and path
func (r *Router) Lookup(method, path string) (int, *Params, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	root := r.methodNode(method)
	if root == nil {
		return 0, nil, false
	}

	if path == "/" {
		if root.handler != 0 {
			return root.handler, &Params{}, true
		}
		return 0, nil, false
	}

	segments := r.parseSegments(path)
	handler, params := r.match(root, segments, 0)
	if handler == 0 {
		return 0, nil, false
	}

	return handler, params, true
}

// match traverses the trie to find handler
func (r *Router) match(current *node, segments []string, index int) (int, *Params) {
	if index >= len(segments) {
		if current.handler != 0 {
			return current.handler, &Params{}
		}
		return 0, nil
	}

	segment := segments[index]

	// Check exact match first
	for _, child := range current.children {
		if child.segment == segment {
			handler, params := r.match(child, segments, index+1)
			if handler != 0 {
				return handler, params
			}
		}
	}

	// Check dynamic match second
	for _, child := range current.children {
		if child.isDynamic {
			handler, params := r.match(child, segments, index+1)
			if handler != 0 {
				// Prepend this parameter
				newParams := &Params{
					Keys:   append([]string{child.paramName}, params.Keys...),
					Values: append([]string{segment}, params.Values...),
				}
				return handler, newParams
			}
		}
	}

	// Check wildcard last (catches everything remaining)
	for _, child := range current.children {
		if child.isWildcard {
			remaining := strings.Join(segments[index:], "/")
			return child.handler, &Params{
				Keys:   []string{child.paramName},
				Values: []string{remaining},
			}
		}
	}

	return 0, nil
}

// RemoveRoute removes a route
func (r *Router) RemoveRoute(method, path string) {
	r.mu.Lock()
	defer r.mu.Unlock()

	root := r.methodNode(method)
	if root == nil {
		return
	}

	if path == "/" {
		root.handler = 0
		return
	}

	segments := r.parseSegments(path)
	r.removeRoute(root, segments, 0)
}

// removeRoute removes a route from the trie
func (r *Router) removeRoute(current *node, segments []string, index int) {
	if index >= len(segments) {
		current.handler = 0
		return
	}

	segment := segments[index]

	for _, child := range current.children {
		if child.segment == segment ||
			(child.isDynamic && strings.HasPrefix(segment, ":")) ||
			(child.isWildcard && strings.HasPrefix(segment, "*")) {
			r.removeRoute(child, segments, index+1)
			break
		}
	}
}