#===============================================================================
# Copyright 2009 Matt Chaput
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#    http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#===============================================================================

"""This module defines writer and reader classes for a fast, immutable
on-disk key-value database format. The current format is identical
to D. J. Bernstein's CDB format (http://cr.yp.to/cdb.html).
"""

from array import array
from bisect import bisect_right
from collections import defaultdict
from cPickle import dumps, loads
from struct import Struct

from whoosh.system import _USHORT_SIZE


def cdb_hash(key):
    h = 5381L
    for c in key:
        h = (h + (h << 5)) & 0xffffffffL ^ ord(c)
    return h

# Read/write convenience functions

_2ints = Struct("!II")
pack2ints = _2ints.pack
def writeints(f, value1, value2):
    f.write(pack2ints(value1, value2))

_unpack2ints = _2ints.unpack
def unpack2ints(s):
    return _unpack2ints(s)

_unpackint = Struct("!I").unpack
def readint(map, offset):
    return _unpackint(map[offset:offset+4])[0]

# Encoders and decoders for storing complex types in
# string -> string hash files.

_int_struct = Struct("!i")
packint = _int_struct.pack
_unpackint = _int_struct.unpack
def unpackint(s):
    return _unpackint(s)[0]

_ushort_struct = Struct("!H")
packushort = _ushort_struct.pack
_unpackushort = _ushort_struct.unpack
def unpackushort(s):
    return _unpackushort(s)[0]

def encode_termkey(term):
    fieldnum, text = term
    return packushort(fieldnum) + text.encode("utf8")

def decode_termkey(key):
    return unpackushort(key[:_USHORT_SIZE]), key[_USHORT_SIZE:].decode("utf8")

_vkey_struct = Struct("!Ii")
_pack_vkey = _vkey_struct.pack
def encode_vectorkey(docandfield):
    return _pack_vkey(*docandfield)

decode_vectorkey = _vkey_struct.unpack
encode_docnum = packint
decode_docnum = unpackint

_terminfo_struct = Struct("!ILI")
_terminfo_pack = _terminfo_struct.pack
def encode_terminfo(cf_offset_df):
    return _terminfo_pack(*cf_offset_df)
decode_terminfo = _terminfo_struct.unpack

def enpickle(data):
    "Encodes a value as a string for storage in a table."
    return dumps(data, -1)

depickle = loads


# Table classes

class FileHashWriter(object):
    def __init__(self, dbfile):
        self.dbfile = dbfile
        dbfile.seek(2048)
        self.hashes = defaultdict(list)
    
    def add_all(self, items):
        dbfile = self.dbfile
        hashes = self.hashes
        pos = dbfile.tell()
        write = dbfile.write
        
        for key, value in items:
            writeints(dbfile, len(key), len(value))
            write(key + value)
            
            h = cdb_hash(key)
            hashes[h & 255].append((h, pos))
            pos += len(key) + len(value) + 8
    
    def add(self, key, value):
        self.add_all(((key, value), ))
    
    def add_key(self, key):
        dbfile = self.dbfile
        
        writeints(dbfile)
    
    def _write_hashes(self):
        dbfile = self.dbfile
        hashes = self.hashes
        directory = self.directory = []
        
        pos = dbfile.tell()
        for i in xrange(0, 256):
            entries = hashes[i]
            numslots = 2 * len(entries)
            directory.append((pos, numslots))
            
            null = (0, 0)
            hashtable = [null] * numslots
            for hashval, position in entries:
                n = (hashval >> 8) % numslots
                while hashtable[n] is not null:
                    n = (n + 1) % numslots
                hashtable[n] = (hashval, position)
            
            for hashval, position in hashtable:
                writeints(dbfile, hashval, position)
                pos += 8
                
        dbfile.flush()
    
    def _write_directory(self):
        dbfile = self.dbfile
        directory = self.directory
        
        dbfile.seek(0)
        for position, numslots in directory:
            writeints(dbfile, position, numslots)
        assert dbfile.tell() == 2048
        dbfile.flush()
    
    def close(self):
        self._write_hashes()
        self._write_directory()
        self.dbfile.close()


class FileHashReader(object):
    def __init__(self, dbfile):
        self.dbfile = dbfile
        self.map = dbfile.map
        self.end_of_data = dbfile.get_uint(0)
        self.is_closed = False
        
    def close(self):
        if self.is_closed:
            raise Exception("Tried to close %r twice" % self)
        del self.map
        self.dbfile.close()
        self.is_closed = True
        
    def read(self, position, length):
        return self.map[position:position+length]

    def read2ints(self, position):
        return unpack2ints(self.map[position:position+8])

    def _ranges(self, pos=2048):
        read2ints = self.read2ints
        eod = self.end_of_data
        
        while pos < eod:
            keylen, datalen = read2ints(pos)
            keypos = pos + 8
            datapos = pos + 8 + keylen
            pos = datapos + datalen
            yield (keypos, keylen, datapos, datalen)

    def __iter__(self):
        return self.items()

    def items(self):
        read = self.read
        for keypos, keylen, datapos, datalen in self._ranges():
            yield (read(keypos, keylen), read(datapos, datalen))

    def keys(self):
        read = self.read
        for keypos, keylen, _, _ in self._ranges():
            yield read(keypos, keylen)

    def values(self):
        read = self.read
        for _, _, datapos, datalen in self._ranges():
            yield read(datapos, datalen)

    def __getitem__(self, key):
        for data in self.all(key):
            return data
        raise KeyError(key)

    def get(self, key, default=None):
        for data in self.all(key):
            return data
        return default

    def _hashtable_info(self, keyhash):
        return self.read2ints(keyhash << 3 & 2047)

    def _key_position(self, key):
        keyhash = cdb_hash(key)
        hpos, hslots = self._hashtable_info(keyhash)
        if not hslots:
            raise KeyError(key)
        slotpos = hpos + (((keyhash >> 8) % hslots) << 3)
        u, pos = self.read2ints(slotpos)
        return pos

    def _get_ranges(self, key):
        read = self.read
        read2ints = self.read2ints
        keyhash = cdb_hash(key)
        hpos, hslots = self._hashtable_info(keyhash)
        if not hslots:
            return
        
        slotpos = hpos + (((keyhash >> 8) % hslots) << 3)
        for _ in xrange(0, hslots):
            u, pos = read2ints(slotpos)
            if not pos:
                return
            
            slotpos += 8
            # If we reach the end of the hashtable, wrap around
            if slotpos == hpos + (hslots << 3):
                slotpos = hpos
            
            if u == keyhash:
                keylen, datalen = read2ints(pos)
                if keylen == len(key):
                    if key == read(pos + 8, keylen):
                        yield (pos + 8 + keylen, datalen)
        
    def all(self, key):
        read = self.read
        for datapos, datalen in self._get_ranges(key):
            yield read(datapos, datalen)
            
    def __contains__(self, key):
        for _ in self._get_ranges(key):
            return True
        return False


class OrderedHashWriter(FileHashWriter):
    def __init__(self, dbfile, blocksize = 100):
        FileHashWriter.__init__(self, dbfile)
        self.blocksize = blocksize
        self.index = []
        self.indexcount = None
        self.lastkey = None
        
    def add_all(self, items):
        dbfile = self.dbfile
        hashes = self.hashes
        pos = dbfile.tell()
        write = dbfile.write
        
        ix = self.index
        ic = self.indexcount
        bs = self.blocksize
        lk = self.lastkey
        
        for key, value in items:
            if key <= lk:
                raise ValueError("Keys must increase: %r .. %r" % (lk, key))
            lk = key
            
            if ic is None:
                ix.append(key)
                ic = 0
            else:
                ic += 1
                if ic == bs:
                    ix.append(key)
                    ic = 0
            
            writeints(dbfile, len(key), len(value))
            write(key + value)
            
            h = cdb_hash(key)
            hashes[h & 255].append((h, pos))
            pos += len(key) + len(value) + 8
        
        self.indexcount = ic
        self.lastkey = lk
        
    def close(self):
        self._write_hashes()
        self.dbfile.write_pickle(self.index)
        self._write_directory()
        self.dbfile.close()

      
class OrderedHashReader(FileHashReader):
    def __init__(self, dbfile):
        FileHashReader.__init__(self, dbfile)
        lastpos, lastnum = self.read2ints(255*8)
        dbfile.seek(lastpos + lastnum * 8)
        self.index = dbfile.read_pickle()

    def _closest_key(self, key):
        index = self.index
        i = max(0, bisect_right(index, key) - 1)
        return index[i]
    
    def _ranges_from(self, key):
        read = self.read
        ckey = self._closest_key(key)
        pos = self._key_position(ckey)
        
        if ckey != key:
            for keypos, keylen, _, _ in self._ranges(pos=pos):
                k = read(keypos, keylen)
                if k >= key:
                    pos = keypos - 8
                    break
        
        return self._ranges(pos=pos)

    def items_from(self, key):
        read = self.read
        for keypos, keylen, datapos, datalen in self._ranges_from(key):
            yield (read(keypos, keylen), read(datapos, datalen))

    def keys_from(self, key):
        read = self.read
        for keypos, keylen, _, _ in self._ranges_from(key):
            yield read(keypos, keylen)

    def values(self, key):
        read = self.read
        for _, _, datapos, datalen in self._ranges_from(key):
            yield read(datapos, datalen)


class FileTableWriter(OrderedHashWriter):
    def __init__(self, dbfile, keycoder=None, valuecoder=None):
        sup = super(FileTableWriter, self)
        sup.__init__(dbfile)
        self.keycoder = keycoder or str
        self.valuecoder = valuecoder or enpickle
        
        self._add = sup.add
        
    def add(self, key, data):
        key = self.keycoder(key)
        data = self.valuecoder(data)
        self._add(key, data)
        

class FileTableReader(OrderedHashReader):
    def __init__(self, dbfile, keycoder=None, keydecoder=None, valuedecoder=None):
        sup = super(FileTableReader, self)
        sup.__init__(dbfile)
        self.keycoder = keycoder or str
        self.keydecoder = keydecoder or int
        self.valuedecoder = valuedecoder or depickle
        
        self._items = sup.items
        self._items_from = sup.items_from
        self._keys = sup.keys
        self._keys_from = sup.keys_from
        self._getitem = sup.__getitem__
        self._contains = sup.__contains__
        
    def __getitem__(self, key):
        k = self.keycoder(key)
        return self.valuedecoder(self._getitem(k))
    
    def __contains__(self, key):
        return self._contains(self.keycoder(key))
    
    def items(self):
        kd = self.keydecoder
        vd = self.valuedecoder
        for key, value in self._items():
            yield (kd(key), vd(value))
            
    def items_from(self, key):
        fromkey = self.keycoder(key)
        kd = self.keydecoder
        vd = self.valuedecoder
        for key, value in self._items_from(fromkey):
            yield (kd(key), vd(value))
            
    def keys(self):
        kd = self.keydecoder
        for k in self._keys():
            yield kd(k)
            
    def keys_from(self, key):
        kd = self.keydecoder
        for k in self._keys_from(self.keycoder(key)):
            yield kd(k)
            

class FileRecordWriter(object):
    def __init__(self, dbfile, format):
        self.dbfile = dbfile
        self.format = format
        self._pack = Struct(format).pack
        
    def close(self):
        self.dbfile.close()
        
    def append(self, args):
        self.dbfile.write(self._pack(*args))


class FileRecordReader(object):
    def __init__(self, dbfile, format):
        self.dbfile = dbfile
        self.map = dbfile.map
        self.format = format
        struct = Struct(format)
        self._unpack = struct.unpack
        self.itemsize = struct.size
    
    def close(self):
        del self.map
        self.dbfile.close()
    
    def record(self, recordnum):
        itemsize = self.itemsize
        return self._unpack(self.map[recordnum * itemsize: recordnum * itemsize + itemsize])
    
    def at(self, recordnum, itemnum):
        return self.record(recordnum)[itemnum]


class FileListWriter(object):
    def __init__(self, dbfile, valuecoder=str):
        self.dbfile = dbfile
        self.positions = array("I")
        self.lengths = array("I")
        dbfile.write_uint(0)
        self.valuecoder = valuecoder
    
    def close(self):
        f = self.dbfile
        directory_pos = f.tell()
        f.write_array(self.positions)
        f.write_array(self.lengths)
        f.flush()
        f.seek(0)
        f.write_uint(directory_pos)
        f.close()
    
    def append(self, value):
        f = self.dbfile
        self.positions.append(f.tell())
        v = self.valuecoder(value)
        self.lengths.append(len(v))
        f.write(v)
        

class FileListReader(object):
    def __init__(self, dbfile, length, valuedecoder=str):
        self.dbfile = dbfile
        self.length = length
        self.valuedecoder = valuedecoder
        
        offset = dbfile.get_uint(0)
        dbfile.seek(offset)
        self.positions = dbfile.read_array("I", length)
        self.lengths = dbfile.read_array("I", length)
    
    def close(self):
        self.dbfile.close()
    
    def __getitem__(self, num):
        position = self.positions[num]
        length = self.lengths[num]
        v = self.dbfile.map[position:position+length]
        return self.valuedecoder(v)
    

# Utility functions

def dump_hash(hashreader):
    dbfile = hashreader.dbfile
    read = hashreader.read
    read2ints = hashreader.read2ints
    eod = hashreader.end_of_data
    
    # Dump hashtables
    for bucketnum in xrange(0, 255):
        pos, numslots = read2ints(bucketnum * 8)
        if numslots:
            print "Bucket %d: %d slots" % (bucketnum, numslots)
            
            dbfile.seek(pos)
            for j in xrange(0, numslots):
                print "  %X : %d" % read2ints(pos)
                pos += 8
    
    # Dump keys and values
    print "-----"
    dbfile.seek(2048)
    pos = 2048
    while pos < eod:
        keylen, datalen = read2ints(pos)
        keypos = pos + 8
        datapos = pos + 8 + keylen
        key = read(keypos, keylen)
        data = read(datapos, datalen)
        print "%d +%d,%d:%r->%r" % (pos, keylen, datalen, key, data)
        pos = datapos + datalen
        



    