Moonshark/core/routers/luarouter.go

package routers

import (
	"errors"
	"os"
	"path/filepath"
	"strings"
	"sync"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// Maximum number of URL parameters per route
const maxParams = 20

// LuaRouter is a filesystem-based HTTP router for Lua files
type LuaRouter struct {
	routesDir string           // Root directory containing route files
	routes    map[string]*node // Method -> route tree
	mu        sync.RWMutex     // Lock for concurrent access to routes
}

// node represents a node in the routing trie
type node struct {
	handler     string           // Path to Lua file (empty if not an endpoint)
	bytecode    []byte           // Pre-compiled Lua bytecode
	paramName   string           // Parameter name (if this is a parameter node)
	staticChild map[string]*node // Static children by segment name
	paramChild  *node            // Parameter/wildcard child
}

// Params holds URL parameters with fixed-size arrays to avoid allocations
type Params struct {
	Keys   [maxParams]string
	Values [maxParams]string
	Count  int
}

// Get returns a parameter value by name
func (p *Params) Get(name string) string {
	for i := 0; i < p.Count; i++ {
		if p.Keys[i] == name {
			return p.Values[i]
		}
	}
	return ""
}

// NewLuaRouter creates a new LuaRouter instance
func NewLuaRouter(routesDir string) (*LuaRouter, error) {
	// Verify routes directory exists
	info, err := os.Stat(routesDir)
	if err != nil {
		return nil, err
	}
	if !info.IsDir() {
		return nil, errors.New("routes path is not a directory")
	}

	r := &LuaRouter{
		routesDir: routesDir,
		routes:    make(map[string]*node),
	}

	// Initialize method trees
	methods := []string{"GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"}
	for _, method := range methods {
		r.routes[method] = &node{
			staticChild: make(map[string]*node),
		}
	}

	// Build routes
	if err := r.buildRoutes(); err != nil {
		return nil, err
	}

	return r, nil
}

// buildRoutes scans the routes directory and builds the routing tree
func (r *LuaRouter) buildRoutes() error {
	return filepath.Walk(r.routesDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip directories
		if info.IsDir() {
			return nil
		}

		// Only process .lua files
		if !strings.HasSuffix(info.Name(), ".lua") {
			return nil
		}

		// Extract method from filename
		method := strings.ToUpper(strings.TrimSuffix(info.Name(), ".lua"))

		// Check if valid method
		root, exists := r.routes[method]
		if !exists {
			return nil // Skip invalid methods
		}

		// Get relative path for URL
		relDir, err := filepath.Rel(r.routesDir, filepath.Dir(path))
		if err != nil {
			return err
		}

		// Build URL path
		urlPath := "/"
		if relDir != "." {
			urlPath = "/" + strings.ReplaceAll(relDir, "\\", "/")
		}

		// Add route to tree
		return r.addRoute(root, urlPath, path)
	})
}

// addRoute adds a route to the routing tree and compiles the Lua file to bytecode
func (r *LuaRouter) addRoute(root *node, urlPath, handlerPath string) error {
	segments := strings.Split(strings.Trim(urlPath, "/"), "/")
	current := root

	for _, segment := range segments {
		if len(segment) >= 2 && segment[0] == '[' && segment[len(segment)-1] == ']' {
			if current.paramChild == nil {
				current.paramChild = &node{
					paramName:   segment[1 : len(segment)-1],
					staticChild: make(map[string]*node),
				}
			}
			current = current.paramChild
		} else {
			// Create or get static child
			child, exists := current.staticChild[segment]
			if !exists {
				child = &node{
					staticChild: make(map[string]*node),
				}
				current.staticChild[segment] = child
			}
			current = child
		}
	}

	// Set handler path
	current.handler = handlerPath

	// Compile Lua file to bytecode
	if err := r.compileHandler(current); err != nil {
		return err
	}

	return nil
}

// Match finds a handler for the given method and path
// Uses the pre-allocated params struct to avoid allocations
func (r *LuaRouter) Match(method, path string, params *Params) (*node, bool) {
	// Reset params
	params.Count = 0

	// Get route tree for method
	r.mu.RLock()
	root, exists := r.routes[method]
	r.mu.RUnlock()

	if !exists {
		return nil, false
	}

	// Split path
	segments := strings.Split(strings.Trim(path, "/"), "/")

	// Match path
	return r.matchPath(root, segments, params, 0)
}

// matchPath recursively matches a path against the routing tree
func (r *LuaRouter) matchPath(current *node, segments []string, params *Params, depth int) (*node, bool) {
	// Base case: no more segments
	if len(segments) == 0 {
		if current.handler != "" {
			return current, true
		}
		return nil, false
	}

	segment := segments[0]
	remaining := segments[1:]

	// Try static child first (exact match takes precedence)
	if child, exists := current.staticChild[segment]; exists {
		if node, found := r.matchPath(child, remaining, params, depth+1); found {
			return node, true
		}
	}

	// Try parameter child
	if current.paramChild != nil {
		// Store parameter
		if params.Count < maxParams {
			params.Keys[params.Count] = current.paramChild.paramName
			params.Values[params.Count] = segment
			params.Count++
		}

		if node, found := r.matchPath(current.paramChild, remaining, params, depth+1); found {
			return node, true
		}

		// Backtrack: remove parameter if no match
		params.Count--
	}

	return nil, false
}

// compileHandler compiles a Lua file to bytecode
func (r *LuaRouter) compileHandler(n *node) error {
	if n.handler == "" {
		return nil
	}

	// Read the Lua file
	content, err := os.ReadFile(n.handler)
	if err != nil {
		return err
	}

	// Compile to bytecode
	state := luajit.New()
	if state == nil {
		return errors.New("failed to create Lua state")
	}
	defer state.Close()

	bytecode, err := state.CompileBytecode(string(content), n.handler)
	if err != nil {
		return err
	}

	// Store bytecode in the node
	n.bytecode = bytecode
	return nil
}

// GetBytecode returns the compiled bytecode for a matched route
func (r *LuaRouter) GetBytecode(method, path string, params *Params) ([]byte, string, bool) {
	node, found := r.Match(method, path, params)
	if !found {
		return nil, "", false
	}
	return node.bytecode, node.handler, true
}

// Refresh rebuilds the router by rescanning the routes directory
func (r *LuaRouter) Refresh() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	// Reset routes
	for method := range r.routes {
		r.routes[method] = &node{
			staticChild: make(map[string]*node),
		}
	}

	// Rebuild routes
	return r.buildRoutes()
}