import os, sys, re
import py
import time

def parse_smaps(lines, attr_ready_callback):
    iterlines = iter(lines)
    headline = iterlines.next()
    try:
        while 1:
            attrs = []
            while 1:
                sizeline = iterlines.next()
                parts = sizeline.split()
                if len(parts) != 3:
                    nextheadline = sizeline
                    break 
                value = int(parts[1])
                attrs.append(value)
            attr_ready_callback(attrs, headline)
            headline = nextheadline
    except StopIteration:
        attr_ready_callback(attrs, headline)

def private_sum(lines):
    private_sum = [0]
    
    def callback(attrs, headline):
        if attrs:
            private_sum[0] += attrs[4] + attrs[5]

    parse_smaps(lines, callback)
    return private_sum[0]

class SmapsRecorder:
    SEPSNAPSHOT = "-"*80

    def __init__(self, pid, stream, filter_system_libs=False, do_flush=True):
        self.stream = stream 
        self.pid = pid
        self.smapspath = py.path.local("/proc/%d/smaps" %(pid,))
        self.filter_system_libs = filter_system_libs
        self.do_flush = do_flush
        assert self.smapspath.check()


    def append_attrs(self, attrs, headline):
        if self.filter_system_libs:
            if headline.split()[-1].strip().startswith("/lib"):
                return
        attrs = attrs[:6] # not interested in "Referenced"
        sizes = ",".join(map(str, attrs))
        assert sizes
        self.stream.write("%-24s %s" %(sizes, headline))

    def _read_lines_from_file(self):
        f = self.smapspath.open()
        lines = f.readlines()
        f.close()
        if not lines:
            # XXX
            raise py.error.ENOENT
        return lines

    def snapshot(self):
        lines = self._read_lines_from_file()
        parse_smaps(lines, self.append_attrs)
        self.stream.write(self.SEPSNAPSHOT + "\n")
        if self.do_flush:
            self.stream.flush()

    def simple_snapshot(self):
        """ Simple version of the above, that stores only private memory
        """
        lines = self._read_lines_from_file()        
        self.stream.write(repr(time.time()) + " " + str(private_sum(lines)) + '\n')

#
# objects gained from reading of the logfile 
#

class Mapping: 
    _attrnames = ("size rss shared_clean shared_dirty "
                  "private_clean private_dirty".split())
    # example '402c1000-402c2000 rw-p 00011000 1f:04 1110 /lib/libnsl-2.5.so'
    _headattrnames = "virtualaddress mode page dev inode filename".split()
    filename = ''

    def __init__(self, line):
        parts = line.split()
        attrs = parts.pop(0)
        for name, value in zip(self._attrnames, map(int, attrs.split(","))):
            setattr(self, name, value)
        for name, value in zip(self._headattrnames, parts):
            setattr(self, name, value)
        self.inode = int(self.inode)

    def __repr__(self):
        attrdata = ",".join([str(getattr(self, name)) for name in self._attrnames])
        headcols = " ".join([str(getattr(self, name)) for name in self._headattrnames])
        return "<Mapping: %s %s>" % (attrdata, headcols)