from paraspace import dexfile
from paraspace.dex_deptracker import prepare_dep_decls
from paraspace.dex_deptracker import restore_dependencies
import os


def _install_dexfile_4_deptracker():
    global dexfile
    import imp, sys
    from paraspace import dex_deptracker
    from paraspace import injection
    
    try:
        new_dexfile = imp.load_compiled('paraspace.dexfile', dexfile.__file__)
    except ImportError:
        new_dexfile = imp.load_source('paraspace.dexfile', dexfile.__file__)
        pass
    dex_deptracker.dexfile = new_dexfile
    dexfile = new_dexfile
    dex_deptracker._nest_types = (dexfile.array, dexfile.cond, dexfile.switch)

    injection.dex_type_2_array_attr_map = \
        injection._saved_dex_type_2_array_attr_map
    
    sys.modules['paraspace.dexfile'] = new_dexfile
    pass


def update_offset_test():
    from paraspace.dex_deptracker import collect_all_dep_decls
    from paraspace.dex_deptracker import build_dependencies
    from paraspace.dex_deptracker import _install_markers
    from paraspace.dex_deptracker import _patch_dex_type_markers
    from paraspace.dex_deptracker import update_offset

    _install_dexfile_4_deptracker()

    all_dep_decls = collect_all_dep_decls()
    _install_markers(all_dep_decls)
    _patch_dex_type_markers(all_dep_decls)

    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    testdatapath = os.path.join(srcroot, 'data', 'testdata1.dex')
    dex = dexfile.DEXFile.open(testdatapath)

    build_dependencies(dex, all_dep_decls)

    offset0 = dex.typeLists.items[0].value.data_offset
    dex.typeLists.items[0].value.data_offset = 0
    offset1 = dex.typeLists.items[1].value.data_offset
    dex.typeLists.items[1].value.data_offset = 0

    classdata1 = dex.classDatas.items[1]
    fields_methods_szs1 = [c.sizeof(c, [None])
                           for c in classdata1.staticFields.items +
                           classdata1.instanceFields.items +
                           classdata1.directMethods.items +
                           classdata1.virtualMethods.items]
    cdoffset0 = dex.classDatas.items[0].data_offset
    dex.classDatas.items[0].data_offset = 0
    cdoffset1 = dex.classDatas.items[1].data_offset
    dex.classDatas.items[1].data_offset = 0
    cdoffset2 = dex.classDatas.items[2].data_offset
    dex.classDatas.items[2].data_offset = 0
    cdoffset4 = dex.classDatas.items[4].data_offset
    dex.classDatas.items[4].data_offset = 0
    cdoffset8 = dex.classDatas.items[8].data_offset
    dex.classDatas.items[8].data_offset = 0
    cdoffset15 = dex.classDatas.items[15].data_offset
    dex.classDatas.items[15].data_offset = 0
    cdoffset30 = dex.classDatas.items[30].data_offset
    dex.classDatas.items[30].data_offset = 0
    cdoffsetlast = dex.classDatas.items[-1].data_offset
    dex.classDatas.items[-1].data_offset = 0

    dex.typeIds.items[12].data_idx = 0

    update_offset(dex, all_dep_decls)

    classdata1 = dex.classDatas.items[1]
    for c1, c2 in map(None, classdata1.staticFields.items + \
                          classdata1.instanceFields.items + \
                          classdata1.directMethods.items + \
                          classdata1.virtualMethods.items, \
                          fields_methods_szs1):
        assert c1.sizeof(c1, [None]) == c2
        pass
    assert dex.typeLists.items[0].value.data_offset == offset0
    assert dex.typeLists.items[1].value.data_offset == offset1
    assert dex.classDatas.items[0].data_offset == cdoffset0
    assert dex.classDatas.items[1].data_offset == cdoffset1
    assert dex.classDatas.items[2].data_offset == cdoffset2
    assert dex.classDatas.items[4].data_offset == cdoffset4
    assert dex.classDatas.items[8].data_offset == cdoffset8
    assert dex.classDatas.items[15].data_offset == cdoffset15
    assert dex.classDatas.items[30].data_offset == cdoffset30
    assert dex.classDatas.items[-1].data_offset == cdoffsetlast
    assert dex.typeIds.items[12].data_idx == 12
    pass


def restore_dependencies_test():
    _install_dexfile_4_deptracker()
    
    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    datadir = os.path.join(srcroot, 'data')
    helloworld_path = os.path.join(datadir, 'helloworld.dex')
    
    decls = prepare_dep_decls()
    
    hello_dex = dexfile.DEXFile.open(helloworld_path)
    strids_size_saved = hello_dex.header.stringIdsSize
    typeids_off_saved = hello_dex.header.typeIdsOff
    data_off_saved = hello_dex.header.dataOff
    map_off_saved = hello_dex.header.mapOff
    file_sz_saved = hello_dex.header.fileSize

    hello_linked = dexfile.DEXFile_linked.build_dependencies(hello_dex, decls)
    
    first_strid = hello_linked.stringIds.items[0]
    strid = dexfile._DEX_StringId()
    strid.stringDataOff = first_strid.stringDataOff
    hello_linked.stringIds.items.append(strid)

    restore_dependencies(hello_linked, decls)

    assert hello_linked.header.fileSize == (file_sz_saved + 4)
    assert hello_linked.header.stringIdsSize == (strids_size_saved + 1)
    assert hello_linked.header.typeIdsOff == (typeids_off_saved + 4)
    assert hello_linked.header.dataOff == (data_off_saved + 4)
    assert hello_linked.header.mapOff == (map_off_saved + 4)
    pass


def sizeof_after_build_dep_test():
    from paraspace.dex_deptracker import update_offset
    
    _install_dexfile_4_deptracker()

    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    datadir = os.path.join(srcroot, 'data')
    helloworld_path = os.path.join(datadir, 'helloworld.dex')
    
    decls = prepare_dep_decls()
    
    hello_dex = dexfile.DEXFile.open(helloworld_path)
    update_offset(hello_dex, decls)
    
    get_sz = lambda x: x.sizeof(x, [hello_dex])
    assert get_sz(hello_dex.header) == 112
    assert get_sz(hello_dex.stringIds) == 0x168
    assert get_sz(hello_dex.typeIds) == (0x250-0x1d8)
    assert get_sz(hello_dex.protoIds) == (0x31c-0x250)
    assert get_sz(hello_dex.fieldIds) == (0x354-0x31c)
    assert get_sz(hello_dex.methodIds) == (0x46c-0x354)
    assert get_sz(hello_dex.classDefs) == (0x58c-0x46c)
    assert get_sz(hello_dex.annotationSetItems) == (0x5dc-0x58c)
    assert get_sz(hello_dex.codeItems) == (0x8b8-0x5dc)
    assert get_sz(hello_dex.annotationsDirectoryItems) == (0x928-0x8b8)
    assert get_sz(hello_dex.typeLists) == (0x98a-0x928)
    assert get_sz(hello_dex.stringDataItems) == (0xe7d-0x98a)
    assert get_sz(hello_dex.debugInfoItems) == (0xf31-0xe7d)
    assert get_sz(hello_dex.annotationItems) == (0xf82-0xf31)
    assert get_sz(hello_dex.encodedArrayItems) == (0xf9d-0xf82)
    assert get_sz(hello_dex.classDatas) == (0x1024-0xf9d-1)
    assert hello_dex.maps.sizeof(hello_dex.maps, [hello_dex]) == 0xd0

    hello_dex.compute_size()
    unlinked_sz = hello_dex.sizeof()
    unlinked_cdsize = hello_dex.classDatas.sizeof(hello_dex.classDatas, [hello_dex])
    assert unlinked_sz == hello_dex.header.fileSize

    hello_linked = dexfile.DEXFile_linked.build_dependencies(hello_dex, decls)

    hello_linked.compute_size()
    linked_sz = hello_linked.sizeof()
    
    get_sz = lambda x: x.sizeof(x, [hello_linked])
    assert get_sz(hello_linked.header) == 112
    assert get_sz(hello_linked.stringIds) == 0x168
    assert get_sz(hello_linked.typeIds) == (0x250-0x1d8)
    assert get_sz(hello_linked.protoIds) == (0x31c-0x250)
    assert get_sz(hello_linked.fieldIds) == (0x354-0x31c)
    assert get_sz(hello_linked.methodIds) == (0x46c-0x354)
    assert get_sz(hello_linked.classDefs) == (0x58c-0x46c)
    assert get_sz(hello_linked.annotationSetItems) == (0x5dc-0x58c)
    assert get_sz(hello_linked.codeItems) == (0x8b8-0x5dc)
    assert get_sz(hello_linked.annotationsDirectoryItems) == (0x928-0x8b8)
    assert get_sz(hello_linked.typeLists) == (0x98a-0x928)
    assert get_sz(hello_linked.stringDataItems) == (0xe7d-0x98a)
    assert get_sz(hello_linked.debugInfoItems) == (0xf31-0xe7d)
    assert get_sz(hello_linked.annotationItems) == (0xf82-0xf31)
    assert get_sz(hello_linked.encodedArrayItems) == (0xf9d-0xf82)

    linked_cdsize = hello_linked.classDatas.sizeof(hello_linked.classDatas,
                                                   [hello_linked])
    assert linked_cdsize == unlinked_cdsize # test for depend_idx_rel
    assert linked_sz == unlinked_sz
    
    first_strid = hello_linked.stringIds.items[0]
    strid = dexfile._DEX_StringId()
    strid.stringDataOff = first_strid.stringDataOff
    hello_linked.stringIds.items.append(strid)
    
    hello_linked.compute_size()
    linked_sz = hello_linked.sizeof()
    
    assert get_sz(hello_linked.header) == 112
    assert get_sz(hello_linked.stringIds) == 0x16c
    assert get_sz(hello_linked.typeIds) == (0x250-0x1d8)
    assert get_sz(hello_linked.protoIds) == (0x31c-0x250)
    assert get_sz(hello_linked.fieldIds) == (0x354-0x31c)
    assert get_sz(hello_linked.methodIds) == (0x46c-0x354)
    assert get_sz(hello_linked.classDefs) == (0x58c-0x46c)
    assert get_sz(hello_linked.annotationSetItems) == (0x5dc-0x58c)
    assert get_sz(hello_linked.codeItems) == (0x8b8-0x5dc)
    assert get_sz(hello_linked.annotationsDirectoryItems) == (0x928-0x8b8)
    assert get_sz(hello_linked.typeLists) == (0x98a-0x928)
    assert get_sz(hello_linked.stringDataItems) == (0xe7d-0x98a)
    assert get_sz(hello_linked.debugInfoItems) == (0xf31-0xe7d)
    assert get_sz(hello_linked.annotationItems) == (0xf82-0xf31)
    assert get_sz(hello_linked.encodedArrayItems) == (0xf9d-0xf82)

    assert linked_sz == (unlinked_sz + 4)
    pass


def dex_sort_sorted_arrays_test():
    from paraspace.dex_deptracker import dex_sort_sorted_arrays
    
    _install_dexfile_4_deptracker()

    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    datadir = os.path.join(srcroot, 'data')
    helloworld_path = os.path.join(datadir, 'helloworld.dex')
    
    decls = prepare_dep_decls()
    
    hello_dex = dexfile.DEXFile.open(helloworld_path)
    hello_linked = \
        dexfile.DEXFile_linked.build_dependencies(hello_dex, decls)

    first = hello_linked.typeIds.items[0]
    last = hello_linked.typeIds.items[-1]
    hello_linked.typeIds.items[0] = last
    hello_linked.typeIds.items[-1] = first

    typeid10 = hello_linked.typeIds.items[10]
    typeid12 = hello_linked.typeIds.items[12]
    typeid10.descriptorIdx, typeid12.descriptorIdx = \
        typeid12.descriptorIdx, typeid10.descriptorIdx
    
    dex_sort_sorted_arrays(hello_linked)

    assert hello_linked.typeIds.items[0] == first
    assert hello_linked.typeIds.items[-1] == last
    
    assert typeid10 > typeid12
    assert hello_linked.typeIds.items[10] == typeid12
    assert hello_linked.typeIds.items[12] == typeid10

    for protoid0, protoid1 in map(None,
                                  hello_linked.protoIds.items[:-1],
                                  hello_linked.protoIds.items[1:]):
        assert protoid0 < protoid1
        pass
    pass