lucky_bot/magic_admin/ams_ast/ast_router_root_enter.go

package ams_ast

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"os"
	"path/filepath"
)

// 根路由注册器

type RootRouterRegister struct {
	// demo : /Users/***/app_manage_system/go_server/router/enter.go
	TargetFile  string   `json:"targetFile"`  // 目标文件-固定-构建时使用绝对路径 go_server/router/enter.go
	ModuleNames []string `json:"moduleNames"` // 模块名
	StructName  string   `json:"structName"`  // 结构体名称 -- RouterGroup
	ValKey      string   `json:"valKey"`      // 默认值  priRouters
	BaseAst
}

// 服务根路由入口文件地址

func ServerRootRouterEnterFile() string {
	return serverRootRouterEnterFile
}

// 支持同时注册多个模块 moduleNames

func BuildRootRouterRegister(alias []string) *RootRouterRegister {
	return &RootRouterRegister{
		TargetFile:  serverRootRouterEnterFile,
		ModuleNames: alias,
		StructName:  "RouterGroup",
		ValKey:      "priRouters",
	}
}

func (r *RootRouterRegister) routerModulePkg(module string) string {
	return fmt.Sprintf("go_server/router/%s", module)
}

// 总路由注册

func (r *RootRouterRegister) Inspect() error {
	// 读取目标文件
	src, err := os.ReadFile(r.TargetFile)
	if err != nil {
		return err
	}
	// 创建 FileSet 用于跟踪文件位置信息
	fSet := token.NewFileSet()
	// 解析源代码为 AST
	f, err := parser.ParseFile(fSet, "", src, 0)
	if err != nil {
		return err
	}
	ast.Inspect(f, func(n ast.Node) bool {
		if decl, ok := n.(*ast.GenDecl); ok && decl.Tok == token.VAR {
			for _, spec := range decl.Specs {
				if valueSpec, ok := spec.(*ast.ValueSpec); ok {
					for i, name := range valueSpec.Names {
						if name.Name == r.ValKey && len(valueSpec.Values) > i {
							if compLit, ok := valueSpec.Values[i].(*ast.CompositeLit); ok {
								// 元素确定
								if !r.isGlobalContextInterface(compLit) {
									return false
								}
								for _, moduleName := range r.ModuleNames {
									// 添加新的元素 ModelName.RouterGroup{}
									newElement := &ast.CompositeLit{
										Type: &ast.SelectorExpr{
											X:   ast.NewIdent(moduleName),
											Sel: ast.NewIdent(r.StructName),
										},
									}
									// 检查是否已注入
									if !r.IsContainsModule(compLit.Elts, moduleName, r.StructName) {
										compLit.Elts = append(compLit.Elts, newElement)
									}
									// 确保导入了 模块对应的包
									normalizedPkg := filepath.ToSlash(moduleName)
									r.AddImport(f, r.routerModulePkg(normalizedPkg))
								}
							}
						}
					}
				}
			}
		}
		return true
	})
	var out []byte
	bf := bytes.NewBuffer(out)
	err = printer.Fprint(bf, fSet, f)
	if err != nil {
		return err
	}
	return os.WriteFile(r.TargetFile, bf.Bytes(), 0666)
}

// 可做进一步检查检查类型是否是 []global.ContextInterface

func (r *RootRouterRegister) isGlobalContextInterface(compLit *ast.CompositeLit) bool {
	if arrType, ok := compLit.Type.(*ast.ArrayType); ok {
		if selExpr, ok := arrType.Elt.(*ast.SelectorExpr); ok {
			if xIdent, ok := selExpr.X.(*ast.Ident); ok && xIdent.Name == "global" {
				if selExpr.Sel.Name == "ContextInterface" {
					// 确认是我们要找的类型
					return true
				}
			}
		}
	}
	return false
}

// 总路由回滚 -- 仅支持单一模块回滚移除

func (r *RootRouterRegister) RollbackRootRouter() error {
	// 读取目标文件
	src, err := os.ReadFile(r.TargetFile)
	if err != nil {
		return err
	}
	// 创建 FileSet 用于跟踪文件位置信息
	fSet := token.NewFileSet()
	// 解析源代码为 AST
	f, err := parser.ParseFile(fSet, "", src, 0)
	if err != nil {
		return err
	}
	ast.Inspect(f, func(n ast.Node) bool {
		if decl, ok := n.(*ast.GenDecl); ok && decl.Tok == token.VAR {
			for _, spec := range decl.Specs {
				if valueSpec, ok := spec.(*ast.ValueSpec); ok {
					for i, name := range valueSpec.Names {
						if name.Name == r.ValKey && len(valueSpec.Values) > i {
							if compLit, ok := valueSpec.Values[i].(*ast.CompositeLit); ok {
								// 元素确定
								if !r.isGlobalContextInterface(compLit) {
									return false
								}
								newEls := make([]ast.Expr, 0, len(compLit.Elts))
								removeImport := false
								for _, elt := range compLit.Elts {
									if !r.IsTargetModule(elt, r.ModuleNames[0], r.StructName) {
										newEls = append(newEls, elt)
									} else {
										removeImport = true
									}
									if removeImport {
										// 确保删除模块对应的包
										r.RemoveImport(f, r.routerModulePkg(r.ModuleNames[0]))
									}
								}
								compLit.Elts = newEls

							}
						}
					}
				}
			}
		}
		return true
	})
	var out []byte
	bf := bytes.NewBuffer(out)
	err = printer.Fprint(bf, fSet, f)
	if err != nil {
		return err
	}
	err = os.Remove(r.TargetFile)
	if err != nil {
		return err
	}
	return os.WriteFile(r.TargetFile, bf.Bytes(), 0666)
}