from typing import Optional, Tuple

from mtkclient.Library.Exploit.exptools.arch import ArchTools

class Aarch64Tools(ArchTools):
    def force_return(self, data: bytearray, offset: int, value: int, thumb: bool = False) -> bytearray:
        _ = thumb
        mov_x0 = 0xD2800000 | ((value & 0xFFFF) << 5)
        ret = 0xD65F03C0
        data[offset:offset + 4] = mov_x0.to_bytes(4, 'little')
        data[offset+4:offset + 8] = ret.to_bytes(4, 'little')
        return data

    def decode_movt(self, instr: int) -> Optional[Tuple[int, int]]:
        """
        Decode MOVT (move top): rd = (rd & 0xFFFF) | (imm16 << 16)
        Opcode mask: 0x0FF00000 == 0x03400000
        """
        if (instr & 0x0FF00000) == 0x03400000:
            rd = (instr >> 12) & 0xF
            imm4 = (instr >> 16) & 0xF
            imm12 = instr & 0xFFF
            imm16 = (imm4 << 12) | imm12
            return rd, imm16
        return None

    def decode_movw(self, instr: int) -> Optional[Tuple[int, int]]:
        """
        Decode MOVW (move wide): rd = imm16
        Opcode mask: 0x0FF00000 == 0x03000000
        """
        if (instr & 0x0FF00000) == 0x03000000:
            rd = (instr >> 12) & 0xF
            imm4 = (instr >> 16) & 0xF
            imm12 = instr & 0xFFF
            imm16 = (imm4 << 12) | imm12
            return rd, imm16
        return None

    def find_function_from_string(self, s: str) -> Optional[int]:
        xref = self.find_string_xref(s)
        if xref is None:
            return None
        return self.find_function_start_from_off(xref)

    def decode_adrp(self, instr: int, pc: int) -> Optional[Tuple[int, int]]:
        if (instr & 0x9F000000) != 0x90000000:
            return None
        rd = instr & 0x1F
        immlo = (instr >> 29) & 0x3
        immhi = (instr >> 5) & 0x7FFFF
        imm = (immhi << 2) | immlo
        if imm & 0x100000 != 0:
            imm -= 0x200000
        page = pc & ~0xFFF
        return page + (imm << 12), rd

    def decode_add_imm(self, instr: int) -> Optional[Tuple[int, int, int]]:
        if (instr & 0xFF000000) != 0x91000000:
            return None
        rn = (instr >> 5) & 0x1F
        rd = instr & 0x1F
        imm12 = (instr >> 10) & 0xFFF
        shift = (instr >> 22) & 1
        imm = imm12 << 12 if shift != 0 else imm12
        return rn, rd, imm

    def decode_bl(self, instr: int, pc: int) -> Optional[int]:
        opcode = instr & 0xFC000000
        if opcode not in (0x94000000,0x14000000):
            return None
        off = instr & 0x03FFFFFF
        if off & 0x02000000 != 0:
            off -= 0x04000000
        return pc + (off * 4)

    def decode_mov_register(self, instr: int) -> Optional[Tuple[int, int]]:
        """ Decode mov register and return its value """
        if (instr & 0xFFE0FFE0) != 0xAA0003E0:
            return None
        return (instr >> 16) & 0x1F, instr & 0x1F

    def find_string(self, s: str) -> Optional[int]:
        target_bytes = s.encode('utf-8')
        with_null = target_bytes + b'\x00'
        pos = self.data.find(with_null)
        if pos != -1:
            return pos
        pos = self.data.find(target_bytes)
        if pos != -1:
            return pos
        return None

    def find_string_xref(self, s: str) -> Optional[int]:
        off = self.find_string(s)
        if off is None:
            return None
        va = self.base_addr + off
        page = va & ~0xFFF
        for off in range(0, len(self.data), 4):
            instr = self.read_u32(off)
            if instr is None:
                continue
            pc = self.base_addr + off
            adrp = self.decode_adrp(instr, pc)
            if adrp is None:
                continue
            adrp_page, reg = adrp
            if adrp_page != page:
                continue
            add_off = self.find_matching_add(off, reg, va)
            if add_off is not None:
                return add_off
        return None

    def find_matching_add(self, base_off: int, reg: int, va: int) -> Optional[int]:
        for off in range(base_off + 4, base_off + 64, 4):
            instr = self.read_u32(off)
            if instr is None:
                continue
            add = self.decode_add_imm(instr)
            if add is None:
                continue
            rn, rd, imm = add
            _ = rd
            if rn == reg and (va & ~0xFFF) + imm == va:
                return off
        return None

    def find_function_start_from_off(self, from_off: int) -> Optional[int]:
        mask = 0xFFC07FFF
        pattern = 0xA9807BFD
        mask_fallback = 0xFFC003FF
        pattern_fallback = 0xA9007BFD
        limit = 0x5000
        end = max(0, from_off - limit)
        for off in range(from_off, end, -4):
            instr = self.read_u32(off)
            if instr is None:
                continue
            if off >= 4:
                prev = self.read_u32(off - 4)
                if prev is None:
                    continue
                # Pointer authentication
                if prev in (0xD503233F,0xD503237F):
                    return off - 4
                # Sub SP before STP
                if ((prev & mask_fallback) == 0xD10003FF and
                        (instr & mask_fallback) == pattern_fallback & mask_fallback):
                    return off - 4
            if (instr & mask) == pattern:
                return off
        return None

    def decode_sub_reg(self, instr: int) -> Optional[Tuple[int, int, int]]:
        """ SUB (register, no shift/extend) 64-bit: 0xCB000000 | (rm << 16) | (rn << 5) | rd """
        if (instr & 0xFF200000) == 0xCB000000:
            rm = (instr >> 16) & 0x1F
            rn = (instr >> 5) & 0x1F
            rd = instr & 0x1F
            return rn, rm, rd  # rn - rm -> rd
        return None

    def decode_sub_imm(self, instr: int) -> Optional[Tuple[int, int, int]]:
        """ SUB (immediate): 0xD1000000 mask """
        if (instr & 0xFF000000) == 0xD1000000:
            rn = (instr >> 5) & 0x1F
            rd = instr & 0x1F
            imm12 = (instr >> 10) & 0xFFF
            shift = (instr >> 22) & 0x1  # 1 = <<12, 0 = no shift
            imm = imm12 << (12 if shift else 0)
            return rn, rd, imm
        return None

    def resolve_register_value_back(self, at: int, reg: int, lookback: int) -> Optional[int]:
        start = max(0, at - lookback * 4)
        registers = [None] * 32  # Track X0-X31
        current_off=at
        while current_off>start:
            instr = self.read_u32(current_off)
            if instr is None:
                current_off -= 4
                continue
            pc = self.offset_to_va(current_off)

            # ADRP Xd, page
            adrp = self.decode_adrp(instr, pc)
            if adrp is not None:
                page, rd = adrp
                if registers[rd] is None:
                    registers[rd] = page
                else:
                    registers[rd] = registers[rd] + page
                current_off -= 4
                continue

            # ADD Xd, Xn, imm
            add = self.decode_add_imm(instr)
            if add is not None:
                rn, rd, imm = add
                if registers[rn] is not None:
                    if registers[rd] is None:
                        registers[rd] = 0
                    registers[rd] = registers[rn] + imm
                elif registers[rd] is None:
                    registers[rd] = imm
                current_off -= 4
                continue

            # SUB Xd, Xn, Xm  (register subtract)
            sub_reg = self.decode_sub_reg(instr)
            if sub_reg is not None:
                rn, rm, rd = sub_reg
                if registers[rn] is not None and registers[rm] is not None:
                    registers[rd] = registers[rn] - registers[rm]
                current_off -= 4
                continue
            current_off -= 4
        return registers[reg]

    def resolve_register_value(self, at: int, reg: int, lookahead: int) -> Optional[int]:
        """
        Resolve a register value by simulating forward from the call target (heap_init entry)
        when the called function is a stub that computes constants and tail-calls (unconditional B).
        This handles the exact pattern in the provided disassembly:
          ADRP+ADD for X0 (base)
          ADRP+ADD for X8 (end)
          SUB X1, X8, X0  (size = end - base)
          B <real_init>
        It stops at the first unconditional B (0x14xxxxxx)
        and returns the tracked value of the target register.
        """
        start_off = at

        registers = [None] * 32  # Track X0-X31

        current_off = start_off
        max_instructions = 100  # Safety limit (heap_init stubs are tiny)

        for _ in range(max_instructions):
            instr = self.read_u32(current_off)
            if instr is None:
                return None

            pc = self.offset_to_va(current_off)

            # ADRP Xd, page
            adrp = self.decode_adrp(instr, pc)
            if adrp is not None:
                page, rd = adrp
                registers[rd] = page
                current_off += 4
                continue

            # ADD Xd, Xn, imm
            add = self.decode_add_imm(instr)
            if add is not None:
                rn, rd, imm = add
                if registers[rn] is not None:
                    if registers[rd] is None:
                        registers[rd] = 0
                    registers[rd] = registers[rn] + imm
                current_off += 4
                continue

            # SUB Xd, Xn, Xm  (register subtract)
            sub_reg = self.decode_sub_reg(instr)
            if sub_reg is not None:
                rn, rm, rd = sub_reg
                if registers[rn] is not None and registers[rm] is not None:
                    registers[rd] = registers[rn] - registers[rm]
                current_off += 4
                continue

            # Unconditional B (0x14000000) - tail call, return current register value
            if (instr & 0xFC000000) == 0x14000000:
                return registers[reg]

            # BL handling
            if (instr & 0xFC000000) == 0x94000000:
                return registers[reg]

            # RET (common exit) - also return current value
            if instr == 0xC0035FD6:  # RET X30
                return registers[reg]

            current_off += 4

        return None  # Did not hit exit

    def resolve_adrp_part(self, from_off: int, limit: int, reg: int, imm: int) -> Optional[int]:
        """ Resolve an adrp part by simulating forward from the call target """
        off = from_off
        while off > limit:
            instr = self.read_u32(off)
            if instr is None:
                break
            pc = self.base_addr + off
            adrp = self.decode_adrp(instr, pc)
            if adrp is not None:
                page, rd = adrp
                if rd == reg:
                    return page + imm
            off -= 4
        return None