Moonshark/core/routers/LuaRouter.go

package routers

import (
	"encoding/binary"
	"errors"
	"hash/fnv"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
	"github.com/VictoriaMetrics/fastcache"
)

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

// Default cache sizes
const (
	defaultBytecodeMaxBytes = 32 * 1024 * 1024 // 32MB for bytecode cache
	defaultRouteMaxBytes    = 8 * 1024 * 1024  // 8MB for route match cache
)

// 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 ""
}

// 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
	failedRoutes map[string]*RouteError // Track failed routes
	mu           sync.RWMutex           // Lock for concurrent access to routes

	// Cache for route matches and bytecode
	routeCache    *fastcache.Cache // Cache for route lookups
	bytecodeCache *fastcache.Cache // Cache for compiled bytecode
}

// node represents a node in the routing trie
type node struct {
	handler     string           // Path to Lua file (empty if not an endpoint)
	paramName   string           // Parameter name (if this is a parameter node)
	staticChild map[string]*node // Static children by segment name
	paramChild  *node            // Parameter/wildcard child
	err         error            // Compilation error if any
	modTime     time.Time        // Last modification time
}

// 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),
		failedRoutes:  make(map[string]*RouteError),
		routeCache:    fastcache.New(defaultRouteMaxBytes),
		bytecodeCache: fastcache.New(defaultBytecodeMaxBytes),
	}

	// 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
	err = r.buildRoutes()

	// If some routes failed to compile, return the router with a warning error
	if len(r.failedRoutes) > 0 {
		return r, ErrRoutesCompilationErrors
	}

	return r, err
}

// buildRoutes scans the routes directory and builds the routing tree
func (r *LuaRouter) buildRoutes() error {
	// Clear failed routes map
	r.failedRoutes = make(map[string]*RouteError)

	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 - continue even if there are errors
		r.addRoute(root, urlPath, path, info.ModTime())
		return nil
	})
}

// addRoute adds a route to the routing tree
func (r *LuaRouter) addRoute(root *node, urlPath, handlerPath string, modTime time.Time) 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 and mod time
	current.handler = handlerPath
	current.modTime = modTime

	// Compile Lua file to bytecode
	if err := r.compileHandler(current); err != nil {
		// Track the failure but don't fail the entire process
		routeKey := getRouteKey(urlPath, handlerPath)
		r.failedRoutes[routeKey] = &RouteError{
			Path:       urlPath,
			ScriptPath: handlerPath,
			Err:        err,
		}
	}

	return nil
}

// getRouteKey generates a unique key for a route
func getRouteKey(path, handler string) string {
	return path + ":" + handler
}

// hashString generates a hash for a string
func hashString(s string) uint64 {
	h := fnv.New64a()
	h.Write([]byte(s))
	return h.Sum64()
}

// uint64ToBytes converts a uint64 to bytes for cache key
func uint64ToBytes(n uint64) []byte {
	b := make([]byte, 8)
	binary.LittleEndian.PutUint64(b, n)
	return b
}

// bytesToUint64 converts bytes to uint64
func bytesToUint64(b []byte) uint64 {
	return binary.LittleEndian.Uint64(b)
}

// getCacheKey generates a cache key for a method and path
func getCacheKey(method, path string) []byte {
	// Simple concatenation with separator to create a unique key
	key := hashString(method + ":" + path)
	return uint64ToBytes(key)
}

// getBytecodeKey generates a cache key for a handler path
func getBytecodeKey(handlerPath string) []byte {
	key := hashString(handlerPath)
	return uint64ToBytes(key)
}

// 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 {
		n.err = err // Store the error in the node
		return err
	}

	// Compile to bytecode
	state := luajit.New()
	if state == nil {
		compileErr := errors.New("failed to create Lua state")
		n.err = compileErr // Store the error in the node
		return compileErr
	}
	defer state.Close()

	bytecode, err := state.CompileBytecode(string(content), n.handler)
	if err != nil {
		n.err = err // Store the error in the node
		return err
	}

	// Store bytecode in cache
	bytecodeKey := getBytecodeKey(n.handler)
	r.bytecodeCache.Set(bytecodeKey, bytecode)

	n.err = nil // Clear any previous error
	return nil
}

// GetBytecode returns the compiled bytecode for a matched route
// Uses FastCache for both route matching and bytecode retrieval
func (r *LuaRouter) GetBytecode(method, path string, params *Params) ([]byte, string, bool) {
	// Check route cache first
	routeCacheKey := getCacheKey(method, path)
	routeCacheData := r.routeCache.Get(nil, routeCacheKey)

	if len(routeCacheData) > 0 {
		// Cache hit - first 8 bytes are bytecode hash, rest is handler path
		handlerPath := string(routeCacheData[8:])
		bytecodeKey := routeCacheData[:8]

		// Get bytecode from cache
		bytecode := r.bytecodeCache.Get(nil, bytecodeKey)
		if len(bytecode) > 0 {
			return bytecode, handlerPath, true
		}

		// Bytecode not found, check if file was modified
		n, exists := r.nodeForHandler(handlerPath)
		if !exists {
			// Handler no longer exists
			r.routeCache.Del(routeCacheKey)
			return nil, "", false
		}

		// Check if file was modified
		fileInfo, err := os.Stat(handlerPath)
		if err != nil || fileInfo.ModTime().After(n.modTime) {
			// Recompile if file was modified
			if err := r.compileHandler(n); err != nil {
				return nil, handlerPath, true // Return with error
			}

			// Update cache
			newBytecodeKey := getBytecodeKey(handlerPath)
			bytecode = r.bytecodeCache.Get(nil, newBytecodeKey)

			// Update route cache
			newCacheData := make([]byte, 8+len(handlerPath))
			copy(newCacheData[:8], newBytecodeKey)
			copy(newCacheData[8:], handlerPath)
			r.routeCache.Set(routeCacheKey, newCacheData)

			return bytecode, handlerPath, true
		}

		// Strange case - bytecode not in cache but file not modified
		// Recompile
		if err := r.compileHandler(n); err != nil {
			return nil, handlerPath, true
		}
		bytecode = r.bytecodeCache.Get(nil, bytecodeKey)
		return bytecode, handlerPath, true
	}

	// Cache miss - do normal routing
	node, found := r.Match(method, path, params)
	if !found {
		return nil, "", false
	}

	// If the route exists but has a compilation error
	if node.err != nil {
		return nil, node.handler, true
	}

	// Get bytecode from cache
	bytecodeKey := getBytecodeKey(node.handler)
	bytecode := r.bytecodeCache.Get(nil, bytecodeKey)

	if len(bytecode) == 0 {
		// Compile if not in cache
		if err := r.compileHandler(node); err != nil {
			return nil, node.handler, true
		}
		bytecode = r.bytecodeCache.Get(nil, bytecodeKey)
	}

	// Add to route cache
	cacheData := make([]byte, 8+len(node.handler))
	copy(cacheData[:8], bytecodeKey)
	copy(cacheData[8:], node.handler)
	r.routeCache.Set(routeCacheKey, cacheData)

	return bytecode, node.handler, true
}

// nodeForHandler finds a node by its handler path
func (r *LuaRouter) nodeForHandler(handlerPath string) (*node, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	for _, root := range r.routes {
		if node := findNodeByHandler(root, handlerPath); node != nil {
			return node, true
		}
	}

	return nil, false
}

// findNodeByHandler finds a node by its handler path
func findNodeByHandler(current *node, handlerPath string) *node {
	if current == nil {
		return nil
	}

	if current.handler == handlerPath {
		return current
	}

	// Check static children
	for _, child := range current.staticChild {
		if node := findNodeByHandler(child, handlerPath); node != nil {
			return node
		}
	}

	// Check param child
	if current.paramChild != nil {
		if node := findNodeByHandler(current.paramChild, handlerPath); node != nil {
			return node
		}
	}

	return nil
}

// 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),
		}
	}

	// Clear failed routes
	r.failedRoutes = make(map[string]*RouteError)

	// Rebuild routes
	err := r.buildRoutes()

	// If some routes failed to compile, return a warning error
	if len(r.failedRoutes) > 0 {
		return ErrRoutesCompilationErrors
	}

	return err
}

// ReportFailedRoutes returns a list of routes that failed to compile
func (r *LuaRouter) ReportFailedRoutes() []*RouteError {
	r.mu.RLock()
	defer r.mu.RUnlock()

	result := make([]*RouteError, 0, len(r.failedRoutes))
	for _, re := range r.failedRoutes {
		result = append(result, re)
	}

	return result
}

// ClearCache clears both the route and bytecode caches
func (r *LuaRouter) ClearCache() {
	r.routeCache.Reset()
	r.bytecodeCache.Reset()
}

// GetCacheStats returns statistics about the cache
func (r *LuaRouter) GetCacheStats() map[string]interface{} {
	var routeStats fastcache.Stats
	var bytecodeStats fastcache.Stats

	r.routeCache.UpdateStats(&routeStats)
	r.bytecodeCache.UpdateStats(&bytecodeStats)

	return map[string]interface{}{
		"routeEntries":       routeStats.EntriesCount,
		"routeBytes":         routeStats.BytesSize,
		"routeCollisions":    routeStats.Collisions,
		"bytecodeEntries":    bytecodeStats.EntriesCount,
		"bytecodeBytes":      bytecodeStats.BytesSize,
		"bytecodeCollisions": bytecodeStats.Collisions,
	}
}

// GetRouteStats returns statistics about the router
func (r *LuaRouter) GetRouteStats() (int, int64) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	routeCount := 0
	bytecodeBytes := int64(0)

	// Count routes and estimate bytecode size
	for _, root := range r.routes {
		count, bytes := countNodesAndBytecode(root)
		routeCount += count
		bytecodeBytes += bytes
	}

	return routeCount, bytecodeBytes
}

// countNodesAndBytecode traverses the tree and counts nodes and bytecode size
func countNodesAndBytecode(n *node) (count int, bytecodeBytes int64) {
	if n == nil {
		return 0, 0
	}

	// Count this node if it has a handler
	if n.handler != "" {
		count = 1
		// Estimate bytecode size (average of 2KB per script)
		bytecodeBytes = 2048
	}

	// Count static children
	for _, child := range n.staticChild {
		childCount, childBytes := countNodesAndBytecode(child)
		count += childCount
		bytecodeBytes += childBytes
	}

	// Count parameter child
	if n.paramChild != nil {
		childCount, childBytes := countNodesAndBytecode(n.paramChild)
		count += childCount
		bytecodeBytes += childBytes
	}

	return count, bytecodeBytes
}

type NodeWithError struct {
	ScriptPath string
	Error      error
}

// GetNodeWithError returns the node with its error for a given path
func (r *LuaRouter) GetNodeWithError(method, path string, params *Params) (*NodeWithError, bool) {
	// Try route cache first
	routeCacheKey := getCacheKey(method, path)
	routeCacheData := r.routeCache.Get(nil, routeCacheKey)

	if len(routeCacheData) > 0 {
		// Cache hit - get handler path
		handlerPath := string(routeCacheData[8:])

		// Find the node for this handler
		node, found := r.nodeForHandler(handlerPath)
		if found {
			return &NodeWithError{
				ScriptPath: node.handler,
				Error:      node.err,
			}, true
		}
	}

	// Cache miss - do normal routing
	node, found := r.Match(method, path, params)
	if !found {
		return nil, false
	}

	return &NodeWithError{
		ScriptPath: node.handler,
		Error:      node.err,
	}, true
}