import struct
import io
from .leb128 import read_uleb128_fast
import bisect

_STRUCT_I = struct.Struct('<I')
_STRUCT_H = struct.Struct('<H')
_STRUCT_III = struct.Struct('<III')
_STRUCT_HHI = struct.Struct('<HHI')
_STRUCT_HHHHII = struct.Struct('<HHHHII')

def _decode_mutf8_fallback(raw : bytes) -> str:
    raw = raw.replace(b'\xc0\x80', b'\x00')
    try:
        return (
            raw.decode('utf-8', errors='surrogatepass')
            .encode('utf-16-le', errors='surrogatepass')
            .decode('utf-16-le', errors='replace')
        )
    except (UnicodeDecodeError, UnicodeEncodeError):
        return raw.decode('utf-8', errors='replace')


class DEXHeader:
    def __init__(self, buf):
        # (off, size)
        # 0x3c offset == ids_off, 0x38 offset == ids_size
        self.strings = (_STRUCT_I.unpack_from(buf, 0x3C)[0], _STRUCT_I.unpack_from(buf, 0x38)[0])
        # 0x44 offset == ids_off, 0x40 offset == ids_size
        self.types = (_STRUCT_I.unpack_from(buf, 0x44)[0], _STRUCT_I.unpack_from(buf, 0x40)[0])
        self.prototypes = (_STRUCT_I.unpack_from(buf, 0x4C)[0], _STRUCT_I.unpack_from(buf, 0x48)[0])
        self.fields = (_STRUCT_I.unpack_from(buf, 0x54)[0], _STRUCT_I.unpack_from(buf, 0x50)[0])
        self.methods = (_STRUCT_I.unpack_from(buf, 0x5C)[0], _STRUCT_I.unpack_from(buf, 0x58)[0])
        self.classes = (_STRUCT_I.unpack_from(buf, 0x64)[0], _STRUCT_I.unpack_from(buf, 0x60)[0])
        self.mapoff = _STRUCT_I.unpack_from(buf, 0x34)[0]

class PrimitiveTypes:
    VOID_T = 0
    BOOLEAN = 1
    BYTE = 2
    SHORT = 3
    CHAR = 4
    INT = 5
    LONG = 6
    FLOAT = 7
    DOUBLE = 8

class TypeTypes:
    PRIMITIVE = 1
    CLASS = 2
    ARRAY = 3

class Type:
    PRIMITIVES = PrimitiveTypes
    TYPES = TypeTypes

    def __init__(self, descriptor):
        self.descriptor = descriptor
        self._dim = 0
        self._underlying_array_type = None
        self._value = None
        self._type = None
        self._parsed = False

    def _parse(self):
        if self._parsed:
            return
        self._parsed = True
        desc = self.descriptor
        while desc.startswith('['):
            self._dim += 1
            desc = desc[1:]
        
        if self._dim > 0:
            self._type = self.TYPES.ARRAY
            self._underlying_array_type = Type(desc)
            return

        if desc == 'V':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.VOID_T
        elif desc == 'Z':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.BOOLEAN
        elif desc == 'B':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.BYTE
        elif desc == 'S':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.SHORT
        elif desc == 'C':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.CHAR
        elif desc == 'I':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.INT
        elif desc == 'J':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.LONG
        elif desc == 'F':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.FLOAT
        elif desc == 'D':
            self._type = self.TYPES.PRIMITIVE
            self._value = self.PRIMITIVES.DOUBLE
        else:
            self._type = self.TYPES.CLASS
            self._value = desc

    @property
    def dim(self):
        self._parse()
        return self._dim

    @property
    def underlying_array_type(self):
        self._parse()
        return self._underlying_array_type

    @property
    def value(self):
        self._parse()
        return self._value

    @property
    def type(self):
        self._parse()
        return self._type

    def __str__(self):
        return self.descriptor

class DexString:
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return self.value

class DexPrototype:
    def __init__(self, dex, proto_idx):
        offset = dex.header.prototypes[0] + proto_idx * 12
        self.shorty_idx, self.return_type_idx, self.parameters_off = _STRUCT_III.unpack_from(dex.buf, offset)
        self.dex = dex

    @property
    def parameters_type(self):
        if not hasattr(self, '_parameters_type'):
            if self.parameters_off == 0:
                self._parameters_type = []
            else:
                size = _STRUCT_I.unpack_from(self.dex.buf, self.parameters_off)[0]
                params = []
                off = self.parameters_off + 4
                for _ in range(size):
                    type_idx = _STRUCT_H.unpack_from(self.dex.buf, off)[0]
                    off += 2
                    params.append(self.dex.get_type(type_idx))
                self._parameters_type = params
        return self._parameters_type

class DexField:
    def __init__(self, dex, field_idx):
        self.index = field_idx
        self.dex = dex
        offset = dex.header.fields[0] + field_idx * 8
        self.class_idx, self.type_idx, self.name_idx = _STRUCT_HHI.unpack_from(dex.buf, offset)
        self.access_flags = 0
        self.is_static = False

    @property
    def cls(self):
        if not hasattr(self, '_cls'):
            class Cls:
                def __init__(self, fullname):
                    self.fullname = fullname
            self._cls = Cls(self.dex.get_type(self.class_idx).descriptor)
        return self._cls

    @property
    def type(self):
        return self.dex.get_type(self.type_idx)

    @property
    def name(self):
        return self.dex.strings[self.name_idx]

class DexMethod:
    def __init__(self, dex, method_idx):
        self.index = method_idx
        self.dex = dex
        offset = dex.header.methods[0] + method_idx * 8
        self.class_idx, self.proto_idx, self.name_idx = _STRUCT_HHI.unpack_from(dex.buf, offset)
        self._bytecode = None
        self._code_off = 0
        self.access_flags = 0
        self.is_direct = False
        self.is_virtual = False

    @property
    def cls(self):
        if not hasattr(self, '_cls'):
            class Cls:
                def __init__(self, fullname):
                    self.fullname = fullname
            self._cls = Cls(self.dex.get_type(self.class_idx).descriptor)
        return self._cls

    @property
    def name(self):
        return self.dex.strings[self.name_idx]

    @property
    def prototype(self):
        if not hasattr(self, '_prototype'):
            self._prototype = self.dex.get_prototype(self.proto_idx)
        return self._prototype

    @property
    def code_offset(self):
        return self._code_off

    @property
    def bytecode(self):
        if self._bytecode is not None:
            return self._bytecode
        if self._code_off == 0:
            self._bytecode = []
            return self._bytecode
        
        # Parse code_item
        off = self._code_off
        registers_size, ins_size, outs_size, tries_size, debug_info_off, insns_size = _STRUCT_HHHHII.unpack_from(self.dex.buf, off)
        off += 16
        
        insns_bytes = self.dex.buf[off : off + insns_size * 2]
        self._bytecode = list(insns_bytes)
        return self._bytecode

class DexClass:
    def __init__(self, dex, class_idx, class_def_off, class_index):
        self.dex = dex
        self.class_idx = class_idx
        self.index = class_index
        self._class_def_off = class_def_off
        
        self.class_data_off = _STRUCT_I.unpack_from(dex.buf, class_def_off + 24)[0]
        
        self._methods = []
        self._fields = []
        self._parsed = False

    @property
    def fullname(self):
        if not hasattr(self, '_fullname'):
            self._fullname = self.dex.get_type(self.class_idx).descriptor
        return self._fullname

    def _parse_class_data(self):
        if self._parsed:
            return
        self._parsed = True
        
        if self.class_data_off == 0:
            return
            
        data = self.dex.buf
        pos = self.class_data_off
        static_fields_size, c = read_uleb128_fast(data, pos); pos += c
        instance_fields_size, c = read_uleb128_fast(data, pos); pos += c
        direct_methods_size, c = read_uleb128_fast(data, pos); pos += c
        virtual_methods_size, c = read_uleb128_fast(data, pos); pos += c
        
        field_idx = 0
        for _ in range(static_fields_size):
            field_idx_diff, c = read_uleb128_fast(data, pos); pos += c
            field_idx += field_idx_diff
            access_flags, c = read_uleb128_fast(data, pos); pos += c
            f = self.dex.get_field(field_idx)
            f.access_flags = access_flags
            f.is_static = True
            self._fields.append(f)
            
        field_idx = 0
        for _ in range(instance_fields_size):
            field_idx_diff, c = read_uleb128_fast(data, pos); pos += c
            field_idx += field_idx_diff
            access_flags, c = read_uleb128_fast(data, pos); pos += c
            f = self.dex.get_field(field_idx)
            f.access_flags = access_flags
            f.is_static = False
            self._fields.append(f)
            
        method_idx = 0
        for _ in range(direct_methods_size):
            method_idx_diff, c = read_uleb128_fast(data, pos); pos += c
            method_idx += method_idx_diff
            access_flags, c = read_uleb128_fast(data, pos); pos += c
            code_off, c = read_uleb128_fast(data, pos); pos += c
            
            m = self.dex.get_method(method_idx)
            m._code_off = code_off
            m.access_flags = access_flags
            m.is_direct = True
            self._methods.append(m)
            
        method_idx = 0
        for _ in range(virtual_methods_size):
            method_idx_diff, c = read_uleb128_fast(data, pos); pos += c
            method_idx += method_idx_diff
            access_flags, c = read_uleb128_fast(data, pos); pos += c
            code_off, c = read_uleb128_fast(data, pos); pos += c
            
            m = self.dex.get_method(method_idx)
            m._code_off = code_off
            m.access_flags = access_flags
            m.is_direct = False
            m.is_virtual = True
            self._methods.append(m)

    @property
    def methods(self):
        self._parse_class_data()
        return self._methods

    @property
    def fields(self):
        self._parse_class_data()
        return self._fields

    @property
    def superclass_idx(self):
        if not hasattr(self, '_superclass_idx'):
            self._superclass_idx = _STRUCT_I.unpack_from(self.dex.buf, self._class_def_off + 8)[0]
        return self._superclass_idx

    @property
    def interface_type_idxs(self) -> list:
        if not hasattr(self, '_interface_type_idxs'):
            interfaces_off = _STRUCT_I.unpack_from(self.dex.buf, self._class_def_off + 12)[0]
            if interfaces_off == 0:
                self._interface_type_idxs = []
            else:
                size = _STRUCT_I.unpack_from(self.dex.buf, interfaces_off)[0]
                idxs = []
                off = interfaces_off + 4
                for _ in range(size):
                    idxs.append(_STRUCT_H.unpack_from(self.dex.buf, off)[0])
                    off += 2
                self._interface_type_idxs = idxs
        return self._interface_type_idxs

class DEX:
    @staticmethod
    def parse(buf, name=""):
        return DEX(buf, name)

    def __init__(self, buf, name=""):
        self.buf = memoryview(buf)
        self.name = name
        self.header = DEXHeader(self.buf)
        
        self._strings = None
        self._types = None
        self._prototypes = None
        self._methods = None
        self._fields = None
        self._classes = {}
        self._classes_by_name = {}

        self._strings_proxy = self.StringsProxy(self)
        self._types_proxy = self.TypesProxy(self)
        self._methods_proxy = self.MethodsProxy(self)
        self._fields_proxy = self.FieldsProxy(self)
        self._classes_proxy = self.ClassesProxy(self)

    # lazy parse
    def get_string(self, str_idx):
        if self._strings is None:
            self._strings = {}

        if str_idx in self._strings:
            return self._strings[str_idx]

        str_idx_off = self.header.strings[0]
        string_off = _STRUCT_I.unpack_from(self.buf, str_idx_off + str_idx * 4)[0]
        utf16_size, c = read_uleb128_fast(self.buf, string_off)
        data_start = string_off + c
        end = data_start + utf16_size
        # ASCII keeps the old O(1) boundary calculation; MUTF-8 needs its terminator.
        # refer to https://github.com/MG1937/ASC/issues/33
        # 20260923 DEX.get_string slices utf16_size bytes out of a string_data_item, but that field counts UTF-16 code units
        if end >= len(self.buf) or self.buf[end] != 0:
            raw_buf = self.buf.obj
            if hasattr(raw_buf, 'find') and len(raw_buf) == len(self.buf):
                end = raw_buf.find(b'\x00', data_start, len(self.buf))
            else:
                end = data_start
                while end < len(self.buf) and self.buf[end] != 0:
                    end += 1
                if end == len(self.buf):
                    end = -1
            if end < 0:
                raise ValueError("unterminated string_data_item")

        raw = bytes(self.buf[data_start:end])
        try:
            s = raw.decode('utf-8')
        except UnicodeDecodeError:
            s = _decode_mutf8_fallback(raw)
        self._strings[str_idx] = s
        return s

#    def _init_string_offsets(self):
#        if not hasattr(self, '_string_offsets') or self._string_offsets is None:
#            off = self.header.strings[0]
#            size = self.header.strings[1]
#            if size > 0:
#                self._string_offsets = struct.unpack_from(f'<{size}I', self.buf, off)
#            else:
#                self._string_offsets = ()

    class StringsProxy:
        def __init__(self, dex):
            self.dex = dex
        def __getitem__(self, idx):
            return self.dex.get_string(idx)
        def __len__(self):
            return self.dex.header.strings[1]

    @property
    def strings(self):
        return self._strings_proxy

    def get_type(self, type_idx):
        if self._types is None:
            self._types = {}
        
        if type_idx not in self._types:
            off = self.header.types[0]
            type_off = off + type_idx * 4
            str_idx = _STRUCT_I.unpack_from(self.buf, type_off)[0]
            self._types[type_idx] = Type(self.strings[str_idx])
            
        return self._types[type_idx]

    def get_prototype(self, proto_idx):
        if self._prototypes is None:
            self._prototypes = {}
        if proto_idx not in self._prototypes:
            self._prototypes[proto_idx] = DexPrototype(self, proto_idx)
        return self._prototypes[proto_idx]

    class TypesProxy:
        def __init__(self, dex):
            self.dex = dex
        def __getitem__(self, idx):
            return self.dex.get_type(idx)
        def __len__(self):
            return self.dex.header.types[1]

    @property
    def types(self):
        return self._types_proxy

    def get_method(self, method_idx):
        if self._methods is None:
            self._methods = {}
        if method_idx not in self._methods:
            self._methods[method_idx] = DexMethod(self, method_idx)
        return self._methods[method_idx]

    def get_field(self, field_idx):
        if self._fields is None:
            self._fields = {}
        if field_idx not in self._fields:
            self._fields[field_idx] = DexField(self, field_idx)
        return self._fields[field_idx]

        # We simulate the list access if someone does dex.methods[idx]
    class MethodsProxy:
        def __init__(self, dex):
            self.dex = dex
        def __getitem__(self, idx):
            return self.dex.get_method(idx)
        def __len__(self):
            return self.dex.header.methods[1]

    @property
    def methods(self):
        return self._methods_proxy

    class FieldsProxy:
        def __init__(self, dex):
            self.dex = dex
        def __getitem__(self, idx):
            return self.dex.get_field(idx)
        def __len__(self):
            return self.dex.header.fields[1]

    @property
    def fields(self):
        return self._fields_proxy

    class ClassesProxy:
        def __init__(self, dex):
            self.dex = dex
            self.size = dex.header.classes[1]
            self.off = dex.header.classes[0]
        def __getitem__(self, idx):
            if idx in self.dex._classes:
                return self.dex._classes[idx]
            class_def_off = self.off + idx * 32
            class_idx = _STRUCT_I.unpack_from(self.dex.buf, class_def_off)[0]
            class_obj = DexClass(self.dex, class_idx, class_def_off, idx)
            self.dex._classes[idx] = class_obj
            return class_obj
        def __len__(self):
            return self.size

    @property
    def classes(self):
        return self._classes_proxy

    """
    # still not lazy
    @property
    def classes(self):
        if self._classes is None:
            self._classes = []
            off = self.header.classes[0]
            size = self.header.classes[1]
            
            for i in range(size):
                class_def_off = off + i * 32
                class_idx = _STRUCT_I.unpack_from(self.buf, class_def_off)[0]
                self._classes.append(DexClass(self, class_idx, class_def_off, i))
        return self._classes
    """

    # only return uleb128 prefix byte, only fit class string
    # we dont care other situation!!!!
    def _get_uleb128_prefix(self, lens : int):
        if lens < 128:
            return bytes([lens])
        else:
            # class len never longer then 0x807f!!
            return bytes([lens | 0x80, lens >> 7])

    def get_class(self, fullname):
        old = self._classes_by_name.get(fullname)
        if old is not None:
            return old
        raw_bytes = self.buf.obj if isinstance(self.buf, memoryview) else self.buf

        off = self.header.classes[0]
        size = self.header.classes[1]
        type_ids_off = self.header.types[0]
        type_ids_size = self.header.types[1]
        type_idx = -1

        left, right = 0, type_ids_size - 1
        while left <= right:
            mid = (left + right) // 2
            desc_idx = _STRUCT_I.unpack_from(raw_bytes, type_ids_off + mid * 0x4)[0]
            string = self.get_string(desc_idx)
            if string == fullname:
                type_idx = mid
                break
            elif string < fullname:
                left = mid + 1
            else:
                right = mid - 1
        
        if type_idx == -1:
            # DEX sorts by UTF-16 code units, while Python compares code points.
            for mid in range(type_ids_size):
                desc_idx = _STRUCT_I.unpack_from(raw_bytes, type_ids_off + mid * 0x4)[0]
                if self.get_string(desc_idx) == fullname:
                    type_idx = mid
                    break

        if type_idx == -1:
            return None

        for class_idx in range(size):
            class_def_off = off + class_idx * 0x20
            if _STRUCT_I.unpack_from(raw_bytes, class_def_off)[0] != type_idx:
                continue
            clazz = self.classes[class_idx]
            self._classes_by_name[fullname] = clazz
            return clazz
        return None
        """
        if self._classes is None:
            self.classes
        return self._classes[class_idx]
        """