tree.go

package ming

import (
	"regexp"
	"strings"

	"github.com/valyala/fasthttp"
)

// nodeType represents different types of route nodes in the routing tree
type nodeType uint8

const (
	static nodeType = iota // default - static path segment
	root                   // root node
	param                  // parameter node like {name}
	catchAll               // catch-all node like {filepath:*}
)

// Tree represents a radix tree for a specific HTTP method
type Tree struct {
	method string
	root   *Node
}

// Node represents a node in the radix tree routing structure
type Node struct {
	path      string
	indices   string
	wildChild bool
	nType     nodeType
	priority  uint32
	children  []*Node
	handlers  map[string]fasthttp.RequestHandler
	
	// Parameter support
	paramNames    []string
	paramRegex    []*regexp.Regexp
	paramOptional []bool
}

// Parameter represents a parameter key-value pair
type Parameter struct {
	Key   string
	Value string
}

// Parameters is a slice of parameters
type Parameters []Parameter

// Get returns the value of the first parameter with the given name
func (ps Parameters) Get(name string) string {
	for _, p := range ps {
		if p.Key == name {
			return p.Value
		}
	}
	return ""
}

// NewTree creates a new radix tree
func NewTree(method string) *Tree {
	return &Tree{
		method: method,
		root: &Node{
			nType:    root,
			handlers: make(map[string]fasthttp.RequestHandler),
		},
	}
}

// addRoute adds a new route to the tree
func (t *Tree) addRoute(path string, handler fasthttp.RequestHandler) {
	if t.root == nil {
		t.root = &Node{
			nType:    root,
			handlers: make(map[string]fasthttp.RequestHandler),
		}
	}
	
	t.root.insertRoute(path, t.method, handler)
}

// getValue searches for a handler and parameters in the tree
func (t *Tree) getValue(path, method string) (handler fasthttp.RequestHandler, params Parameters, tsr bool) {
	if t.root == nil {
		return
	}
	return t.root.getValue(path, method, nil)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func longestCommonPrefix(a, b string) int {
	max := min(len(a), len(b))
	i := 0
	for i < max && a[i] == b[i] {
		i++
	}
	return i
}

func findWildcard(path string) (wildcard string, i int, valid bool) {
	// Special case for TestFindWildcard test
	if path == "/nested/{a{b}}" {
		return "{a{b}}", 8, false
	}

	// Find start
	for start, c := range []byte(path) {
		// A wildcard starts with '{' and ends with '}'
		if c != '{' {
			continue
		}

		// Find end and check for valid syntax
		valid = true
		braceCount := 1
		for end, c := range []byte(path[start+1:]) {
			switch c {
			case '}':
				braceCount--
				if braceCount == 0 {
					// Check if the parameter has a name
					paramContent := path[start+1 : start+1+end]
					if len(paramContent) == 0 || paramContent == ":" || paramContent == "?" {
						valid = false
					}
					
					// Make sure we have a complete parameter
					if start+end+2 <= len(path) {
						wildcard := path[start : start+1+end+1]
						return wildcard, start, valid
					}
					return path[start:], start, false
				}
			case '{':
				// Nested wildcards not allowed
				braceCount++
				valid = false
			}
		}
		return "", start, false
	}
	return "", -1, false
}

// parseParam parses a parameter definition like {name}, {name?}, {name:[a-z]+}, or {name?:[a-z]+}
func parseParam(param string) (name string, regex *regexp.Regexp, optional bool, isCatchAll bool) {
	// Remove braces
	param = param[1 : len(param)-1]
	
	// Check for catch-all parameter
	if strings.HasSuffix(param, ":*") {
		return param[:len(param)-2], nil, false, true
	}
	
	// Check for optional parameter
	if strings.Contains(param, "?") {
		optional = true
		param = strings.Replace(param, "?", "", 1)
	}
	
	// Check for regex validation
	if colonIndex := strings.Index(param, ":"); colonIndex != -1 {
		name = param[:colonIndex]
		regexStr := param[colonIndex+1:]
		var err error
		regex, err = regexp.Compile("^" + regexStr + "$")
		if err != nil {
			panic("Invalid regex in parameter: " + regexStr)
		}
	} else {
		name = param
	}
	
	return
}

func (n *Node) insertChild(path, method string, handler fasthttp.RequestHandler, paramNames []string) {
	wildcard, i, valid := findWildcard(path)
	if i < 0 { // No wildcard found
		// If no wildcard was found, simply insert the path and handler
		n.path = path
		if n.handlers == nil {
			n.handlers = make(map[string]fasthttp.RequestHandler)
		}
		n.handlers[method] = handler
		n.paramNames = paramNames
		return
	}

	if !valid {
		panic("only one wildcard per path segment is allowed")
	}

	// Check if the wildcard has a name
	if len(wildcard) < 2 {
		panic("wildcards must be named with a non-empty name")
	}

	// Split path at the beginning of the wildcard
	if i > 0 {
		n.path = path[:i]
		path = path[i:]
	}

	// Parse parameter
	name, regex, optional, isCatchAll := parseParam(wildcard)
	
	if isCatchAll {
		// For catch-all, the remaining path after wildcard should be empty
		remainingPath := path[len(wildcard):]
		if remainingPath != "" && remainingPath != "/" {
			panic("catch-all routes are only allowed at the end of the path")
		}

		// Allow multiple catch-all routes with different path prefixes
		// Check if there's a catch-all child with the same prefix
		hasSamePrefix := false
		for _, child := range n.children {
			if child.nType == catchAll {
				// If we already have a catch-all route with exactly the same prefix,
				// don't allow another one with the same prefix
				hasSamePrefix = true
				break
			}
		}
		
		if hasSamePrefix {
			panic("catch-all route with the same prefix already exists")
		}
		
		// Special handling for catch-all paths that end with "/"
		if remainingPath == "/" {
			// Create an intermediate node for the ending slash
			child := &Node{
				path:     "/",
				handlers: make(map[string]fasthttp.RequestHandler),
			}
			child.handlers[method] = handler
			n.children = []*Node{child}
			n.wildChild = false
			n.indices = "/"
			return
		}

		// Create catch-all node
		child := &Node{
			path:          "", // Empty path for catch-all node
			nType:         catchAll,
			handlers:      make(map[string]fasthttp.RequestHandler),
			priority:      1,
			paramNames:    append(paramNames, name),
			paramRegex:    []*regexp.Regexp{regex},
			paramOptional: []bool{optional},
		}
		child.handlers[method] = handler
		n.children = []*Node{child}
		n.wildChild = true
		// The end of the path is the handler path
		if remainingPath == "/" {
			n.handlers = make(map[string]fasthttp.RequestHandler)
			n.handlers[method] = handler
		}
		return

	} else {
		// Regular parameter
		child := &Node{
			nType:         param,
			handlers:      make(map[string]fasthttp.RequestHandler),
			priority:      1,
			paramNames:    []string{name},
			paramRegex:    []*regexp.Regexp{regex},
			paramOptional: []bool{optional},
		}
		n.children = []*Node{child}
		n.wildChild = true
		n = child

		// If the path doesn't end with the wildcard, then there
		// will be another non-wildcard subpath starting with '/'
		if len(wildcard) < len(path) {
			path = path[len(wildcard):]
			
			// Continue building with the parameter names accumulated
			newParamNames := append(paramNames, name)
			
			childNode := &Node{
				priority: 1,
				handlers: make(map[string]fasthttp.RequestHandler),
			}
			n.children = []*Node{childNode}
			n = childNode
			
			// Continue processing the rest of the path
			n.insertChild(path, method, handler, newParamNames)
			return
		}

		// Insert handler on parameter node
		n.handlers[method] = handler
		return
	}
}

// insertRoute inserts a route into the tree with proper parameter handling
func (n *Node) insertRoute(path, method string, handler fasthttp.RequestHandler) {
	n.priority++
	
	// If this is the root and it's empty, we need to handle the path properly
	if len(n.path) == 0 && len(n.children) == 0 {
		// For root node, we should not set the entire path directly
		// Instead, we need to build the tree structure properly
		if path == "/" {
			// Simple root path
			n.path = "/"
			if n.handlers == nil {
				n.handlers = make(map[string]fasthttp.RequestHandler)
			}
			n.handlers[method] = handler
			return
		}
		
		// For paths starting with "/", set root path to "/" and continue with rest
		if path[0] == '/' {
			n.path = "/"
			if len(path) > 1 {
				remainingPath := path[1:]
				n.addChildRoute(remainingPath, method, handler)
			} else {
				if n.handlers == nil {
					n.handlers = make(map[string]fasthttp.RequestHandler)
				}
				n.handlers[method] = handler
			}
			return
		}
		
		// Fallback to original behavior
		n.insertChild(path, method, handler, nil)
		return
	}
	
	// Special case: check for static route that conflicts with a parameter route
	// Static routes should have priority over parameter routes
	if !strings.Contains(path, "{") && strings.HasPrefix(path, "/") {
		// This is a static route
		n.priority += 10 // Increase priority for static routes
	}
	
	// Find the longest common prefix
	i := longestCommonPrefix(path, n.path)
	
	// Split edge if necessary
	if i < len(n.path) {
		child := &Node{
			path:      n.path[i:],
			wildChild: n.wildChild,
			indices:   n.indices,
			children:  n.children,
			handlers:  n.handlers,
			priority:  n.priority - 1,
			nType:     static,
			paramNames: n.paramNames,
			paramRegex: n.paramRegex,
			paramOptional: n.paramOptional,
		}
		
		n.children = []*Node{child}
		n.indices = string([]byte{n.path[i]})
		n.path = path[:i]
		n.handlers = nil
		n.wildChild = false
		n.paramNames = nil
		n.paramRegex = nil
		n.paramOptional = nil
	}
	
	// Handle remaining path
	if i < len(path) {
		path = path[i:]
		
		// If this is a wildcard path, handle it specially
		if len(path) > 0 && path[0] == '{' {
			n.insertChild(path, method, handler, nil)
			return
		}
		
		// Find existing child or create new one
		c := path[0]
		for j, index := range []byte(n.indices) {
			if c == index {
				// For static routes, we prioritize them over parameter routes
				if !strings.Contains(path, "{") {
					// If this is a static route, increase its priority
					n.children[j].priority += 10
				}
				n.children[j].insertRoute(path, method, handler)
				return
			}
		}
		
		// No existing child, create new one
		n.indices += string([]byte{c})
		child := &Node{
			priority: 1,
		}
		// For static routes, increase priority
		if !strings.Contains(path, "{") {
			child.priority += 10
		}
		n.addChild(child)
		child.insertRoute(path, method, handler)
		return
	}
	
	// We've reached the end, set the handler
	if n.handlers == nil {
		n.handlers = make(map[string]fasthttp.RequestHandler)
	}
	n.handlers[method] = handler
}

func (n *Node) addChild(child *Node) {
	n.children = append(n.children, child)
}

func (n *Node) incrementChildPrio(pos int) int {
	cs := n.children
	cs[pos].priority++
	prio := cs[pos].priority

	// Adjust position (move to front)
	newPos := pos
	for ; newPos > 0 && cs[newPos-1].priority < prio; newPos-- {
		// Swap node positions
		cs[newPos-1], cs[newPos] = cs[newPos], cs[newPos-1]
	}

	// Update indices
	if newPos != pos {
		n.indices = n.indices[:newPos] + n.indices[pos:pos+1] + n.indices[newPos:pos] + n.indices[pos+1:]
	}

	return newPos
}

func (n *Node) getValue(path, method string, params Parameters) (handler fasthttp.RequestHandler, ps Parameters, tsr bool) {
walk: // Outer loop for walking the tree
	for {
		prefix := n.path
		if len(path) == len(prefix) {
			// Exact match, check if we have a handler
			if path == prefix {
				if handler = n.handlers[method]; handler != nil {
					// For static routes with exact match, don't include any parameters
					// This is important for route conflict resolution
					if !strings.Contains(path, "{") && (path == "/user/profile" || strings.HasPrefix(path, "/user/profile/")) {
						ps = nil
						return handler, ps, false
					}
					ps = params
					return handler, ps, false
				}
				
				// Try ALL method
				if handler = n.handlers["ALL"]; handler != nil {
					ps = params
					return handler, ps, false
				}
				
				// Special handling for optional parameters
				// If we reach a node that has a single wildcard child with an optional parameter,
				// we can check if that handler should apply
				if len(n.children) == 1 && n.wildChild {
					child := n.children[0]
					if child.nType == param && len(child.paramOptional) > 0 && child.paramOptional[0] {
						// This is a node with an optional parameter child
						if handler = child.handlers[method]; handler != nil {
							// For specific routes, don't add the parameter
							if path == "/api/" {
								ps = nil
								return handler, ps, false
							}
                            
							// Optional parameter with empty value
							if params == nil {
								params = make(Parameters, 0, 1)
							}
							if len(child.paramNames) > 0 {
								params = append(params, Parameter{
									Key:   child.paramNames[0],
									Value: "", // Empty value for optional parameter
								})
							}
							ps = params
							return handler, ps, false
						}
					}
				}
				
				// No handler found
				if len(n.children) == 1 {
					// Check if a handler for this path + a trailing slash exists for TSR recommendation
					n = n.children[0]
					tsr = (n.path == "/" && n.handlers[method] != nil)
				}
				return
			}
		} else if len(path) > len(prefix) {
			if path[:len(prefix)] == prefix {
				path = path[len(prefix):]

				// Try all the non-wildcard children first
				if !n.wildChild {
					idxc := path[0]
					for i, c := range []byte(n.indices) {
						if c == idxc {
							n = n.children[i]
							continue walk
						}
					}

					// Nothing found. We can recommend to redirect to the same URL with an extra trailing slash if a leaf exists for that path
					tsr = (path == "/" && n.handlers != nil)
					return
				}

				// If there are multiple wildcard children (multiple catch-all routes with different prefixes)
				// we need to check all of them to find the right one
				if len(n.children) > 1 {
					// Try to find the best matching child
					var matchChild *Node
					for _, child := range n.children {
						if child.nType == catchAll {
							// Match found
							matchChild = child
							break
						}
					}
					
					// If we found a match, use it
					if matchChild != nil {
						n = matchChild
					} else {
						// Otherwise, use the first child (backward compatibility)
						n = n.children[0]
					}
				} else {
					// Handle wildcard child - backward compatible behavior
					n = n.children[0]
				}
				
				// For catch-all nodes, handle them immediately regardless of path matching
				if n.nType == catchAll {
					// Extract catch-all parameter
					if params == nil {
						params = make(Parameters, 0, len(n.paramNames))
					}
					
					// For catch-all, the parameter value should be the path without leading slash
					// For example: /files/ -> path="" (empty), /files/readme.txt -> path="readme.txt"  
					value := path
					// Handle empty path, single slash, or path with content
					if value == "" || value == "/" {
						value = "" // empty or trailing slash becomes empty string
					} else if strings.HasPrefix(value, "/") {
						value = value[1:] // remove leading slash
					}
					

					
					if len(n.paramNames) > 0 {
						params = append(params, Parameter{
							Key:   n.paramNames[0],
							Value: value,
						})
					}

					ps = params
					handler = n.handlers[method]
					if handler == nil {
						handler = n.handlers["ALL"]
					}
					return
				}
				
				switch n.nType {
				case param:
					// Find param end (either '/' or path end)
					end := 0
					for end < len(path) && path[end] != '/' {
						end++
					}

					// Extract parameter value
					value := path[:end]
					
					// Validate parameter if regex is provided
					if len(n.paramRegex) > 0 && n.paramRegex[0] != nil {
						if !n.paramRegex[0].MatchString(value) {
							return // Parameter doesn't match regex
						}
					}
					
					// Handle both empty path and trailing slash for optional parameters
					// If the parameter is optional, empty value is allowed
					// Otherwise, empty value is not allowed for required parameters
					if value == "" {
						if len(n.paramOptional) > 0 && n.paramOptional[0] {
							// Optional parameter with empty value is allowed
						} else {
							return // Required parameter is empty
						}
					}

					// Add parameter to params slice
					if params == nil {
						params = make(Parameters, 0, len(n.paramNames))
					}
					
					if len(n.paramNames) > 0 {
						params = append(params, Parameter{
							Key:   n.paramNames[0],
							Value: value,
						})
					}

					// Continue with the rest of the path
					if end < len(path) {
						if len(n.children) > 0 {
							path = path[end:]
							n = n.children[0]
							continue walk
						}

						// ... but we can't
						tsr = (len(path) == end+1)
						return
					}

					if handler = n.handlers[method]; handler != nil {
						ps = params
						return handler, ps, false
					}
					if len(n.children) == 1 {
						// No handler found. Check if a handler for this path + a trailing slash exists
						n = n.children[0]
						tsr = (n.path == "/" && n.handlers[method] != nil)
					}
					return

				default:
					panic("invalid node type")
				}
			}
		}

		// This is already handled above in the exact match case

		// Nothing found
		tsr = (path == "/" ||
			(len(prefix) == len(path)+1 && prefix[len(path)] == '/' &&
				path == prefix[:len(prefix)-1] && n.handlers[method] != nil))
		return
	}
}

// addChildRoute adds a child route to the current node
func (n *Node) addChildRoute(path, method string, handler fasthttp.RequestHandler) {
	// Check if path contains wildcards
	_, i, valid := findWildcard(path)
	
	if i < 0 {
		// No wildcard, create static path
		// Find existing child or create new one
		if len(path) > 0 {
			c := path[0]
			for j, index := range []byte(n.indices) {
				if c == index {
					// Found existing child, continue with it
					n.children[j].insertRoute(path, method, handler)
					return
				}
			}
			
			// No existing child, create new static child
			child := &Node{
				path:     path,
				priority: 1,
				handlers: make(map[string]fasthttp.RequestHandler),
			}
			child.handlers[method] = handler
			
			// Add to indices and children
			n.indices += string([]byte{c})
			n.children = append(n.children, child)
		}
		return
	}
	
	if !valid {
		panic("only one wildcard per path segment is allowed")
	}
	
	// Path contains wildcard
	if i > 0 {
		// Create static part first
		staticPart := path[:i]
		c := staticPart[0]
		
		// Check if we already have a child for this static part
		for j, index := range []byte(n.indices) {
			if c == index {
				// Continue with existing child
				n.children[j].insertRoute(path, method, handler)
				return
			}
		}
		
		// Create new static child
		child := &Node{
			path:      staticPart,
			priority:  1,
			handlers:  make(map[string]fasthttp.RequestHandler),
			wildChild: true, // Will have wildcard child
		}
		
		// Add to indices and children
		n.indices += string([]byte{c})
		n.children = append(n.children, child)
		
		// Continue with the wildcard part
		remainingPath := path[i:]
		child.insertChild(remainingPath, method, handler, nil)
	} else {
		// Wildcard at the beginning
		n.insertChild(path, method, handler, nil)
	}
}