//go:generate go run main.go -input ../../spec

package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io/fs"
	"log"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
)

var (
	// Helpers to replace the `String()` method of enum.
	stringMethodName       = "String"
	enumMethodName         = "Enum"
	enumOptionsImportAlias = "datatypev1"
	replaceCode            = fmt.Sprintf(`
package main
func this_function_body_will_replace_the_string_method_body_of_enum(x any) string {
    valueDesc := x.Descriptor().Values().ByNumber(x.Number())
    if proto.HasExtension(valueDesc.Options(), %s.E_Name) {
        ext := proto.GetExtension(valueDesc.Options(), %s.E_Name)
        s, ok := ext.(string)
        if ok && s != "" {
            return s
        }
    }
    return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x))
}
`, enumOptionsImportAlias, enumOptionsImportAlias)

	// Helpers to modify imports after replacing the `String()` method of enum.
	protoImportPath       = `"google.golang.org/protobuf/proto"`
	enumOptionsImportPath = `"github.com/tidbcloud/tidb-management-service/api/spec/datatype/v1"`
)

func main() {
	var inputDir string
	flag.StringVar(&inputDir, "input", "", "input dir")
	flag.Parse()

	rewritePbGoFilesFunc := func(path string, info fs.FileInfo, err error) error {
		if info.IsDir() {
			return nil
		}
		if !strings.Contains(path, "pb.go") {
			return nil
		}

		return rewrite(path)
	}

	if err := filepath.Walk(inputDir, rewritePbGoFilesFunc); err != nil {
		log.Fatal("rewrite go file failed", err)
	}
}

func rewrite(inputPath string) error {
	// Parse the input file to an AST.
	fset := token.NewFileSet()
	insertionFile, err := parser.ParseFile(fset, "insertion.go", replaceCode, parser.ParseComments)
	if err != nil {
		return err
	}
	targetFile, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments)
	if err != nil {
		return err
	}

	if !hasDatatypeImport(targetFile) {
		return nil
	}

	// Find all enum types in the file.
	enumTypesInFile := findEnumTypesFunc(targetFile)

	// Rewrite function replace enum `String()` method to `replaceCode()` body.
	needToModifyImport := rewriteFunc(targetFile, insertionFile, enumTypesInFile)

	// Add import if needed.
	if needToModifyImport {
		rewriteImport(targetFile)
	}

	// Convert the rewritten AST to string.
	var modifiedCode strings.Builder
	if err := printer.Fprint(&modifiedCode, fset, targetFile); err != nil {
		return err
	}

	// Reformat code.
	formatted, err := format.Source([]byte(modifiedCode.String()))
	if err != nil {
		return err
	}

	// Write back to the input file.
	if err := os.WriteFile(inputPath, formatted, 0o600); err != nil {
		return err
	}

	return nil
}

func rewriteFunc(fast, insertionStmt *ast.File, enumTypesInFile map[string]struct{}) bool {
	var needToImport bool
	astutil.Apply(fast, func(cursor *astutil.Cursor) bool {
		n := cursor.Node()
		funcDecl, ok := n.(*ast.FuncDecl)
		if !ok || funcDecl.Recv == nil {
			return true
		}

		// Check if the receiver type is an enum type with a "replace" tag.
		isFuncReceiverReplaceEnum := false
		for _, field := range funcDecl.Recv.List {
			identType, ok := field.Type.(*ast.Ident)
			if !ok {
				continue
			}
			_, ok = enumTypesInFile[identType.Name]
			if ok {
				isFuncReceiverReplaceEnum = true
				break
			}
		}

		// Replace the `String()` method with the `replaceCode()` body.
		if isFuncReceiverReplaceEnum && funcDecl.Name.Name == stringMethodName {
			needToImport = true
			funcDecl.Body = insertionStmt.Decls[0].(*ast.FuncDecl).Body
			fmt.Printf("replace %s to use name tag \n", fast.Name.String())
		}

		return true
	}, nil)
	return needToImport
}

func findEnumTypesFunc(fast *ast.File) map[string]struct{} {
	enumTypesInFile := make(map[string]struct{})
	astutil.Apply(fast, func(cursor *astutil.Cursor) bool {
		n := cursor.Node()
		funcDecl, ok := n.(*ast.FuncDecl)
		if !ok || funcDecl.Recv == nil {
			return true
		}

		// Enum type is defined in a method with the name "Enum".
		if funcDecl.Name.Name == enumMethodName {
			for _, field := range funcDecl.Recv.List {
				identType, ok := field.Type.(*ast.Ident)
				if !ok {
					continue
				}
				enumTypesInFile[identType.Name] = struct{}{}
			}
		}

		return true
	}, nil)
	return enumTypesInFile
}

func hasDatatypeImport(f *ast.File) bool {
	var importDecl *ast.GenDecl
	for _, decl := range f.Decls {
		if d, ok := decl.(*ast.GenDecl); ok && d.Tok == token.IMPORT {
			importDecl = d
		}
	}
	for _, spec := range importDecl.Specs {
		if importSpec, ok := spec.(*ast.ImportSpec); ok {
			if importSpec.Path.Value == enumOptionsImportPath {
				return true
			}
		}
	}
	return false
}

func rewriteImport(f *ast.File) {
	// Get the import declaration.
	var importDecl *ast.GenDecl
	for _, decl := range f.Decls {
		if d, ok := decl.(*ast.GenDecl); ok && d.Tok == token.IMPORT {
			importDecl = d
		}
	}

	hasProtoImport := false

	for _, spec := range importDecl.Specs {
		if importSpec, ok := spec.(*ast.ImportSpec); ok {
			if importSpec.Path.Value == protoImportPath {
				hasProtoImport = true
			}
			// Rewrite the enum option import alias.
			if importSpec.Path.Value == enumOptionsImportPath {
				importSpec.Name = &ast.Ident{Name: enumOptionsImportAlias}
			}
		}
	}

	if !hasProtoImport {
		// Add the proto import if it is not already present.
		importDecl.Specs = append(importDecl.Specs, &ast.ImportSpec{
			Path: &ast.BasicLit{
				Kind:  token.STRING,
				Value: protoImportPath,
			},
		})
	}
}