Add logging on function by AST

_ = require("lodash")
esprima = require('esprima')
escodegen = require('escodegen')

handleFunctionBody = (body, newStatements) ->
  body.body = newStatements.concat(body.body)

insertStatementsIntoFunction = (s, newStatements) ->
  if s.type == 'FunctionDeclaration' && s.body && s.body.type == 'BlockStatement'
    handleFunctionBody(s.body, newStatements)
  if _.isArray(s.body)
    _.each(s.body, (b) =>
      insertStatementsIntoFunction(b, newStatements)
    )
  else if s.body
    insertStatementsIntoFunction(s.body, newStatements)

addLogger = (orig) ->
  logger = esprima.parse('console.log("inserted console log", arguments);').body
  tree = esprima.parse(orig)
#  console.log(JSON.stringify(tree))
  insertStatementsIntoFunction(tree, logger)
  escodegen.generate(tree)


origCode = '' +
  'function test(arg) {' +
    'console.log("test function");' +
  '};' +
  'test("arg1");'

newCode = addLogger(origCode)
console.log(origCode)
console.log(newCode)
eval(newCode)