from paraspace import dexfile
import os

def block_sizes_test():
    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    testdatapath = os.path.join(srcroot, 'data', 'testdata1.dex')
    dex = dexfile.DEXFile.open(testdatapath)
    assert dex.stringIds.data_size == 5384
    assert dex.typeIds.data_size == 704
    assert dex.protoIds.data_size == 1776
    assert dex.fieldIds.data_size == 1944
    assert dex.methodIds.data_size == 3552
    assert dex.classDefs.data_size == 1600
    assert dex.classDatas.data_size == 2052
    assert dex.typeLists.data_size == 838
    assert dex.codeItems.data_size == 30688
    assert dex.annotationSetItems.data_size == 412
    assert dex.annotationsDirectoryItems.data_size == 792
    assert dex.annotationItems.data_size == 642
    assert dex.encodedArrayItems.data_size == 139
    assert dex.debugInfoItems.data_size == 6353
    assert dex.stringDataItems.data_size == 23492
    pass


def dependencies_test():
    from paraspace import dex_deptracker
    
    deps = dex_deptracker.collect_all_dep_decls()
    assert deps['_DEX_AnnotationItem.typeIdx'][0] == dexfile.depend_idx
    assert deps['_DEX_AnnotationItem.typeIdx'][1] == 'DEXFile.typeIds'
    assert deps['_DEX_FieldId.typeIdx'][0] == dexfile.depend_idx
    assert deps['_DEX_FieldId.typeIdx'][1] == 'DEXFile.typeIds'
    assert deps['_DEX_ClassDef.staticValuesOffRef.value'][0] == \
        dexfile.depend_off
    assert deps['_DEX_ClassDef.staticValuesOffRef.value'][1] == \
        '_DEX_EncodedArrayItem'
    assert deps['_DEX_Try.handlerOff'][0] == dexfile.depend_off_rel
    assert deps['_DEX_Try.handlerOff'][1] == '_DEX_Catch'
    assert deps['_DEX_Try.handlerOff'][2] == '_DEX_Code.handlers_size'
    pass


def resolve_name_path_test():
    from paraspace.dex_deptracker import _resolve_name_path
    
    obj, parent = _resolve_name_path('_DEX_ClassData.staticFields.items')
    assert obj == list
    obj, parent = _resolve_name_path('_DEX_ClassData.staticFields.items.*')
    assert obj == dexfile._DEX_Field
    obj, parent = _resolve_name_path('_DEX_ClassData.staticFields.items.1')
    assert obj == dexfile._DEX_Field
    
    obj, parent = _resolve_name_path('_DEX_Catch.catchAllHandler.value')
    assert obj == dexfile._DEX_CatchAllHandler
    
    key = dexfile._DEX_AnnotationMember_noname.kDexAnnotationNull
    obj, parent = _resolve_name_path('_DEX_AnnotationMember_noname.value.' + 
                                     repr(key))
    assert isinstance(obj, dexfile.abs_value)

    obj, parent = _resolve_name_path('_DEX_AnnotationSetItem.annotationOffs.'
                                     'items.0')
    assert isinstance(obj, dexfile.depend_off)
    pass


def find_dep_decls_from_clazz_test():
    from paraspace.dex_deptracker import _find_dep_decls_from_clazz
    
    dex_types = dict([(dex_type_name, getattr(dexfile, dex_type_name))
                      for dex_type_name in dir(dexfile)
                      if dex_type_name.startswith('_DEX_')])
    deps = _find_dep_decls_from_clazz('_DEX_ProtoId', dexfile._DEX_ProtoId,
                                      dex_types)
    assert len(deps) == 3
    assert '_DEX_ProtoId.shortyIdx' in deps
    assert deps['_DEX_ProtoId.shortyIdx'][0] == dexfile.depend_idx
    assert deps['_DEX_ProtoId.shortyIdx'][1] == 'DEXFile.stringIds'

    assert deps['_DEX_ProtoId.parametersOffRef.value'][1] == '_DEX_TypeList'
    pass


def find_dep_decls_from_clazz__array_test():
    from paraspace.dex_deptracker import _find_dep_decls_from_clazz

    dex_types = dict([(dex_type_name, getattr(dexfile, dex_type_name))
                      for dex_type_name in dir(dexfile)
                      if dex_type_name.startswith('_DEX_')])
    deps = _find_dep_decls_from_clazz('_DEX_AnnotationSetItem',
                                      dexfile._DEX_AnnotationSetItem,
                                      dex_types)
    assert len(deps) == 1
    name_path = '_DEX_AnnotationSetItem.annotationOffs.items.*'
    assert name_path in deps
    assert deps[name_path][0] == dexfile.depend_off
    assert deps[name_path][1] == '_DEX_AnnotationItem'
    pass


def travel_dex_relocatable__array_test():
    from paraspace.dex_deptracker import _travel_dex_relocatable

    srcdir = os.path.dirname(__file__)
    srcroot = os.path.join(srcdir, '..', '..')
    testdatapath = os.path.join(srcroot, 'data', 'testdata1.dex')
    dex = dexfile.DEXFile.open(testdatapath)
    dexroot = dex.typeLists.items[0].value

    itr = _travel_dex_relocatable(dexroot)
    pathes = [v[2] for v in itr]
    assert len(pathes) == 5
    assert '_DEX_TypeList' in pathes
    assert '_DEX_TypeList.num' in pathes
    assert '_DEX_TypeList.typeItems' in pathes
    assert '_DEX_TypeList.typeItems.items' in pathes
    assert '_DEX_TypeList.typeItems.items.0' in pathes
    pass


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


def travel_dex_type_test():
    from paraspace.dex_deptracker import collect_all_dep_decls
    from paraspace.dex_deptracker import _install_markers
    from paraspace.dex_deptracker import _travel_dex_type

    _install_dexfile_4_deptracker()
    
    attr_infos = [attr_info
                  for attr_info in \
                      _travel_dex_type(dexfile._DEX_AnnotationsDirectoryItem,
                                       '_DEX_AnnotationsDirectoryItem')
                  ]
    attr_names = set([attr_info[1]
                      for attr_info in attr_infos])
    assert '_DEX_AnnotationsDirectoryItem.classAnnotationsOff' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.fieldsSize' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.methodsSize' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.parametersSize' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.fieldAnnotationsItems' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.fieldAnnotationsItems.items.*' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.methodAnnotationsItems' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.parameterAnnotationsItems' in attr_names
    
    all_dep_decls = collect_all_dep_decls()
    _install_markers(all_dep_decls)
    
    attr_infos = [attr_info
                  for attr_info in \
                      _travel_dex_type(dexfile._DEX_AnnotationsDirectoryItem,
                                       '_DEX_AnnotationsDirectoryItem')
                  ]
    attr_names = set([attr_info[1]
                      for attr_info in attr_infos])
    assert '_DEX_AnnotationsDirectoryItem.classAnnotationsOff' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.fieldsSize' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.methodsSize' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.parametersSize' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.fieldAnnotationsItems' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.fieldAnnotationsItems.items.*' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.methodAnnotationsItems' in attr_names
    assert '_DEX_AnnotationsDirectoryItem.parameterAnnotationsItems' in attr_names
    pass


def install_markers_test():
    from paraspace.dex_deptracker import collect_all_dep_decls
    from paraspace.dex_deptracker import _install_markers, _idx_marker
    from paraspace.dex_deptracker import _marker
    from paraspace.dex_deptracker import _offset_marker, _rel_offset_marker
    from paraspace.dex_deptracker import _patch_dex_type_markers

    _install_dexfile_4_deptracker()

    all_dep_decls = collect_all_dep_decls()
    _install_markers(all_dep_decls)
    assert isinstance(dexfile.DEXFile.typeIds, _idx_marker)
    assert isinstance(dexfile._DEX_Code.back_type.handlers_size,
                      _rel_offset_marker)
    assert dexfile._DEX_TypeList.__class__ == _offset_marker

    _patch_dex_type_markers(all_dep_decls)
    assert isinstance(dexfile._DEX_TypeList, _offset_marker)
    assert isinstance(dexfile.DEXFile.typeLists.child_type.value,
                      _offset_marker)
    assert isinstance(dexfile._DEX_Catch, _marker)
    pass


def link_dependencies_test():
    from paraspace.dex_deptracker import collect_all_dep_decls
    from paraspace.dex_deptracker import build_dependencies
    from paraspace.dex_deptracker import _install_markers, _idx_marker
    from paraspace.dex_deptracker import _offset_marker, _rel_offset_marker
    from paraspace.dex_deptracker import _patch_dex_type_markers

    _install_dexfile_4_deptracker()

    all_dep_decls = collect_all_dep_decls()
    _install_markers(all_dep_decls)
    assert isinstance(dexfile.DEXFile.typeIds, _idx_marker)
    _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)
    
    code_item = dex.codeItems.items[0]
    assert code_item.debugInfoOff.__class__.__name__ == '_DEX_DebugInfoItem'

    stringid = dex.stringIds.items[0]
    assert isinstance(stringid.stringDataOff, dexfile._DEX_StringDataItem)
    pass


def link_dependencies_cond_value_ref_test():
    from paraspace.dex_deptracker import collect_all_dep_decls
    from paraspace.dex_deptracker import build_dependencies
    from paraspace.dex_deptracker import _install_markers, _idx_marker
    from paraspace.dex_deptracker import _patch_dex_type_markers

    _install_dexfile_4_deptracker()

    all_dep_decls = collect_all_dep_decls()
    _install_markers(all_dep_decls)
    assert isinstance(dexfile.DEXFile.typeIds, _idx_marker)
    _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)
    
    code_item = dex.codeItems.items[0]
    assert code_item.debugInfoOff.__class__.__name__ == '_DEX_DebugInfoItem'
    
    anno_dir_item = dex.annotationsDirectoryItems.items[0]
    assert isinstance(anno_dir_item.classAnnotationsOffRef.value,
                      dexfile._DEX_AnnotationSetItem)
    
    anno_dir_item = dex.annotationsDirectoryItems.items[1]
    assert isinstance(anno_dir_item.classAnnotationsOffRef.value,
                      dexfile._DEX_AnnotationSetItem)

    clazz_def = dex.classDefs.items[1]
    assert isinstance(clazz_def.annotationsOffRef.value,
                      dexfile._DEX_AnnotationsDirectoryItem)
    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
    
    update_offset(dex, all_dep_decls)

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