# -*- coding: utf-8 -*-
import ctypes, io, logging, re, struct, time
import math
import cfg, wbfs2fat, fat

sw2B=struct.Struct(">B")
sw2H=struct.Struct(">H")
sw2I=struct.Struct(">I")
sw2b=struct.Struct(">?")

unit = ["B","KiB","MiB","GiB", "TiB"]

logger = logging.getLogger(__name__)

def hsize(size):
    n=0
    while size>>13:
        size = size>>10
        n+=1
    return "{:>4} {:>3}".format(size, unit[n])


class Wbfs:
    def __init__(self, head, status):
        self.raw = head
        self.status = status
        if self.raw[0:4] != b'WBFS':
            msg = self.status.add(wbfs2fat.Status.err_nomagic)
            logger.warning(msg)
            raise wbfs2fat.WbfsError(msg)
        self.sect_cnt, self.sect_size_p, self.block_size_p, self.version = struct.unpack_from(">IBBBx", self.raw, 4)
        if self.version > 1:
            msg = self.status.add(wbfs2fat.Status.err_unkvers, self.version)
            logger.error(msg)
            raise wbfs2fat.WbfsError(msg)

    def block2adr(self, block):
        return block << self.block_size_p

    def analyse(self,  dev):
        with io.BufferedWriter(io.FileIO((cfg.options.backup_dir or ".")+"/wbfs2fat.debug","wb")) as debug:
            debug.write(self.raw)
        self.block_cnt = self.sect_cnt>>self.block_size_p-self.sect_size_p
        self.clust_size_p = sorted([self.sect_size_p, math.frexp((self.sect_cnt<<self.sect_size_p)/268435444)[1], 15, math.frexp((self.sect_cnt<<self.sect_size_p)/65525)[1]-1, self.sect_size_p+7])[2]
        self.fat_head_sects = fat.Fat.res_sects + (self.sect_cnt-1>>self.clust_size_p-2) + 1
        self.fat_blocks = (self.fat_head_sects-1 >> Game.default_block_size_p-self.sect_size_p) + 1
        self.disc_blocks = 143432*2 >> self.block_size_p - 15
        self.slot_size = ((0x0100 + 2*self.disc_blocks)-1>>self.sect_size_p)+1<<self.sect_size_p
        self.clust_sects_p = self.clust_size_p - self.sect_size_p
        self.block_sects_p = self.block_size_p - self.sect_size_p
        self.block_clusts_p = self.block_size_p - self.clust_size_p
        self.sect_size = 1<<self.sect_size_p
        self.clust_sects = 1<<self.clust_sects_p
        self.block_size = 1<<self.block_size_p
        self.block_sects = 1<<self.block_sects_p
        self.vblock_shrink_p = max(0, self.block_size_p - Game.default_block_size_p)
        self.vblock_size_p = self.block_size_p - self.vblock_shrink_p
        self.vblock_size = 1<<self.vblock_size_p

        used = []
        self.games = []
        for slot in range(self.sect_size-12):
            if sw2b.unpack_from(self.raw,12+slot)[0]:
                try:
                    game = Game(self.raw[self.sect_size+slot*self.slot_size:self.sect_size+(slot+1)*self.slot_size],self, dev)
                    used.extend(game.used)
                    self.games.append(game)
                except Exception:
                    logger.error("Exception caught",exc_info=1)
        self.fat_blocks += 1 + ((len(self.games)+2 << self.clust_size_p) + (len(self.games) << 5) >> self.vblock_size_p)

        used.sort()
        rel=[]
        free=[]
        last=0
        lastgame=""
        err_xlink=""
        err_segv=""
        for (block, game) in used:
            if block >= self.block_cnt << self.vblock_shrink_p:
                err_segv="block "+str(block)+" claimed by "+str(game)+" but only "+str(self.block_cnt)+" blocks in filesystem"
                logger.error(err_segv)
            if block == last:
                err_xlink="block "+str(block)+" claimed by "+str(lastgame)+" and "+str(game)
                logger.warn(err_xlink)
                rel.append((block, game))
            elif block < self.fat_blocks + cfg.options.create_mbr: rel.append((block, game))
            elif block > max(last+1,self.fat_blocks): free.extend(range(max(last+1,self.fat_blocks),block))
            last = block
            lastgame = game
        free.extend(range(last+1, self.block_cnt << self.vblock_shrink_p))
        self.remap = dict(zip(rel, free))
#        print zip(rel, free)
        for game,new in zip(self.games, free[len(rel):]):
            game.block_map[0]=(-1,new)

        logger.info(self.status.add(wbfs2fat.Status.info_wbfs,hsize(self.sect_cnt<<self.sect_size_p),len(self.games)))
        n_blocks = len(rel)+len(self.games)-len(free)
        if n_blocks > 0:
            logger.error(self.status.add(wbfs2fat.Status.err_nofree, hsize(n_blocks<<self.block_size_p), n_blocks, 1<<self.block_size_p%10, unit[self.block_size_p/10]))
        if err_xlink:
            logger.warning(self.status.add(wbfs2fat.Status.err_xlink))
        if err_segv:
            logger.warning(self.status.add(wbfs2fat.Status.err_segv))

#        dev.seek(100+self.sect_cnt << self.sect_size_p)


class Game:
    file_blocks_max = 2047
    default_block_size_p = 21

    def __init__(self, slotdata, wbfs, dev):
        self.wbfs = wbfs
        self.head = slotdata[:0x0080]
        self.id6 = self.head[:0x06].decode("ascii","strict").strip()
        if len(self.id6) != 6 or not self.id6.isalnum():
            raise wbfs2fat.WbfsError("ID6 invalid")
        self.name = self.head[0x20:].decode("latin-1","replace").strip(' \t\r\n\0')
        self.name = re.sub('<','(',self.name)
        self.name = re.sub('>',')',self.name)
        self.name = re.sub('"',"'",self.name)
        self.name = re.sub('\*','~',self.name)
        self.name = re.sub('[?^]',' ',self.name)
        self.name = re.sub('[\/|:]',' - ',self.name)
        self.name = re.sub(' +',' ',self.name)
        self.block_map = [(-1,0)]
        self.used = []
        logger.info("reading "+self.id6)
        for game_block_org in range(wbfs.disc_blocks):
            wbfs_block_org = sw2H.unpack_from(slotdata, 0x0100+2*game_block_org)[0]
            if wbfs_block_org:
                for game_block, wbfs_block in zip(range(game_block_org<<self.wbfs.vblock_shrink_p, game_block_org+1<<self.wbfs.vblock_shrink_p),range(wbfs_block_org<<self.wbfs.vblock_shrink_p, wbfs_block_org+1<<self.wbfs.vblock_shrink_p)):
                    self.block_map.append((game_block,wbfs_block))
                    #logger.debug(str((src,dest)))
        if cfg.options.scrub:
            #bad excuse for real scrubbing but seems to make most files identical to wwt/wit
            tmp = []
            last = 0
            self.block_map.reverse()
            for block in self.block_map[:-1]:
                if block[0] != last-1 or block[0] in [0x7b]: #trial by error
                    dev.seek(block[1]<<self.wbfs.vblock_size_p)
                    data = dev.read(self.wbfs.vblock_size)
                    if not ord(max(data)):
#                        print block[0], "0"
                        continue
#                    print block[0],"x"
                tmp.append(block)
                last = block[0]
            tmp.reverse()
            self.block_map.reverse()
            self.block_map = self.block_map[:1]
            for block in tmp:
                if block[0] != last+1:
                    dev.seek(block[1]<<self.wbfs.vblock_size_p)
                    data = dev.read(self.wbfs.vblock_size)
                    if not ord(max(data)):
#                        print block[0],"0"
                        continue
#                    print block[0],"x"
                self.block_map.append(block)
                last = block[0]
        for block in self.block_map[1:]:
            self.used.append((block[1], self.id6))

    def mkfiles(self, wbfsdir):
        logger.info(self.id6)
        dname = self.name+" ["+self.id6+"]"
        gamedir = wbfsdir.mkdir(self.id6, dname)
        for ext, file_blocks in enumerate(self.block_map[i:i+Game.file_blocks_max] for i in range(0, len(self.block_map), Game.file_blocks_max)):
            #print file_blocks
            logger.debug(str(ext))
            if ext == 0:
                fname = self.id6+".wbfs"
            else:
                fname = self.id6+".wbf"+str(ext)
            extends = []
            for (_,dest) in file_blocks:
                if (dest, self.id6) in self.wbfs.remap: dest=self.wbfs.remap[dest, self.id6]
                adr = (dest-cfg.options.create_mbr) << self.wbfs.vblock_size_p
                extends.append((adr, adr+self.wbfs.vblock_size))
            gamedir.mkfile('{:<8}{:<3}'.format(self.id6[:6], ext), fname, extends)

    def write_head(self, dev):
#        cfg.report(2,"\t"+id6)
        _,new = self.block_map[0]
        buff = ctypes.create_string_buffer(1<<21)
        buff[0:4] = self.wbfs.raw[0:4]
        sw2I.pack_into(buff, 4, len(self.block_map) << self.wbfs.vblock_size_p - self.wbfs.sect_size_p)
        sw2B.pack_into(buff, 8, 9)
        sw2B.pack_into(buff, 9, 21)
        sw2B.pack_into(buff, 10, 1)
        sw2B.pack_into(buff, 12, 0x01)
        buff[self.wbfs.sect_size:self.wbfs.sect_size+len(self.head)]=self.head
        for num, (src, _) in enumerate(self.block_map[1:],1):
            sw2H.pack_into(buff,self.wbfs.sect_size+0x0100+2*src,num)
        #print self.id6, cnt
        freesize = (len(self.block_map)-1 >> 3+self.wbfs.sect_size_p) + 1 << self.wbfs.sect_size_p
        cnt = len(self.block_map)-1
        for adr in range((1<<21)-2*freesize, (1<<21)-freesize, 4):
            if cnt < 0:
                sw2I.pack_into(buff,adr,0xffffffff)
            elif cnt < 32:
                sw2I.pack_into(buff,adr,0xffffffff - (1<<cnt)+1)
            cnt-=32
        dev.seek(new<<self.wbfs.vblock_size_p)
        dev.write(buff)
