try:
    from primitives import Mem
except ImportError:
    from mem import Mem

import sys
if sys.version >= '3':
    xrange = range

class MMU():
    def __init__(self, mem, size=0):
        """ Initialize MMU
        """
        self._enabled = False
        self._mem = mem
        self._wordsize = 4
        self._table = []

    def isEnabled(self):
        return self._enabled

    def enable(self):
        """ Enables MMU
        """
        self._enabled = True

    def disable(self):
        """ Disables MMU
        """
        self._enabled = False

    def getEntries(self, entries, startpos=None):
        """ Get page entries and parse them, handle recursively

        >>> from primitives import Mem
        >>> m = Mem(1024*1024)
        >>> m.setData(0, 0x00000100, 4)
        >>> m.setData(4, 0x00001100, 4)
        >>> m.setData(8, 0x00002100, 4)
        >>> m.setData(12, 0x00003100, 4)
        >>> u = MMU(m)
        >>> entries = [(0, MMU.Flags(solved={'execute': False, 'ok': True, 'size1': False, 'size2': False, 'write': False, 'subtable': True, 'userspace': False, 'size': 4*1024}), 0),
        ... (32768, MMU.Flags(solved={'execute': False, 'ok': True, 'size1': True, 'size2': False, 'write': False, 'subtable': False, 'userspace': False, 'size': 64}), 65536),
        ... (0, MMU.Flags(solved={'execute': False, 'ok': True, 'size1': False, 'size2': False, 'write': False, 'subtable': False, 'userspace': False, 'size': 4}), 131072)]
        >>> u.getEntries(entries)
        [(0, execute=False,ok=True,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 0), (4096, execute=False,ok=True,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 4096), (8192, execute=False,ok=True,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 8192), (12288, execute=False,ok=True,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 12288), (32768, execute=False,ok=True,size=64,size1=True,size2=False,subtable=False,userspace=False,write=False, 65536), (0, execute=False,ok=True,size=4,size1=False,size2=False,subtable=False,u
serspace=False,write=False, 131072)]
        """
        if startpos is None:
            startpos = 0
        subs = []
        for (addr, flags, pos) in entries:
            if flags['subtable']:
                size = flags['size'] * 1024 / 4
                if flags['ok']:
                    tmp = self.readTable(addr, size, pos)
                    entries = self.getEntries(tmp, startpos)
                    subs += entries
                    startpos += flags['size'] * 1024
        
    else:
                if flags['ok']:
                    subs.append((addr, flags, pos))
        return subs

    def initialize(self, tablepos, tablesize):
        """ Initializes MMU with a initial page
        Does recursive parsing

        >>> from primitives import Mem
        >>> m = Mem(1024*1024*10)
        >>> u = MMU(m)
        >>> # Subtable, starts at phys 4k
        >>> m.setData(22, 0x00001111, 4)
        >>> # Page, virtual start at 32k, size 64k
        >>> m.setData(14, 0x00008110, 4)
        >>> # Page, virtual start at 98k, size 4k
        >>> m.setData(18, 0x00018100, 4)
        >>> for i in xrange(1023):
        ...   m.setData(0x1000 + i, 0)
        >>> # Page at 8k, size 4k
        >>> m.setData(0x1000, 0x00002100, 4)
        >>> # Page at 12k, size 1M 
        >>> m.setData(0x1004, 0x00003120, 4)
        >>> u.initialize(14, 3)
        [(32768, execute=False,ok=True,size=64,size1=True,size2=False,subtable=False,userspace=False,write=False, 0), (98304, execute=False,ok=True,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 65536), (8192, execute=False,ok=True,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 69632), (12288, execute=False,ok=True,size=1024,size1=False,size2=True,subtable=False,userspace=False,write=False, 73728)]
        """
        entries = self.readTable(tablepos, tablesize)
        self._table = self.getEntries(entries)
        return self._table

    def diffTime(self, a, b):
        d = a - b
        print (d.seconds*1000*1000 + d.microseconds)

    def readTable(self, tablepos, tablesize, pos=None):
        """ Reads table from memory

        >>> from primitives import Mem
        >>> m = Mem(1024*100)
        >>> u = MMU(m)
        >>> # Subtable, starts at phys 4k
        >>> m.setData(10, 0x00001111, 4)
        >>> # Page, starts at 32k, size 64k
        >>> m.setData(14, 0x00008110, 4)
        >>> for i in xrange(1023):
        ...   m.setData(0x1000 + i, 0)
        >>> tmp = u.readTable(10, 3)
        >>> tmp[0][0]
        4096
        >>> tmp[1][0]
        32768
        >>> tmp[0][1]
        execute=False,ok=True,size=65536,size1=True,size2=False,subtable=True,userspace=False,write=False
        >>> tmp[1][1]
        execute=False,ok=True,size=64,size1=True,size2=False,subtable=False,userspace=False,write=False
        >>> tmp[0]
        (4096, execute=False,ok=True,size=65536,size1=True,size2=False,subtable=True,userspace=False,write=False, 0)
        >>> tmp[1]
        (32768, execute=False,ok=True,size=64,size1=True,size2=False,subtable=False,userspace=False,write=False, 67108864)
        """
        import datetime
        datas = []
        if pos is None:
            pos = 0

        virtpos = tablepos
        cnt = 0

        # Optimized reading in blocks instead of one byte at time
        block = self._mem.getBlock(tablepos, tablesize * 4)
        oldtmp = 0
        items = 0
        preindex = 0
        for (bpos, data) in block:
            if data is None:
                continue

            if preindex > 0:
                # Do we have old data from previous block?
                if preindex == 1:
                    oldtmp += (data[0] << 24)
                if preindex == 2:
                    oldtmp += (data[0] << 16) + (data[1] << 24)
                if preindex == 3:
                    oldtmp += (data[0] << 8) + (data[1] << 16) + (data[2] << 24)

                (ok, pos, res) = self.readEntry(oldtmp, pos)
                if ok:
                    datas.append(res)
                tablepos = preindex

            datalen = len(data)
            l = int(datalen / 4 - 1)
            index = tablepos % 0x1000
            for item in xrange(l):
                tmp = data[index] + (data[index+1] << 8) + (data[index+2] << 16) + (data[index+3] << 24)
                (ok, pos, res) = self.readEntry(tmp, pos)
                if ok:
                    datas.append(res)
                index += 4
                items += 4
                if index > datalen - 4:
                    miss = datalen - index
                    preindex = 0

                    # Check if we didn't read all the data...
                    if miss > 0:
                        oldtmp = data[index]
                        if miss > 1:
                            oldtmp += (data[index+1] << 8)
                        if miss > 2:
                            oldtmp += (data[index+2] << 16)
                        preindex = 4 - miss
                    break
                if items > (tablesize + tablepos):
                    break
        return datas

        """
        for index in xrange(tablesize):
            tmp = self._mem.getData(virtpos, self._wordsize)
            virtpos += self._wordsize
            if tmp > 0:
                print tmp
                cnt += 1
                (ok, pos, res) = self.readEntry(tmp, pos)
                if ok:
                    datas.append(res)

        return datas
        """

    def readEntry(self, data, pos=0):
        """ Read entry from one page table item data

        >>> m = Mem()
        >>> u = MMU(m)
        >>> u.readEntry(0x00000000)
        (False, 0, (0, execute=False,ok=False,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 0))
        >>> u.readEntry(0x00001000)
        (True, 4096, (4096, execute=False,ok=False,size=4,size1=False,size2=False,subtable=False,userspace=False,write=False, 0))
        >>> u.readEntry(0x00001111)
        (True, 67108864, (4096, exe