#!/usr/bin/env python

# RaidGuessFS, a FUSE pseudo-filesystem to guess RAID parameters of a damaged device
# Copyright (C) 2015 Ludovic Pouzenc <ludovic@pouzenc.fr>
#
# This file is part of RaidGuessFS.
#
# RaidGuessFS is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# RaidGuessFS is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with RaidGuessFS. If not, see <http://www.gnu.org/licenses/>

import logging, numpy

class MyRaid():
    """Auxiliary class, managing RAID layer"""
    RAID_TYPES = [ '0', '1', '5', '5+0' ]
    RAID5_LAYOUTS = [ 'la', 'ra', 'ls', 'rs' ]

    def __init__(self, *args, **kwargs):
        self.disks = []
        self.raid_start = 0
        self.raid_end = 0
        self.raid_sector_size = 512
        self.raid_chunk_size = 65536
        self.raid_disk_order = []
        self.raid_disk_count = 0
        self.raid_layout = 'ls'

    def get_raid_start(self):
        return self.raid_start

    def get_raid_end(self):
        return self.raid_end

    def get_raid_chunk_size(self):
        return self.raid_chunk_size

    def get_raid_disk_order(self):
        return self.raid_disk_order

    def get_raid_disk_order_str(self):
        return ' '.join(map(str,self.raid_disk_order))

    def get_raid_layout(self):
        return self.raid_layout

    def set_disks(self, disks):
        self.disks = disks

    def set_raid_start(self, new_raid_start):
        """Update the start offset of raid data on underlying disks"""
        self.raid_start = new_raid_start

    def set_raid_end(self, new_raid_end):
        """Update the end offset of raid data on underlying disks"""
        self.raid_end = new_raid_end

    def set_raid_chunk_size(self, new_raid_chunk_size):
        """Update the size of chucks of data (or slice size)"""
        self.raid_chunk_size = new_raid_chunk_size

    def set_raid_disk_order(self, new_raid_disk_order):
        """Update the raid logical disk order"""
        card=len(self.disks)
        check=[0]*card
        for item in new_raid_disk_order:
            d = int(item)
            if not 0 <= d < card:
                raise ValueError('Value out of range : %i [0,%i]'%(d,card-1))
            check[d]=check[d]+1
        
        for d in range(card):
            if check[d] != 1 and check[d] != 0:
                raise ValueError('Disk %i appears %i times (must be 0 or 1)'%(d,check[d]))
        self.raid_disk_order = new_raid_disk_order
        self.raid_disk_count = len(new_raid_disk_order)

    def set_raid_layout(self, new_raid_layout):
        if new_raid_layout in MyRaid.RAID5_LAYOUTS:
            self.raid_layout = new_raid_layout
        else:
            raise ValueError('raid_layout has to be one of %s'%' '.join(RAID_LAYOUTS))

    def sizeof_raid_result(self, raid_type):
        size = self.raid_end - self.raid_start
        if size <= 0 :
            return 0
        else:
            return {
                '0'  : size * self.raid_disk_count,
                '1'  : size if self.raid_disk_count == 2 else 0,
                '5'  : size * (self.raid_disk_count - 1) if self.raid_disk_count >= 3 else 0,
                '5+0': size * (self.raid_disk_count - 2) if self.raid_disk_count >= 6 and self.raid_disk_count % 2 == 0 else 0,
                }[raid_type]

    def sizeof_disk_xor(self, raid_type):
        size = self.raid_end - self.raid_start
        if size <= 0:
            return 0
        else:
            return {
                '0'  : 0, # TODO Could contain some plain text error message
                '1'  : size if self.raid_disk_count == 2 else 0,
                '5'  : size if self.raid_disk_count >= 3 else 0,
                '5+0': size if self.raid_disk_count >= 6 and self.raid_disk_count % 2 == 0 else 0,
                }[raid_type]

    def sizeof_disk_parity(self, raid_type):
        return self.sizeof_disk_xor(raid_type) / self.raid_sector_size * 16

    def xor_blocks(self,fd_list, offset, size):
        """Compute bitwise XOR against a bunch of disks slice"""
        logging.info("Enter xor_blocks(fd_list,%d,%d)"%(offset, size))
        
        if size % 8 != 0:
            raise ValueError('xor_blocks : size must be multiple of 8')
        dt = numpy.dtype('<Q8')

        fd_list[0].seek(offset)
        str_b1=fd_list[0].read(size)
        numpy_b1 = numpy.fromstring(str_b1, dtype=dt)
        all_zero = (numpy.count_nonzero(numpy_b1) == 0 )
        any_zero = all_zero

        for fd in fd_list[1:]:
            fd.seek(offset)
            str_b2=fd.read(size)
            numpy_b2 = numpy.fromstring(str_b2, dtype=dt)
            b2_zero = (numpy.count_nonzero(numpy_b2) == 0 )
            if all_zero == True:
                all_zero = b2_zero
            if any_zero == False:
                any_zero = b2_zero

            numpy.bitwise_xor(numpy_b1,numpy_b2,numpy_b1)

        if all_zero == True:
            result = 'z'
        elif numpy.count_nonzero(numpy_b1) == 0:
            if any_zero:
                result = 'g'
            else:
                result = 'G'
        else:
            result = 'b'

        logging.info("Exit. xor_blocks(fd_list,%d,%d)"%(offset, size))
        #import binascii
        #logging.warn(binascii.hexlify(numpy_b1))
        return (result,numpy_b1)

    def read_disk_xor(self,raid_type,offset,size):
        raid_disks = [ self.disks[i] for i in self.raid_disk_order ]
        return self.xor_blocks(raid_disks,offset,size)[1].tostring()

    def read_disk_parity(self,raid_type,offset,size):
        """Returns textual information about parity status of each sector"""
        logging.warn("Enter read_disk_parity(%s,%d,%d)"%(raid_type,offset,size))
        raid_disks = [ self.disks[i] for i in self.raid_disk_order ]
        start = self.raid_start + offset * self.raid_sector_size / 16
        end = start + size * self.raid_sector_size / 16

        #TODO : improove for nested levels
        if raid_type in ['1','5', '5+0']:
            result = ''.join(
                    [ '0x%011x %c\n'%( addr, self.xor_blocks(raid_disks, addr, self.raid_sector_size)[0])
                            for addr in range(start, end, self.raid_sector_size)
                    ])
        else:
            result = None

        logging.warn("Exit. read_disk_parity(%s,%d,%d)"%(raid_type,offset,size))
        return result


    def read_raid_result(self,raid_type,offset,size):
        """Returns actual RAID data"""
        raid_disks = [ self.disks[i] for i in self.raid_disk_order ] # TODO A garder en attribut ?
        disk_count = len(self.raid_disk_order) # TODO doublon ?

        if raid_type == '0':
            segment_no = offset / self.raid_chunk_size
            segment_off = offset % self.raid_chunk_size 
            stripe_no = segment_no / disk_count
            par_disk = -1
            data_disk = segment_no % disk_count
            off_disk = self.raid_start + stripe_no * self.raid_chunk_size + segment_off
            size2 = min(size, (segment_no+1) * self.raid_chunk_size - offset)

        elif raid_type == '1':
            segment_no = -1
            segment_off = -1
            stripe_no = -1
            par_disk = 1
            data_disk = 0
            off_disk = self.raid_start + offset
            size2 = size

        elif raid_type == '5':
            segment_no = offset / self.raid_chunk_size
            segment_off = offset % self.raid_chunk_size 
            stripe_no = segment_no / (disk_count-1)

            if self.raid_layout in ['ls','la']:
                par_disk = (disk_count-1) - (stripe_no % disk_count)
            else: # self.raid_layout in ['rs','ra']:
                par_disk = stripe_no % disk_count

            if self.raid_layout in ['ls','rs']:
                data_disk = (par_disk+1 + (segment_no % (disk_count-1)) ) % disk_count
            else: # self.raid_layout in ['la','ra']:
                data_disk = segment_no % (disk_count-1)
                if data_disk >= par_disk:
                    data_disk = data_disk + 1

            off_disk = self.raid_start + stripe_no * self.raid_chunk_size + segment_off
            size2 = min(size, (segment_no+1) * self.raid_chunk_size - offset)

        logging.info("raid.read_result(%s): offset=%d,segment_no=%d,segment_off=%d,stripe_no=%d,par_disk=%d,data_disk=%d,off_disk=%d,size2=%d,segment_off+size2=%d" 
        % (raid_type,offset,segment_no,segment_off,stripe_no,par_disk,data_disk,off_disk,size2,segment_off+size2) )

        #TODO recorver from parity if damaged sectors in data_disk
        data_fd = raid_disks[data_disk]
        data_fd.seek(off_disk)
        data = data_fd.read(size2)

        # This kills performance but don't make short reads before EOF
        #if size2 > 0 and size2 < size:
        #    data += self.read_result(self,raid_type,offset+size2,size-size2)

        return data