
#===============================================================================
# Copyright 2010 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.
#===============================================================================

from __future__ import with_statement
import os
import tempfile
from array import array
from collections import defaultdict
from heapq import heapify, heappop, heapreplace
from marshal import load, dump
#import sqlite3 as sqlite

from whoosh.filedb.filetables import LengthWriter, LengthReader
from whoosh.util import length_to_byte


try:
    from sys import getsizeof
except ImportError:
    # If this is Python 2.5, rig up a guesstimated version of getsizeof
    def getsizeof(obj):
        if obj is None:
            return 8
        t = type(obj)
        if t is int:
            return 12
        elif t is float:
            return 16
        elif t is long:
            return 16
        elif t is str:
            return 21 + len(obj)
        elif t is unicode:
            return 26 + 2 * len(obj)


try:
    from heapq import merge
    def imerge(iterables):
        return merge(*iterables)
except ImportError:
    def imerge(iterables):
        """Merge-sorts items from a list of iterators.
        """
        
        _heappop, _heapreplace, _StopIteration = heappop, heapreplace, StopIteration
    
        h = []
        h_append = h.append
        for itnum, it in enumerate(map(iter, iterables)):
            try:
                next = it.next
                h_append([next(), itnum, next])
            except _StopIteration:
                pass
        heapify(h)
    
        while 1:
            try:
                while 1:
                    v, itnum, next = s = h[0]   # raises IndexError when h is empty
                    yield v
                    s[0] = next()               # raises StopIteration when exhausted
                    _heapreplace(h, s)          # restore heap condition
            except _StopIteration:
                _heappop(h)                     # remove empty iterator
            except IndexError:
                return


def read_run(filename, count, atatime=100):
    with open(filename, "rb") as f:
        while count:
            buff = []
            take = min(atatime, count)
            for _ in xrange(take):
                buff.append(load(f))
            count -= take
            for item in buff:
                yield item


DEBUG_DIR = False


class PoolBase(object):
    def __init__(self, schema, dir=None, basename=''):
        self.schema = schema
        self._using_tempdir = False
        self.dir = dir
        self._using_tempdir = dir is None
        self.basename = basename
        
        self.length_arrays = {}
        self._fieldlength_totals = defaultdict(int)
        self._fieldlength_maxes = {}
    
    def _make_dir(self):
        if self.dir is None:
            self.dir = tempfile.mkdtemp(".whoosh")
            
            if DEBUG_DIR:
                dfile = open(self._filename("DEBUG.txt"), "wb")
                import traceback
                traceback.print_stack(file=dfile)
                dfile.close()
    
    def _filename(self, name):
        return os.path.abspath(os.path.join(self.dir, self.basename + name))
    
    def _clean_temp_dir(self):
        if self._using_tempdir and self.dir and os.path.exists(self.dir):
            if DEBUG_DIR:
                os.remove(self._filename("DEBUG.txt"))
            
            try:
                os.rmdir(self.dir)
            except OSError:
                # directory didn't exist or was not empty -- don't
                # accidentially delete data
                pass
    
    def cleanup(self):
        self._clean_temp_dir()
    
    def cancel(self):
        pass
    
    def fieldlength_totals(self):
        return dict(self._fieldlength_totals)
    
    def fieldlength_maxes(self):
        return self._fieldlength_maxes
    
    def add_posting(self, fieldname, text, docnum, weight, valuestring):
        raise NotImplementedError
    
    def add_field_length(self, docnum, fieldname, length):
        self._fieldlength_totals[fieldname] += length
        if length > self._fieldlength_maxes.get(fieldname, 0):
            self._fieldlength_maxes[fieldname] = length
        
        if fieldname not in self.length_arrays:
            self.length_arrays[fieldname] = array("B")
        arry = self.length_arrays[fieldname]
        
        if len(arry) <= docnum:
            for _ in xrange(docnum - len(arry) + 1):
                arry.append(0)
        arry[docnum] = length_to_byte(length)
    
    def _fill_lengths(self, doccount):
        for fieldname in self.length_arrays.keys():
            arry = self.length_arrays[fieldname]
            if len(arry) < doccount:
                for _ in xrange(doccount - len(arry)):
                    arry.append(0)
    
    def add_content(self, docnum, fieldname, field, value):
        add_posting = self.add_posting
        termcount = 0
        # TODO: Method for adding progressive field values, ie
        # setting start_pos/start_char?
        for w, freq, weight, valuestring in field.index(value):
            #assert w != ""
            add_posting(fieldname, w, docnum, weight, valuestring)
            termcount += freq
        
        if field.scorable and termcount:
            self.add_field_length(docnum, fieldname, termcount)
            
        return termcount
    
    def _write_lengths(self, lengthfile, doccount):
        self._fill_lengths(doccount)
        lw = LengthWriter(lengthfile, doccount, lengths=self.length_arrays)
        lw.close()


class TempfilePool(PoolBase):
    def __init__(self, schema, limitmb=32, dir=None, basename='', **kw):
        super(TempfilePool, self).__init__(schema, dir=dir, basename=basename)
        
        self.limit = limitmb * 1024 * 1024
        
        self.size = 0
        self.count = 0
        self.postings = []
        self.runs = []
        
    def add_posting(self, fieldname, text, docnum, weight, valuestring):
        if self.size >= self.limit:
            self.dump_run()

        tup = (fieldname, text, docnum, weight, valuestring)
        # 48 bytes for tuple overhead (28 bytes + 4 bytes * 5 items) plus the
        # sizes of the objects inside the tuple, plus 4 bytes overhead for
        # putting the tuple in the postings list
        #self.size += 48 + sum(getsizeof(o) for o in tup) + 4
        valsize = len(valuestring) if valuestring else 0
        self.size += 48 + len(fieldname) + 22 + len(text) + 26 + 16 + 16 + valsize + 22 + 4
        self.postings.append(tup)
        self.count += 1
    
    def dump_run(self):
        if self.size > 0:
            self._make_dir()
            fd, filename = tempfile.mkstemp(".run", dir=self.dir)
            runfile = os.fdopen(fd, "w+b")
            self.postings.sort()
            for p in self.postings:
                dump(p, runfile)
            runfile.close()
            
            self.runs.append((filename, self.count))
            self.postings = []
            self.size = 0
            self.count = 0
    
    def run_filenames(self):
        return [filename for filename, _ in self.runs]
    
    def cancel(self):
        self.cleanup()
    
    def cleanup(self):
        for filename in self.run_filenames():
            if os.path.exists(filename):
                try:
                    os.remove(filename)
                except IOError:
                    pass
        
        self._clean_temp_dir()
        
    def finish(self, termswriter, doccount, lengthfile):
        self._write_lengths(lengthfile, doccount)
        lengths = LengthReader(None, doccount, self.length_arrays)
        
        if self.postings or self.runs:
            if self.postings and len(self.runs) == 0:
                self.postings.sort()
                postiter = iter(self.postings)
            elif not self.postings and not self.runs:
                postiter = iter([])
            else:
                self.dump_run()
                postiter = imerge([read_run(runname, count)
                                   for runname, count in self.runs])
        
            termswriter.add_iter(postiter, lengths.get)
        self.cleanup()
        

# Alternative experimental and testing pools

class SqlitePool(PoolBase):
    def __init__(self, schema, dir=None, basename='', limitmb=32, **kwargs):
        super(SqlitePool, self).__init__(schema, dir=dir, basename=basename)
        self._make_dir()
        self.postbuf = defaultdict(list)
        self.bufsize = 0
        self.limit = limitmb * 1024 * 1024
        self.fieldnames = set()
        self._flushed = False
    
    def _field_filename(self, name):
        return self._filename("%s.sqlite" % name)
    
    def _con(self, name):
        import sqlite3 as sqlite
        
        filename = self._field_filename(name)
        con = sqlite.connect(filename)
        if name not in self.fieldnames:
            self.fieldnames.add(name)
            con.execute("create table postings (token text, docnum int, weight float, value blob)")
            #con.execute("create index postix on postings (token, docnum)")
        return con
    
    def flush(self):
        for fieldname, lst in self.postbuf.iteritems():
            con = self._con(fieldname)
            con.executemany("insert into postings values (?, ?, ?, ?)", lst)
            con.commit()
            con.close()
        self.postbuf = defaultdict(list)
        self.bufsize = 0
        self._flushed = True
        print "flushed"
    
    def add_posting(self, fieldname, text, docnum, weight, valuestring):
        self.postbuf[fieldname].append((text, docnum, weight, valuestring))
        self.bufsize += len(text) + 8 + len(valuestring)
        if self.bufsize > self.limit:
            self.flush()
    
    def readback(self):
        for name in sorted(self.fieldnames):
            con = self._con(name)
            con.execute("create index postix on postings (token, docnum)")
            for text, docnum, weight, valuestring in con.execute("select * from postings order by token, docnum"):
                yield (name, text, docnum, weight, valuestring)
            con.close()
            os.remove(self._field_filename(name))
        
        if self._using_tempdir and self.dir:
            try:
                os.rmdir(self.dir)
            except OSError:
                # directory didn't exist or was not empty -- don't
                # accidentially delete data
                pass
    
    def readback_buffer(self):
        for fieldname in sorted(self.postbuf.keys()):
            lst = self.postbuf[fieldname]
            lst.sort()
            for text, docnum, weight, valuestring in lst:
                yield (fieldname, text, docnum, weight, valuestring)
            del self.postbuf[fieldname]
            
    def finish(self, termswriter, doccount, lengthfile):
        self._write_lengths(lengthfile, doccount)
        lengths = LengthReader(None, doccount, self.length_arrays)
        
        if not self._flushed:
            gen = self.readback_buffer()
        else:
            if self.postbuf:
                self.flush()
            gen = self.readback()
        
        termswriter.add_iter(gen, lengths.get)
    

class NullPool(PoolBase):
    def __init__(self, *args, **kwargs):
        self._fieldlength_totals = {}
        self._fieldlength_maxes = {}
    
    def add_content(self, *args):
        pass
    
    def add_posting(self, *args):
        pass
    
    def add_field_length(self, *args, **kwargs):
        pass
    
    def finish(self, *args):
        pass
        

class MemPool(PoolBase):
    def __init__(self, schema, **kwargs):
        super(MemPool, self).__init__(schema)
        self.schema = schema
        self.postbuf = []
        
    def add_posting(self, *item):
        self.postbuf.append(item)
        
    def finish(self, termswriter, doccount, lengthfile):
        self._write_lengths(lengthfile, doccount)
        lengths = LengthReader(None, doccount, self.length_arrays)
        self.postbuf.sort()
        termswriter.add_iter(self.postbuf, lengths.get)


#class UnixSortPool(PoolBase):
#    def __init__(self, schema, dir=None, basename='', limitmb=32, **kwargs):
#        super(UnixSortPool, self).__init__(schema, dir=dir, basename=basename)
#        self._make_dir()
#        fd, self.filename = tempfile.mkstemp(".run", dir=self.dir)
#        self.sortfile = os.fdopen(fd, "wb")
#        self.linebuffer = []
#        self.bufferlimit = 100
#        
#    def add_posting(self, *args):
#        self.sortfile.write(b64encode(dumps(args)) + "\n")
#        
#    def finish(self, termswriter, doccount, lengthfile):
#        self.sortfile.close()
#        from whoosh.util import now
#        print "Sorting file...", self.filename
#        t = now()
#        outpath = os.path.join(os.path.dirname(self.filename), "sorted.txt")
#        os.system("sort %s >%s" % (self.filename, outpath))
#        print "...took", now() - t

    
    



























