var esprima = require('esprima'),
    escodegen = require('escodegen');


// Executes visitor on the object and its children (recursively).
var traverse = function traverse(object, visitor) {
    var key, child;

    visitor.call(null, object);
    for (key in object) {
        if (object.hasOwnProperty(key)) {
            child = object[key];
            if (typeof child === 'object' && child !== null) {
                traverse(child, visitor);
            }
        }
    }
};


var powerAssertVariableDeclaration = {
    "type": "VariableDeclaration",
    "declarations": [
        {
            "type": "VariableDeclarator",
            "id": {
                "type": "Identifier",
                "name": "passert"
            },
            "init": {
                "type": "CallExpression",
                "callee": {
                    "type": "Identifier",
                    "name": "require"
                },
                "arguments": [
                    {
                        "type": "Literal",
                        "value": "../lib/power-assert"
                    }
                ]
            }
        }
    ],
    "kind": "var"
};


var isAssertOkCall = function (callee) {
    if (callee.type !== 'MemberExpression' || callee.computed !== false) {
        return false;
    }
    var ob = callee.object,
        prop = callee.property;
    return ((ob.type === 'Identifier' && ob.name === 'assert') && (prop.type === 'Identifier' && prop.name === 'ok'));
};


var isAssertCall = function (callee) {
    return (callee.type === 'Identifier' && callee.name === 'assert');
};


// enclose  in __ident__
var capture = function (target, ident) {
    return {
        "type": "CallExpression",
        "callee": {
            "type": "MemberExpression",
            "computed": false,
            "object": {
                "type": "Identifier",
                "name": "passert"
            },
            "property": {
                "type": "Identifier",
                "name": "__ident__"
            }
        },
        "arguments": [
            {
                "type": "Literal",
                "value": ident.name
            },
            target,
            {
                "type": "Literal",
                "value": ident.loc.start.column
            },
            {
                "type": "Literal",
                "value": ident.loc.end.column
            }
        ]
    };
};


var captureIdent = function (ident) {
    return capture(ident, ident);
};


var captureMember = function (member) {
    member.object = captureIdent(member.object);
    return capture(member, member.property);
};


// enclose assert expression in __expr__
var captureFirstArgument = function (firstArgument, line) {
    return {
        "type": "CallExpression",
        "callee": {
            "type": "MemberExpression",
            "computed": false,
            "object": {
                "type": "Identifier",
                "name": "passert"
            },
            "property": {
                "type": "Identifier",
                "name": "__expr__"
            }
        },
        "arguments": [
            firstArgument,
            {
                "type": "Literal",
                "value": line
            }
        ]
    };
};


var instrument = function (expression, line) {
    var firstArgument = expression.arguments[0];
    if (firstArgument.type === 'Identifier') {
        firstArgument = captureIdent(firstArgument);
    } else if (firstArgument.type === 'BinaryExpression') {
        var left = firstArgument.left,
            right = firstArgument.right;
        if (left.type === 'Identifier') {
            firstArgument.left = captureIdent(left);
        } else if (left.type === 'MemberExpression') {
            // ex: ary.length
            firstArgument.left = captureMember(left);
        }
        if (right.type === 'Identifier') {
            firstArgument.right = captureIdent(right);
        } else if (right.type === 'MemberExpression') {
            // ex: ary.length
            firstArgument.right = captureMember(right);
        }
    } else {
        return;
    }
    expression.arguments[0] = captureFirstArgument(firstArgument, line);
};


var modifyTree = function (syntax, lines) {
    traverse(syntax, function (node) {
        if (typeof node.type === 'undefined' || node.type !== 'ExpressionStatement') {
            return;
        }
        var expression = node.expression;
        if (typeof expression.type === 'undefined' || expression.type !== 'CallExpression') {
            return;
        }
        var callee = expression.callee;
        if (!(isAssertCall(callee) || isAssertOkCall(callee))) {
            return;
        }
        var line = lines[node.loc.start.line - 1];
        instrument(expression, line);
    });
};


var generator = function generator(source, options) {
    options = options || {};
    var lines = source.split("\n"),
        tree = esprima.parse(source, {tolerant: true, loc: true});
    modifyTree(tree, lines, options);

    if (options['module'] === 'CommonJS') {
        // insert "var passert = require('power-assert');"
        tree.body.unshift(powerAssertVariableDeclaration);
    }

    console.log(escodegen.generate(tree));
};


generator.modifyTree = modifyTree;
generator.instrument = instrument;


/*global exports:false*/
if (typeof exports !== 'undefined') {
    module.exports = generator;
}

@@ -3298,26 +3298,22 @@
 ;%0A%0A%0Avar 
-instrument
+capcap
  = funct
@@ -3321,107 +3321,49 @@
 on (
-expression, line) %7B%0A    var firstArgument = expression.arguments%5B0%5D;%0A    if (firstArgument.type ===
+node) %7B%0A    switch (node.type) %7B%0A    case
  'Id
@@ -3363,35 +3363,33 @@
 ase 'Identifier'
-) %7B
+:
 %0A        firstAr
@@ -3385,329 +3385,221 @@
     
-firstArgument = captureIdent(firstArgument);%0A    %7D else if (firstArgument.type === 'BinaryExpression') %7B%0A        var left = firstArgument.left,%0A            right = firstArgument.right;%0A        if (left.type === 'Identifier') %7B%0A            firstArgument.left = captureIdent(left);%0A        %7D else if (left.type === 'MemberE
+return captureIdent(node);%0A        break;%0A    case 'MemberExpression': // ex: ary.length%0A        return captureMember(node);%0A        break;%0A    default:%0A        return node;%0A    %7D%0A%7D;%0A%0A%0Avar instrument = function (e
 xpre
@@ -3599,33 +3599,38 @@
  (expression
-'
+, line
 ) %7B%0A    
         // e
@@ -3621,124 +3621,95 @@
     
-        // ex: ary.length%0A            firstArgument.left = captureMember(left);%0A        %7D%0A        if (right.type ===
+var firstArgument = expression.arguments%5B0%5D;%0A    switch (firstArgument.type) %7B%0A    case
  'Id
@@ -3717,24 +3717,18 @@
 ntifier'
-) %7B%0A    
+:%0A
         
@@ -3744,79 +3744,65 @@
 ment
-.right = captureIdent(right);%0A        %7D else if (right.type === 'Member
+ = capcap(firstArgument);%0A        break;%0A    case 'Binary
 Expr
@@ -3800,35 +3800,33 @@
 inaryExpression'
-) %7B
+:
 %0A            // 
@@ -3822,29 +3822,56 @@
     
-    // ex: ary.length
+firstArgument.left = capcap(firstArgument.left);
 %0A   
@@ -3867,36 +3867,32 @@
 .left);%0A        
-    
 firstArgument.ri
@@ -3904,49 +3904,61 @@
  cap
-tureMember(right);%0A        %7D%0A    %7D else %7B
+cap(firstArgument.right);%0A        break;%0A    default:
 %0A   

switch it