gencode/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/dave/dst"
	"github.com/dave/dst/decorator"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

const (
	configFile = "project_config.json" // 项目配置文件
)

var (
	projectMutexes = make(map[string]*sync.Mutex)
)

type ProjectConfig struct {
	StartCode int `json:"startCode"`
}

var newCodeGenComment = "// Code generated by gen_code. DO NOT EDIT."

func generateBusinessCodes(projectName string) error {
	// 解析 errmsg 包中的文件
	fset := token.NewFileSet()
	filename := "../errmsg/errmsg.go"
	file, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
	if err != nil {
		return err
	}

	// 删除原始的 //go:generate 注释
	file.Comments = removeGoGenerateComment(file.Comments)

	// 添加新的生成注释
	addGeneratedComment(file)

	projectMutex, err := getProjectMutex(projectName)
	if err != nil {
		return err
	}
	projectMutex.Lock()
	defer projectMutex.Unlock()

	// 加载项目配置
	projectConfig, err := loadProjectConfig(projectName)
	if err != nil {
		return err
	}

	for _, decl := range file.Decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok || genDecl.Tok != token.CONST {
			continue
		}

		for _, spec := range genDecl.Specs {
			valueSpec, ok := spec.(*ast.ValueSpec)
			if !ok {
				continue
			}

			for i, name := range valueSpec.Names {
				if strings.HasPrefix(name.Name, "TypeMsg") {
					trimmedName := strings.TrimPrefix(name.Name, "TypeMsg")
					constName := fmt.Sprintf("TypeCode%s", trimmedName)
					constOriValue := valueSpec.Values[i].(*ast.BasicLit).Value
					// 更新常量名
					valueSpec.Names[i] = ast.NewIdent(constName)

					// 生成常量
					valueSpec.Values[i] = &ast.BasicLit{
						Kind:  token.INT,
						Value: fmt.Sprintf("%d", projectConfig.StartCode),
					}

					// 生成 Get 函数
					getFuncName := "Get" + trimmedName
					getFunc := &ast.FuncDecl{
						Name: ast.NewIdent(getFuncName),
						Type: &ast.FuncType{
							Results: &ast.FieldList{
								List: []*ast.Field{
									{
										Type: ast.NewIdent("Resp"),
									},
								},
							},
						},
						Body: &ast.BlockStmt{
							List: []ast.Stmt{
								&ast.ReturnStmt{
									Results: []ast.Expr{
										&ast.CompositeLit{
											Type: ast.NewIdent("Resp"),
											Elts: []ast.Expr{
												&ast.KeyValueExpr{
													Key:   ast.NewIdent("Code"),
													Value: ast.NewIdent(constName),
												},
												&ast.KeyValueExpr{
													Key:   ast.NewIdent("Msg"),
													Value: ast.NewIdent(name.Name),
												},
											},
										},
									},
								},
							},
						},
						// 添加 Doc 注释
						Doc: &ast.CommentGroup{
							List: []*ast.Comment{
								{
									Text: fmt.Sprintf("// 提示内容：%s\n// 状态码：%d\n", constOriValue, projectConfig.StartCode),
								},
							},
						},
					}
					file.Decls = append(file.Decls, getFunc)
					file.Comments = append(file.Comments, getFunc.Doc)
					projectConfig.StartCode++
				}
			}
		}
	}

	//// 保存已分配的错误码
	//err = saveLastErrorCode(projectName, lastCode)
	//if err != nil {
	//	return err
	//}
	//
	// 格式化生成的代码
	var generatedCode strings.Builder
	err = format.Node(&generatedCode, fset, file)
	if err != nil {
		return err
	}

	f, err := decorator.Parse(generatedCode.String())
	if err != nil {
		panic(err)
	}

	start := 1
	for _, decl := range f.Decls {
		genDecl, ok := decl.(*dst.FuncDecl)
		if !ok {
			continue
		}

		genDecl.Decorations().Start = f.Decorations().Start[start : start+2]
		start += 2
	}

	f.Decorations().Start = f.Decorations().Start[:1]

	if err := decorator.Print(f); err != nil {
		panic(err)
	}

	// 将生成的代码写入文件（与 errmsg 包同级目录）
	// 使用绝对路径来指定代码文件的位置
	outputFilePath := filepath.Join("", "generated_codes.go")
	//err = os.WriteFile(outputFilePath, []byte(generatedCode.String()), 0644)
	//if err != nil {
	//	return err
	//}
	printFile, _ := os.Create(outputFilePath)

	err = decorator.Fprint(printFile, f)
	if err != nil {
		panic(err)
	}

	return nil
}

func getProjectMutex(projectName string) (*sync.Mutex, error) {
	mutex, exists := projectMutexes[projectName]
	if exists {
		return mutex, nil
	}

	mutex = &sync.Mutex{}
	projectMutexes[projectName] = mutex
	return mutex, nil
}

func loadProjectConfig(projectName string) (*ProjectConfig, error) {
	configFilePath := filepath.Join("..", "", configFile)
	configFileData, err := os.ReadFile(configFilePath)
	if err != nil {
		return nil, err
	}

	configs := make(map[string]ProjectConfig)
	err = json.Unmarshal(configFileData, &configs)
	if err != nil {
		return nil, err
	}

	config, exists := configs[projectName]
	if !exists {
		return nil, fmt.Errorf("project config not found for %s", projectName)
	}

	return &config, nil
}

func loadLastErrorCode(projectName string) (int, error) {
	filename := fmt.Sprintf("%s_last_code.txt", projectName)
	data, err := os.ReadFile(filename)
	if err != nil {
		if os.IsNotExist(err) {
			return 0, nil // 初始错误码
		}
		return 0, err
	}

	var lastCode int
	_, err = fmt.Sscanf(string(data), "%d", &lastCode)
	if err != nil {
		return 0, err
	}

	return lastCode, nil
}

func saveLastErrorCode(projectName string, code int) error {
	filename := fmt.Sprintf("%s_last_code.txt", projectName)
	return os.WriteFile(filename, []byte(fmt.Sprintf("%d", code)), 0644)
}

// 删除原始的 //go:generate 注释
func removeGoGenerateComment(comments []*ast.CommentGroup) []*ast.CommentGroup {
	var newComments []*ast.CommentGroup
	for _, group := range comments {
		var newGroup ast.CommentGroup
		for _, comment := range group.List {
			if !strings.HasPrefix(comment.Text, "//go:generate") {
				newGroup.List = append(newGroup.List, comment)
			}
		}
		if len(newGroup.List) > 0 {
			newComments = append(newComments, &newGroup)
		}
	}
	return newComments
}

// 添加新的生成注释
func addGeneratedComment(file *ast.File) {
	newComment := &ast.Comment{
		Text:  newCodeGenComment,
		Slash: 0, // 设置为 0，即在第一行添加注释
	}

	if len(file.Comments) > 0 {
		file.Comments[0].List = append(file.Comments[0].List, newComment)
	} else {
		file.Comments = append(file.Comments, &ast.CommentGroup{
			List: []*ast.Comment{newComment},
		})
	}
}

func Run(ctx context.Context, projectName string) error {
	err := generateBusinessCodes(projectName)
	if err != nil {
		return err
	}

	fmt.Println("Business codes generated successfully")
	return nil
}

func main() {
	projectName := "myproject" // 你的项目名
	err := Run(context.Background(), projectName)
	if err != nil {
		fmt.Println(err)
	}
}