package main

import (
	"fmt"
	"os"
	"io/ioutil"
	"strings"
	"regexp"
)

var functionPrefix = "Complexity"

var testingFunc = `
func main(){
	nVals := []int{1, 10, 100, 1000, 10000, 100000}
	for _, n := range nVals {
		dataSet := getDataSet(n)
		start := time.Now().UnixNano()
		<complexityFunc>(dataSet)
		fmt.Printf("\n(%d, %d)\n", n, time.Now().UnixNano() - start)
	}
}`

func main() {
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		fmt.Println(err)
	}
	code := string(data)

	// Split the code at each function declaration and exclude the first element
	funcBodys := strings.Split(code, "\nfunc ")[1:]

	// Parse the funtionBodys to retrieve a list of all functions in code
	funcNames := getFuncs(funcBodys)

	// Rename original main function so that it doesn't conflict with our new main()
	code = renameMainFunc(code, "main2")

	// Add in our new main method and swap out the placeholder for the function call
	code = code + strings.Replace(testingFunc, "<complexityFunc>", getTestFunc(funcNames), -1)

	// Print to stdout
	fmt.Println(code)
}

// Parses slice of function bodys and returns a slice with all function names
func getFuncs(funcBodys []string) []string {
	funcNames := make([]string, len(funcBodys))
	for i := 0; i < len(funcBodys); i++ {
		funcName := strings.Split(funcBodys[i], "(")[0]
		funcNames[i] = funcName
	}
	return funcNames
}

// Selects the first function in the list of functions that starts with a predefined function prefix
// Returns an empty string if fitting function is not found
func getTestFunc(funcNames []string) string {
	for _, funcName := range funcNames {
		if funcName[0:len(functionPrefix)] == functionPrefix {
			return funcName
		}
	}
	return ""
}

// Uses regex to find and replace the main method declaration with newFuncName
func renameMainFunc(code, newFuncName string) string {
	funcRegex, _ := regexp.Compile("func main ?\\( ?\\) ?{")
	return funcRegex.ReplaceAllString(code, fmt.Sprintf("func %s() {", newFuncName))
}