package main

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"
)

func main() {
	if err := process(os.Args[1]); err != nil {
		log.Fatal(err)
	}
}

func process(path string) error {
	pkg, err := build.Import(path, "", 0)
	if err != nil {
		return err
	}

	fset := token.NewFileSet()
	var files []*ast.File
	for _, file := range pkg.GoFiles {
		f, err := parser.ParseFile(fset, filepath.Join(pkg.Dir, file), nil, parser.ParseComments)
		if err != nil {
			return err
		}
		files = append(files, f)
	}

	for _, file := range files {
		var changed bool
		ast.Inspect(file, func(n ast.Node) bool {
			bs, ok := n.(*ast.BlockStmt)
			if !ok {
				return true
			}

			for i := range bs.List {
				if i < 2 {
					continue
				}
				ds, ok := bs.List[i-2].(*ast.DeclStmt)
				if !ok {
					continue
				}
				as, ok := bs.List[i-1].(*ast.AssignStmt)
				if !ok {
					continue
				}
				ret, ok := bs.List[i].(*ast.ReturnStmt)
				if !ok {
					continue
				}
				if as.Tok != token.ADD_ASSIGN {
					continue
				}
				gd, ok := ds.Decl.(*ast.GenDecl)
				if !ok {
					continue
				}
				if len(gd.Specs) != 1 {
					continue
				}
				vs, ok := gd.Specs[0].(*ast.ValueSpec)
				if !ok {
					continue
				}
				if len(vs.Names) != 1 {
					continue
				}
				if len(as.Lhs) != 1 {
					continue
				}
				id, ok := as.Lhs[0].(*ast.Ident)
				if !ok {
					continue
				}
				if vs.Names[0].Obj != id.Obj {
					continue
				}
				if len(ret.Results) != 1 {
					fmt.Println(len(ret.Results))
					continue
				}
				retid, ok := ret.Results[0].(*ast.Ident)
				if !ok {
					continue
				}
				if retid.Obj != id.Obj {
					fmt.Println("diffobj")
					continue
				}
				copy(bs.List[i-2:], bs.List[i:])
				bs.List = bs.List[:len(bs.List)-2]
				ret.Results = as.Rhs
				changed = true
				f := fset.File(bs.Pos())
				f.MergeLine(fset.Position(as.Pos()).Line)
				f.MergeLine(fset.Position(bs.Pos()).Line)
			}
			return true
		})

		if changed {
			fname := fset.Position(file.Pos()).Filename
			c, err := os.Create(fname)
			if err != nil {
				log.Fatal(err)
			}

			err = format.Node(c, fset, file)
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	return nil
}