from structs import *

"""
Probably buggy
HAX, only works on case SENSITIVE
"""

class BTree(object):
    def __init__(self, file, keyStruct, dataStruct):
        self.file = file
        self.keyStruct = keyStruct
        self.dataStruct = dataStruct
        block0 = self.file.readBlock(0)
        btnode = BTNodeDescriptor.parse(block0)
        assert btnode.kind == kBTHeaderNode
        self.header = BTHeaderRec.parse(block0[BTNodeDescriptor.sizeof():])
        assert self.header.keyCompareType == 0 or self.header.keyCompareType == 0 or kHFSBinaryCompare
        #TODO: do more testing when nodeSize != blockSize
        self.nodeSize = self.header.nodeSize
        self.nodesInBlock = file.blockSize / self.header.nodeSize
        self.blocksForNode = self.header.nodeSize / file.blockSize
        #print file.blockSize , self.header.nodeSize
        self.lastRecordNumber = 0
        type, (hdr, maprec) = self.readBtreeNode(0)
        assert len(maprec) == self.nodeSize - 256
        if self.header.totalNodes / 8 > len(maprec):
            pass #TODO: handle map records
        self.maprec = maprec

    def isNodeInUse(self, nodeNumber):
        thisByte = ord(self.maprec[nodeNumber / 8])
        return (thisByte & (1 << (7 - (nodeNumber % 8)))) != 0
    
    def readEmptySpace(self):
        res = ""
        z = 0
        for i in xrange(self.header.totalNodes):
            if not self.isNodeInUse(i):
                z += 1
                res += self.readNode(i)
        assert z == self.header.freeNodes
        return res
    
    #convert construct structure to tuple
    def getComparableKey(self, k):
        raise Exception("implement in subclass")
    
    def compareKeys(self, k1, k2):
        k2 = self.getComparableKey(k2) 
        if k1 == k2:
            return 0
        return -1 if k1 < k2 else 1
    
    def printLeaf(self, key, data):
        print key, data

    def readNode(self, nodeNumber):
        node = ""
        for i in xrange(self.blocksForNode):
            node += self.file.readBlock(nodeNumber * self.blocksForNode + i)
        return node
    
    def readBtreeNode(self, nodeNumber):
        self.lastnodeNumber = nodeNumber
        node = self.readNode(nodeNumber)
        self.lastbtnode = btnode = BTNodeDescriptor.parse(node)

        if btnode.kind == kBTHeaderNode:
            assert btnode.numRecords == 3
            end = self.nodeSize - 8 #2*4
            offsets = Array(btnode.numRecords+1, UBInt16("off")).parse(node[end:])
            assert offsets[-4] == end
            hdr = BTHeaderRec.parse(node[BTNodeDescriptor.sizeof():])
            maprec = node[offsets[-3]:end]
            return kBTHeaderNode, [hdr, maprec]
        elif btnode.kind == kBTIndexNode:
            recs = []
            offsets = Array(btnode.numRecords, UBInt16("off")).parse(node[-2*btnode.numRecords:])
            for i in xrange(btnode.numRecords):
                off = offsets[btnode.numRecords-i-1]
                k = self.keyStruct.parse(node[off:])
                off += 2 + k.keyLength
                k.childNode = UBInt32("nodeNumber").parse(node[off:off+4])
                recs.append(k)
            return kBTIndexNode, recs
        elif btnode.kind == kBTLeafNode:
            recs = []
            offsets = Array(btnode.numRecords, UBInt16("off")).parse(node[-2*btnode.numRecords:])
            for i in xrange(btnode.numRecords):
                off = offsets[btnode.numRecords-i-1]
                k = self.keyStruct.parse(node[off:])
                off += 2 + k.keyLength
                d = self.dataStruct.parse(node[off:])
                recs.append((k,d))
            return kBTLeafNode, recs
        else:
            raise Exception("Invalid node type " + str(btnode)) 

    def search(self, searchKey, node=None):
        if node == None:
            node = self.header.rootNode
            
        type, stuff = self.readBtreeNode(node)
        if len(stuff) == 0:
            return None, None
        
        if type == kBTIndexNode: 
            for i in xrange(len(stuff)):
                if self.compareKeys(searchKey, stuff[i]) < 0:
                    if i > 0:
                        i = i - 1
                    return self.search(searchKey, stuff[i].childNode)
            return self.search(searchKey, stuff[len(stuff)-1].childNode)
        elif type == kBTLeafNode:
            self.lastRecordNumber = 0
            for k,v in stuff:
                res = self.compareKeys(searchKey, k)
                if res == 0:
                    return k, v
                if res < 0:
                    return None, None
                self.lastRecordNumber += 1
        return None, None

    def traverse(self, node=None, count=0, callback=None):
        if node == None:
            node = self.header.rootNode
   
        type, stuff = self.readBtreeNode(node)
        
        if type == kBTIndexNode: 
            for i in xrange(len(stuff)):
                count += self.traverse(stuff[i].childNode, callback=callback)
        elif type == kBTLeafNode:
            for k,v in stuff:
                if callback:
                    callback(k,v)
                else:
                    self.printLeaf(k, v)
                count += 1
        return count
    
    def traverseLeafNodes(self, callback=None):
        nodeNumber = self.header.firstLeafNode
        count = 0
        while nodeNumber != 0:
            _, stuff = self.readBtreeNode(nodeNumber)
            count += len(stuff)
            for k,v in stuff:
                if callback:
                    callback(k,v)
                else:
                    self.printLeaf(k, v)
            nodeNumber = self.lastbtnode.fLink
        return count
    
    #XXX
    def searchMultiple(self, searchKey, filterKeyFunction=lambda x:False):
        self.search(searchKey)
        nodeNumber = self.lastnodeNumber
        recordNumber = self.lastRecordNumber
        kv = []
        while nodeNumber != 0:
            _, stuff = self.readBtreeNode(nodeNumber)
            for k,v in stuff[recordNumber:]:
                if filterKeyFunction(k):
                    kv.append((k,v))
                else:
                    return kv
            nodeNumber = self.lastbtnode.fLink
            recordNumber = 0
        return kv

    def getLBAsHax(self):
        nodes = [self.lastnodeNumber]
        n = self.lastbtnode
        for i in xrange(2):
            nodes.append(self.lastbtnode.bLink)
            self.readBtreeNode(self.lastbtnode.bLink)
        self.lastbtnode = n
        for i in xrange(2):
            nodes.append(self.lastbtnode.fLink)
            self.readBtreeNode(self.lastbtnode.fLink)
        res = []
        for n in nodes:
            res.append(self.file.getLBAforBlock(n * self.blocksForNode))
        return res
        
class CatalogTree(BTree):
    def __init__(self, file, volume):
        super(CatalogTree,self).__init__(file, HFSPlusCatalogKey, HFSPlusCatalogData)
        self.volume = volume
    
    def printLeaf(self, k, d):
        if d.recordType == kHFSPlusFolderRecord or d.recordType == kHFSPlusFileRecord:
            print getString(k)
   
    def getComparableKey(self, k2):
        #XXX http://dubeiko.com/development/FileSystems/HFSPLUS/tn1150.html#StringComparisonAlgorithm
        return (k2.parentID, getString(k2))
    
    def searchByCNID(self, cnid):
        threadk, threadd = self.search((cnid, ""))
        return self.search((threadd.data.parentID, getString(threadd.data))) if threadd else (None, None)
    
    def getFolderContents(self, cnid):
        return self.searchMultiple((cnid, ""), lambda k:k.parentID == cnid)
    
    def getRecordFromPath(self, path):
        if not path.startswith("/"):
            return None, None
        if path == "/":
            return self.searchByCNID(kHFSRootFolderID)
        parentId=kHFSRootFolderID
        i = 1
        k, v = None, None
        for p in path.split("/")[1:]:
            if p == "":
                break
            k,v  = self.search((parentId, p))
            if (k,v) == (None, None):
                return None, None

            if v.recordType == kHFSPlusFolderRecord:
                parentId = v.data.folderID
            elif v.recordType == kHFSPlusFileRecord and is_symlink(v.data):
                linkdata = self.volume.readFileByRecord(v)
                print "symlink %s => %s" % (p, linkdata)
                if not linkdata:
                    return None, None
                t = path.split("/")
                t[i] = linkdata
                newpath = "/".join(t)
                return self.getRecordFromPath(newpath)
            else:
                break
            i += 1
        return k,v
    
class ExtentsOverflowTree(BTree):
    def __init__(self, file):
        super(ExtentsOverflowTree,self).__init__(file, HFSPlusExtentKey, HFSPlusExtentRecord)
    
    def getComparableKey(self, k2):
        return (k2.fileID, k2.forkType, k2.startBlock)
    
    def searchExtents(self, fileID, forkType, startBlock):
        return self.search((fileID, forkType, startBlock))

class AttributesTree(BTree):
    def __init__(self, file):
        super(AttributesTree,self).__init__(file, HFSPlusAttrKey, HFSPlusAttrData)
    
    def printLeaf(self, k, d):
        print k.fileID, getString(k), d.data.encode("hex")
    
    def getComparableKey(self, k2):
        return (k2.fileID, getString(k2))
    
    def searchXattr(self, fileID, name):
        k,v = self.search((fileID, name))
        return v.data if v else None
    
    def getAllXattrs(self, fileID):
        res = {}
        for k,v in self.searchMultiple((fileID, ""), lambda k:k.fileID == fileID):
            res[getString(k)] = v.data
        return res