#!/usr/bin/env python3

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 addr not 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)  # unused
                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)
