神奇代码

package httpserver

import (
	"net/http"
	"sort"
	"zqzd-safe/modules/gateway/pkg/ast"

	iradix "github.com/hashicorp/go-immutable-radix/v2"
)

// RouteIndexer 路由索引器，用于快速查找路由
type RouteIndexer struct {
	// Host索引: Host -> PathIndexer
	hostIndex map[string]*PathIndexer

	// 默认路径索引（没有Host限制的路由）
	defaultPathIndex *PathIndexer

	// 全量路由列表（fallback用）
	allRoutes []*RouteEntry
}

// PathIndexer 路径索引器（使用基数树）
type PathIndexer struct {
	radixTree *iradix.Tree[*[]*RouteEntry]
}

// NewRouteIndexer 创建路由索引器
func NewRouteIndexer() *RouteIndexer {
	return &RouteIndexer{
		hostIndex:        make(map[string]*PathIndexer),
		defaultPathIndex: &PathIndexer{radixTree: iradix.New[*[]*RouteEntry]()},
		allRoutes:        make([]*RouteEntry, 0),
	}
}

// NewPathIndexer 创建路径索引器
func NewPathIndexer() *PathIndexer {
	return &PathIndexer{radixTree: iradix.New[*[]*RouteEntry]()}
}

// IndexRoute 索引一个路由
func (idx *RouteIndexer) IndexRoute(entry *RouteEntry) {
	// 保存到全量列表
	idx.allRoutes = append(idx.allRoutes, entry)

	// 对全量列表按优先级排序
	sortRoutesByPriority(idx.allRoutes)

	// 分析AST，提取可索引的信息
	hosts, pathPrefixes := analyzeAST(entry.AST)

	// 如果有Host限制，添加到Host索引
	if len(hosts) > 0 {
		for _, host := range hosts {
			pathIndexer, exists := idx.hostIndex[host]
			if !exists {
				pathIndexer = NewPathIndexer()
				idx.hostIndex[host] = pathIndexer
			}
			pathIndexer.addRoute(pathPrefixes, entry)
		}
	} else {
		// 没有Host限制，添加到默认索引
		idx.defaultPathIndex.addRoute(pathPrefixes, entry)
	}
}

// addRoute 将路由添加到路径索引中
func (pi *PathIndexer) addRoute(pathPrefixes []string, entry *RouteEntry) {
	// 如果没有路径前缀，使用"/"作为key
	if len(pathPrefixes) == 0 {
		pathPrefixes = []string{"/"}
	}

	for _, prefix := range pathPrefixes {
		key := []byte(prefix)
		existing, _ := pi.radixTree.Get(key)

		var routes []*RouteEntry
		if existing != nil {
			routes = *existing
		}

		routes = append(routes, entry)

		// 按优先级排序（优先级高的在前）
		sortRoutesByPriority(routes)

		tree, _, _ := pi.radixTree.Insert(key, &routes)
		pi.radixTree = tree
	}
}

// FindRoute 查找匹配的路由（优先级优先 + 索引优化）
func (idx *RouteIndexer) FindRoute(r *http.Request) *RouteEntry {
	// 策略：
	// 1. 先用索引找到候选路由中的最高优先级
	// 2. 然后遍历全量列表，只检查优先级 >= 候选最高优先级的路由
	// 3. 这样既保证了优先级正确性，又利用了索引优化

	var maxCandidatePriority int = -1
	var candidateMatch *RouteEntry

	// 第一步：根据Host选择PathIndexer
	var pathIndexer *PathIndexer
	if pi, exists := idx.hostIndex[r.Host]; exists {
		pathIndexer = pi
	} else {
		pathIndexer = idx.defaultPathIndex
	}

	// 第二步：使用基数树查找路径前缀匹配的候选路由
	candidates := pathIndexer.findCandidates(r.URL.Path)

	// 第三步：在候选路由中进行AST求值，找到最高优先级的匹配
	for _, route := range candidates {
		if route.AST.Eval(r) {
			if candidateMatch == nil || route.Config.Priority > maxCandidatePriority {
				candidateMatch = route
				maxCandidatePriority = route.Config.Priority
			}
			// 找到第一个匹配的候选路由后，记录其优先级
			// 候选列表已按优先级排序，所以第一个匹配的就是最高优先级
			break
		}
	}

	// 第四步：遍历全量列表，检查是否有更高优先级的路由
	// 只检查优先级 > 候选最高优先级的路由
	for _, route := range idx.allRoutes {
		// 如果当前路由优先级不大于已找到的候选路由，后续也不会有更高的了
		if candidateMatch != nil && route.Config.Priority <= maxCandidatePriority {
			break
		}

		if route.AST.Eval(r) {
			return route
		}
	}

	// 返回候选路由中找到的匹配（如果有）
	return candidateMatch
}

// findCandidates 查找路径前缀匹配的候选路由
func (pi *PathIndexer) findCandidates(path string) []*RouteEntry {
	// 使用基数树的LongestPrefix查找
	_, value, found := pi.radixTree.Root().LongestPrefix([]byte(path))
	if !found {
		return nil
	}

	return *value
}

// analyzeAST 分析AST，提取可索引的Host和PathPrefix
func analyzeAST(node ast.Node) (hosts []string, pathPrefixes []string) {
	switch n := node.(type) {
	case *ast.FuncCall:
		if n.Name == "Host" && len(n.Args) == 1 {
			hosts = append(hosts, n.Args[0])
		} else if n.Name == "PathPrefix" && len(n.Args) == 1 {
			pathPrefixes = append(pathPrefixes, n.Args[0])
		} else if n.Name == "Path" && len(n.Args) == 1 {
			pathPrefixes = append(pathPrefixes, n.Args[0])
		}

	case *ast.BinaryOp:
		// AND: 取两边的交集（都要满足）
		if n.Op == "&&" {
			leftHosts, leftPaths := analyzeAST(n.Left)
			rightHosts, rightPaths := analyzeAST(n.Right)

			hosts = append(hosts, leftHosts...)
			hosts = append(hosts, rightHosts...)
			pathPrefixes = append(pathPrefixes, leftPaths...)
			pathPrefixes = append(pathPrefixes, rightPaths...)
		}

		// OR: 取两边的并集（满足任一）
		if n.Op == "||" {
			leftHosts, leftPaths := analyzeAST(n.Left)
			rightHosts, rightPaths := analyzeAST(n.Right)

			hosts = append(hosts, leftHosts...)
			hosts = append(hosts, rightHosts...)
			pathPrefixes = append(pathPrefixes, leftPaths...)
			pathPrefixes = append(pathPrefixes, rightPaths...)
		}

	case *ast.UnaryOp:
		// NOT: 无法提取索引信息
		return nil, nil
	}

	return hosts, pathPrefixes
}

// sortRoutesByPriority 按优先级对路由进行排序（优先级高的在前，相同优先级保持添加顺序）
func sortRoutesByPriority(routes []*RouteEntry) {
	sort.SliceStable(routes, func(i, j int) bool {
		// 优先级高的在前（降序）
		return routes[i].Config.Priority > routes[j].Config.Priority
	})
}