#!/usr/bin/env python2

from __future__ import print_function

from collections import Counter
from operator import itemgetter
from optparse import OptionParser
from re import match
from subprocess import PIPE, CalledProcessError, Popen
from sys import stderr, stdin, stdout

addr_map = {}


def check_output(args):
    proc = Popen(args, stdout=PIPE)
    out, _ = proc.communicate()
    res = proc.poll()

    if res != 0:
        raise CalledProcessError(res, args[0])

    return out


def map_address(addr, asd_path):
    if not addr in addr_map:
        try:
            resolv = check_output(["addr2line", "-e", asd_path, addr])

        except OSError:
            print("The addr2line command is not available. You might have to install the binutils package.", file=stderr)
            raise SystemExit(1)

        except CalledProcessError:
            print("The addr2line command failed.", file=stderr)
            raise SystemExit(1)

        addr_map[addr] = resolv.rstrip()

    return addr_map[addr]


def parse_mem(in_path):
    time = None
    aggr = Counter()

    try:
        f = stdin if in_path == '-' else open(in_path, "r")
        for line in f:
            line = line.rstrip()
            m = match("^-+ ([^-]+) -+$", line)

            if m:
                new_time = m.group(1)
                print("Processing memory dump of {0}".format(
                    new_time), file=stderr)

                if time:
                    print("More than one memory dump in {0}".format(
                        in_path), file=stderr)
                    raise SystemExit(1)

                time = new_time
                continue

            m = match(
                "^0x([0-9a-f]+) +([0-9]+) +0x([0-9a-f]+) +0x([0-9a-f]+)$", line)

            if m:
                (addr_s, tid_s, size_hi_s, size_lo_s) = m.group(1, 2, 3, 4)

                addr = int(addr_s, 16)
                tid = int(tid_s, 16)
                size_hi = int(size_hi_s, 16)
                size_lo = int(size_lo_s, 16)

                size = size_hi * 2 ** 64 + size_lo

                if (size >= 2 ** 127):
                    size -= 2 ** 128

                aggr[addr] += size
                continue

            print("Invalid line in {0}: {1}".format(
                in_path, line), file=stderr)
            raise SystemExit(1)

    except IOError:
        print("Cannot open file {0}".format(in_path), file=stderr)
        raise SystemExit(1)

    return aggr


parser = OptionParser()
parser.add_option("-i", "--in", dest="in_path", default="-", metavar="FILE",
                  help="read memory accounting info from FILE; defaults to \"-\", which means stdin")
parser.add_option("-o", "--out", dest="out_path", default="-", metavar="FILE",
                  help="write aggregated memory accounting info to FILE; defaults to \"-\", which means stdout")
parser.add_option("-a", "--asd", dest="asd_path", default="/usr/bin/asd", metavar="ASD",
                  help="read line number information from Aerospike server executable ASD; defaults to /usr/bin/asd")
parser.add_option("-d", "--disable-site-eval", dest="disable_site", action="store_true",
                  help="disable evaluating line number information from the Aerospike server executable")

(opts, args) = parser.parse_args()

aggr = parse_mem(opts.in_path)

try:
    f = stdout if opts.out_path == '-' else open(opts.out_path, "w")

    for addr, size in sorted(aggr.items(), key=itemgetter(1)):
        if size > 0:
            addr_s = "0x{0:x}".format(addr)

            if (opts.disable_site):
                print("{}|{}".format(addr_s, aggr[addr]), file=f)
            else:
                site = map_address(addr_s, opts.asd_path)
                print("{}|{}|{}".format(site, addr_s, aggr[addr]), file=f)
except IOError:
    print("Cannot open file {}".format(opts.out_path), file=stderr)
    raise SystemExit(1)

raise SystemExit(0)
