package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"sort"

	"github.com/MakeNowJust/heredoc/v2"
	"github.com/cloudspannerecosystem/memefish/tools/util/astcatalog"
)

var (
	usage = heredoc.Doc(`
		Usage of tools/gen-ast-walk.go

		A generator of ast/walk_internal.go.

		Example:

		  $ go run ./tools/gen-ast-walk/main.go -astfile ast/ast.go -constfile ast/ast_const.go
		        Print the generated ast/walk_internal.go to stdout.

		Flags:
	`)
	prologue = heredoc.Doc(`
		// Code generated by tools/gen-ast-walk; DO NOT EDIT.

		package ast

		func walkInternal(node Node, v Visitor, stack []*stackItem) []*stackItem {
			switch n := node.(type) {
	`)
)

var (
	astfile   = flag.String("astfile", "ast/ast.go", "path to ast/ast.go")
	constfile = flag.String("constfile", "ast/ast_const.go", "path to ast/ast_const.go")
	outfile   = flag.String("outfile", "", "output filename (if it is not specified, the result is printed to stdout.)")
)

func main() {
	flag.Usage = func() {
		fmt.Print(usage)
		flag.PrintDefaults()
	}

	flag.Parse()

	catalog, err := astcatalog.Load(*astfile, *constfile)
	if err != nil {
		log.Fatal(err)
	}

	structs := make([]*astcatalog.NodeStructDef, 0, len(catalog.Structs))
	for _, structDef := range catalog.Structs {
		structs = append(structs, structDef)
	}
	sort.Slice(structs, func(i, j int) bool {
		return structs[i].SourcePos < structs[j].SourcePos
	})

	var buffer bytes.Buffer
	buffer.WriteString(prologue)

	for i, structDef := range structs {
		fmt.Fprintf(&buffer, "\tcase *%s:\n", structDef.Name)

		output := false

		for i := len(structDef.Fields) - 1; i >= 0; i-- {
			field := structDef.Fields[i]

			if _, ok := extractNodeStructPointer(field.Type); ok {
				fmt.Fprintf(&buffer, "\t\tstack = append(stack, &stackItem{node: wrapNode(n.%s), visitor: v.Field(%q)})\n", field.Name, field.Name)
				output = true
				continue
			}

			if _, ok := extractNodeStructSlice(field.Type); ok {
				fmt.Fprintf(&buffer, "\t\tstack = append(stack, &stackItem{nodes: wrapNodes(n.%s), visitor: v.Field(%q)})\n", field.Name, field.Name)
				output = true
				continue
			}

			if _, ok := extractNodeInterface(field.Type); ok {
				fmt.Fprintf(&buffer, "\t\tstack = append(stack, &stackItem{node: wrapNode(n.%s), visitor: v.Field(%q)})\n", field.Name, field.Name)
				output = true
				continue
			}

			if _, ok := extractNodeInterfaceSlice(field.Type); ok {
				fmt.Fprintf(&buffer, "\t\tstack = append(stack, &stackItem{nodes: wrapNodes(n.%s), visitor: v.Field(%q)})\n", field.Name, field.Name)
				output = true
				continue
			}
		}

		if !output {
			fmt.Fprintln(&buffer, "\t\t// nothing to do")
		}

		if i < len(structs)-1 {
			fmt.Fprintln(&buffer)
		}
	}

	fmt.Fprintln(&buffer, "\t}")
	fmt.Fprintln(&buffer, "\treturn stack")
	fmt.Fprintln(&buffer, "}")

	if *outfile == "" {
		fmt.Print(buffer.String())
		return
	}

	err = os.WriteFile(*outfile, buffer.Bytes(), 0666)
	if err != nil {
		log.Fatal(err)
	}
}

func extractNodeStructPointer(ty astcatalog.Type) (astcatalog.NodeStructType, bool) {
	if ty, ok := ty.(*astcatalog.PointerType); ok {
		if ty, ok := ty.Type.(astcatalog.NodeStructType); ok {
			return ty, true
		}
	}

	return "", false
}

func extractNodeStructSlice(ty astcatalog.Type) (astcatalog.NodeStructType, bool) {
	if ty, ok := ty.(*astcatalog.SliceType); ok {
		return extractNodeStructPointer(ty.Type)
	}

	return "", false
}

func extractNodeInterface(ty astcatalog.Type) (astcatalog.NodeInterfaceType, bool) {
	if ty, ok := ty.(astcatalog.NodeInterfaceType); ok {
		return ty, ok
	}

	return "", false
}

func extractNodeInterfaceSlice(ty astcatalog.Type) (astcatalog.NodeInterfaceType, bool) {
	if ty, ok := ty.(*astcatalog.SliceType); ok {
		return extractNodeInterface(ty.Type)
	}

	return "", false
}
