package trace

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strings"
)

func Write(dir string) {
	_ = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		if !strings.HasSuffix(path, ".go") {
			return nil
		}
		if strings.HasSuffix(path, "_test.go") {
			return nil
		}
		astFile(path)
		return nil
	})

	for s, m := range modifyFile {
		if len(m) == 0 {
			continue
		}
		file, _ := os.ReadFile(s)
		lines := bytes.Split(file, []byte("\n"))
		for i, line := range lines {
			if _, ok := m[i+1]; ok {
				lines[i] = append(line, []byte(wrapDeferFuncName)...)
			}
		}
		lines[0] = append(lines[0], []byte(wrapPackage)...)
		_ = os.WriteFile(s, bytes.Join(lines, []byte("\n")), 0644)
	}
}

var modifyFile = make(map[string]map[int]bool)

func astFile(path string) {
	file, _ := os.Open(path)
	fSet := token.NewFileSet()
	node, err := parser.ParseFile(fSet, path, file, parser.ParseComments)
	if err != nil {
		panic(err)
	}
	ast.Inspect(node, func(node ast.Node) bool {
		fmt.Println(reflect.TypeOf(node))
		switch nFile := node.(type) {
		case *ast.File:
			//for i, decl := range nFile.Imports {
			//	fmt.Println(i, reflect.TypeOf(decl))
			//}
			for _, decl := range nFile.Decls {
				if decl, ok := decl.(*ast.FuncDecl); ok && len(decl.Body.List) > 0 {
					//line0 := fSet.PositionFor(decl.Pos(), false).Line
					line1 := fSet.Position(decl.Body.List[0].Pos()).Line
					body0 := decl.Body.List[0]
					if stmt, ok := body0.(*ast.DeferStmt); ok {
						if ident, ok := stmt.Call.Fun.(*ast.Ident); ok {
							if ident.Name == "AstTrace" {
								continue
							}
						}
					}
					if _, ok := modifyFile[path]; !ok {
						modifyFile[path] = make(map[int]bool)
					}
					modifyFile[path][line1-1] = true
				}
			}
		case *ast.FuncDecl:
		default:
			log.Println(reflect.TypeOf(node).Name())
		}
		return false
	})
}