"""
Builds on top of blockfs to offer storage of strings of bytes of any size.
Note that any string uses an integral number of 4096-bytes blocks, so
it's not suited to small strings.
"""

import struct, md5
from blockfs import BlockFS


class StringFS(object):

    def __init__(self, blockfs):
        if not isinstance(blockfs, BlockFS):
            blockfs = BlockFS(blockfs)
        self.blockfs = blockfs
        self.blocksize = blockfs.blocksize
        self.digest_size = blockfs.digest_size

    def close(self):
        self.blockfs.close()

    def __getitem__(self, key):
        bs = self.blocksize
        ds = self.digest_size
        buf = self.blockfs[key]
        count, = struct.unpack("!H", buf[-2:])
        data = buf[:count & 0x7fff]
        while count & 0x8000:
            count, = struct.unpack("!H", data[-2:])
            p = len(data) - 2 - (count & 0x7fff)
            tail = data[p:-2]
            subkeys = [data[i:i+ds] for i in range(0, p, ds)]
            result = self.blockfs.getitems(subkeys)
            result.append(tail)
            data = ''.join(result)
            del result
        return data

    def _store(self, buf, reserved=2):
        flags = 0
        bs = self.blocksize
        while 1:
            pad = bs - reserved - len(buf)
            if pad >= 0:
                count = len(buf) | flags
                buf += '\x00' * pad + struct.pack("!H", count)
                return buf
            keys = []
            n = len(buf) // bs
            for i in range(n):
                keys.append(self.blockfs.store(buf[i*bs:(i+1)*bs]))
            tail = buf[n*bs:]
            keys.append(tail)
            count = len(tail) | flags
            keys.append(struct.pack("!H", count))
            buf = ''.join(keys)
            flags = 0x8000

    def store(self, buf):
        return self.blockfs.store(self._store(buf))


def xor16(a, b):
    a0, a1, a2, a3 = struct.unpack("iiii", a)
    b0, b1, b2, b3 = struct.unpack("iiii", b)
    return struct.pack("iiii", a0^b0, a1^b1, a2^b2, a3^b3)

class TweakedBlockFS(BlockFS):
    MAGIC = 'BlkFs\x02\x0c'
    digest_size = md5.digest_size
    assert digest_size == 16

    def digest(self, blockdata):
        m = md5.new(blockdata[:-18])
        m.update(blockdata[-2:])
        return xor16(m.digest(), blockdata[-18:-2])

class MD5KeyStringFS(StringFS):
    """A variant of StringFS where the keys returned by store() are
    precisely the MD5 checksum of the string.
    """
    BlockFS = TweakedBlockFS

    def __init__(self, path):
        StringFS.__init__(self, self.BlockFS(path))

    def store(self, buf):
        checksum = md5.new(buf).digest()
        buf = self._store(buf, 18)
        bufbase = buf[:-2]
        buftail = buf[-2:]
        m = md5.new(bufbase)
        m.update(buftail)
        buf = bufbase + (xor16(m.digest(), checksum) + buftail)
        key = self.blockfs.store(buf)
        assert key == checksum
        return checksum