# Written by Bram Cohen
# see LICENSE.txt for license information

from sha import sha
from threading import Event
true = 1
false = 0

def dummy_status(fractionDone = None, activity = None):
    pass

def dummy_data_flunked(size):
    pass

class StorageWrapper:
    def __init__(self, storage, request_size, hashes, 
            piece_size, finished, failed, 
            statusfunc = dummy_status, flag = Event(), check_hashes = true,
            data_flunked = dummy_data_flunked, backfunc = None,
            config = {}, unpauseflag = None):
        self.storage = storage
        self.request_size = request_size
        self.hashes = hashes
        self.piece_size = piece_size
        self.piece_length = piece_size
        self.data_flunked = data_flunked
        self.backfunc = backfunc
        self.config = config
        self.alloc_type = config.get('alloc_type','normal')
        self.bgalloc_enabled = false
        self.bgalloc_active = false
        self.alloc_buf = chr(0xFF) * piece_size
        self.total_length = storage.get_total_length()
        self.amount_left = self.total_length
        if self.total_length <= piece_size * (len(hashes) - 1):
            raise ValueError, 'bad data from tracker - total too small'
        if self.total_length > piece_size * len(hashes):
            raise ValueError, 'bad data from tracker - total too big'
        self.finished = finished
        self.failed = failed
        self.numactive = [0] * len(hashes)
        self.inactive_requests = [1] * len(hashes)
        self.amount_inactive = self.total_length
        self.endgame = false
        self.have = [false] * len(hashes)
        self.waschecked = [check_hashes] * len(hashes)
        self.places = {}
        self.holes = []
        self.stat_active = {}
        self.stat_new = {}
        self.dirty = {}
        self.stat_numflunked = 0
        self.stat_numdownloaded = 0
        self.stat_numfound = 0
        
        if len(hashes) == 0:
            finished()
            return
        targets = {}
        total = len(hashes)
        for i in xrange(len(hashes)):
            if not self._waspre(i):
                if not targets.has_key(hashes[i]):
                    targets[hashes[i]] = [i]
                else:
                    targets[hashes[i]] = [] # in case of a hash collision, discard
                total -= 1
        numchecked = 0.0
        if total and check_hashes:
            statusfunc(activity = 'checking existing data', fractionDone = 0)
        def markgot(piece, pos, self = self, check_hashes = check_hashes):
            self.places[piece] = pos
            self.have[piece] = true
            self.amount_left -= self._piecelen(piece)
            self.amount_inactive -= self._piecelen(piece)
            self.inactive_requests[piece] = None
            self.waschecked[piece] = check_hashes
            self.stat_numfound += 1
        lastlen = self._piecelen(len(hashes) - 1)
        out_of_place = 0
        updatenum = int(total/300)+1
        updatecount = 0
        for i in xrange(len(hashes)):
            if not self._waspre(i):
                self.holes.append(i)
            elif not check_hashes:
                markgot(i, i)
            else:
                try:
                    v = self.storage.read(piece_size * i, self._piecelen(i))
                except IOError, e:
                        self.failed('IO Error ' + str(e))
                        return
                sh = sha(v[:lastlen])
                sp = sh.digest()
                sh.update(v[lastlen:])
                s = sh.digest()
                if s == hashes[i]:
                    markgot(i, i)
                elif targets.get(s) and self._piecelen(i) == self._piecelen(targets[s][-1]):
                    markgot(targets[s].pop(), i)
                    out_of_place += 1
                elif not self.have[-1] and sp == hashes[-1] and (i == len(hashes) - 1 or not self._waspre(len(hashes) - 1)):
                    markgot(len(hashes) - 1, i)
                    out_of_place += 1
                else:
                    self.places[i] = i
                if unpauseflag is not None and not unpauseflag.isSet():
                    unpauseflag.wait()
                if flag.isSet():
                    return
                numchecked += 1
                updatecount += 1
                if updatecount >= updatenum:
                    updatecount = 0
                    statusfunc(fractionDone = numchecked / total)
        statusfunc(fractionDone = 1.0)
        if self.amount_left == 0:
            finished()
        elif self.alloc_type == 'pre-allocate' and self.holes:
            numholes = len(self.holes)
            statusfunc(activity = 'allocating disk space', fractionDone = 1.0)
            updatenum = int(len(self.holes)/300)+1
            updatecount = 0
            while self.holes:
                if unpauseflag is not None and not unpauseflag.isSet():
                    unpauseflag.wait()
                if flag.isSet():
                    return
                self._doalloc()
                updatecount += 1
                if updatecount >= updatenum:
                    updatecount = 0
                    statusfunc(fractionDone = float(len(self.holes)) / numholes)
            self.storage.flush()
            statusfunc(fractionDone = 0.0)
        elif self.alloc_type == 'sparse':
            self.storage.top_off()  # sets file lengths to their final size
            self.holes = []
            if out_of_place > 0:
                statusfunc(activity = 'moving data', fractionDone = 1.0)
            tomove = out_of_place
            updatenum = int(out_of_place/300)+1
            updatecount = 0
            for i in xrange(len(hashes)):
                if unpauseflag is not None and not unpauseflag.isSet():
                    unpauseflag.wait()
                if flag.isSet():
                    return
                if not self.places.has_key(i):
                    self.places[i] = i
                elif self.places[i] != i:
                    try:
                        old = self.storage.read(self.piece_size * self.places[i], self._piecelen(i))
                        self.storage.write(self.piece_size * i, old)
                    except IOError, e:
                        self.failed('IO Error ' + str(e))
                        return
                    self.places[i] = i
                    tomove -= 1
                    updatecount += 1
                    if updatecount >= updatenum:
                        updatecount = 0
                        statusfunc(fractionDone = float(tomove)/out_of_place)
            self.storage.flush()
            statusfunc(fractionDone = 0.0)
        elif self.alloc_type == 'background' and self.backfunc is not None:
            self.bgalloc()
            

    def bgalloc(self):
        if self.holes and not self.bgalloc_enabled:
            self.bgalloc_enabled = true
            self.bgalloc_active = true
            self.backfunc(self._bgalloc,0.1)
        else:
            self.storage.flush()    # force a flush whenever the "finish allocation" button
                                    # is hit

    def _bgalloc(self):
        if self.holes:
            self._doalloc()
            self.backfunc(self._bgalloc,
                  float(self.piece_size)/(self.config.get('alloc_rate',1.0)*1048576))
        else:
            self.storage.flush()
            self.bgalloc_active = false

    def _doalloc(self):
      try:
        n = self.holes.pop(0)
        if self.places.has_key(n):
            oldpos = self.places[n]
            self.places[oldpos] = oldpos
            old = self.storage.read(self.piece_size * oldpos, self._piecelen(n))
            self.storage.write(self.piece_size * n, old)
        else:
            self.storage.write(self.piece_size * n, self.alloc_buf[:self._piecelen(n)])
        self.places[n] = n
      except IOError, e:
        self.failed('IO Error ' + str(e))

    def _waspre(self, piece):
        return self.storage.was_preallocated(piece * self.piece_size, self._piecelen(piece))

    def _piecelen(self, piece):
        if piece < len(self.hashes) - 1:
            return self.piece_size
        else:
            return self.total_length - piece * self.piece_size

    def get_amount_left(self):
        return self.amount_left

    def do_I_have_anything(self):
        return self.amount_left < self.total_length

    def _make_inactive(self, index):
        length = min(self.piece_size, self.total_length - self.piece_size * index)
        l = []
        x = 0
        while x + self.request_size < length:
            l.append((x, self.request_size))
            x += self.request_size
        l.append((x, length - x))
        self.inactive_requests[index] = l

    def is_endgame(self):
        return self.endgame

    def reset_endgame(self):
        self.endgame = false

    def get_have_list(self):
        return self.have

    def do_I_have(self, index):
        return self.have[index]

    def do_I_have_requests(self, index):
        return not not self.inactive_requests[index]

    def is_unstarted(self, index):
        return ( not self.have[index] and not self.numactive[index]
                 and not self.dirty.has_key(index) )

    def get_hash(self, index):
        return self.hashes[index]

    def new_request(self, index):
        # returns (begin, length)
        if self.inactive_requests[index] == 1:
            self._make_inactive(index)
        self.numactive[index] += 1
        self.stat_active[index] = 1
        if not self.dirty.has_key(index):
            self.stat_new[index] = 1
        rs = self.inactive_requests[index]
        r = min(rs)
        rs.remove(r)
        self.amount_inactive -= r[1]
        if self.amount_inactive == 0:
            self.endgame = true
        return r

    def piece_came_in(self, index, begin, piece):
        try:
            return self._piece_came_in(index, begin, piece)
        except IOError, e:
            self.failed('IO Error ' + str(e))

    def _piece_came_in(self, index, begin, piece):
        if not self.places.has_key(index):
            n = self.holes.pop(0)
            if self.places.has_key(n):
                oldpos = self.places[n]
                old = self.storage.read(self.piece_size * oldpos, self._piecelen(n))
                self.storage.write(self.piece_size * n, old)
                self.places[n] = n
                if index == oldpos or index in self.holes:
                    self.places[index] = oldpos
                else:
                    for p, v in self.places.items():
                        if v == index:
                            break
                    self.places[index] = index
                    self.places[p] = oldpos
                    old = self.storage.read(self.piece_size * index, self.piece_size)
                    self.storage.write(self.piece_size * oldpos, old)
            elif index in self.holes or index == n:
                self.storage.write(self.piece_size * n, self.alloc_buf[:self._piecelen(n)])
                self.places[index] = n
            else:
                for p, v in self.places.items():
                    if v == index:
                        break
                self.places[index] = index
                self.places[p] = n
                old = self.storage.read(self.piece_size * index, self._piecelen(n))
                self.storage.write(self.piece_size * n, old)
        self.storage.write(self.places[index] * self.piece_size + begin, piece)
        self.dirty[index] = 1
        self.numactive[index] -= 1
        if not self.numactive[index]: