import re
import typing
import keystone
import iced_x86


# ==-------------------------------------------------------------------== #
# Global and static variables, constants                                  #
# ==-------------------------------------------------------------------== #
assembly_archs = {
    "x86": keystone.KS_ARCH_X86
}

assembly_bitnesses = {
    64: keystone.KS_MODE_64,
    32: keystone.KS_MODE_32,
    16: keystone.KS_MODE_16,
}


# ==-------------------------------------------------------------------== #
# Classes                                                                 #
# ==-------------------------------------------------------------------== #
class InstructionTranscoder:
    """Creates instance to assemble, disassemble assembly code."""

    def __init__(self, arch: typing.Literal["x86"], bitness: typing.Literal[64, 32, 16]) -> None:
        """Creates instance of instructions transcodes to assemble / disassembly machine code."""

        # If arch is not allowed
        if arch not in (allowed_literals := typing.get_args(self.__init__.__annotations__["arch"])):
            raise Exception("Arch `%s` is not supported, expected one of: `%s`" % (arch, ", ".join(allowed_literals)))

        # If bitness is not allowed
        if bitness not in (allowed_literals := typing.get_args(self.__init__.__annotations__["bitness"])):
            raise Exception("Bitness `%s` is not supported, expected one of: `%s`" % (arch, ", ".join(allowed_literals)))

        # Save argument as instance attributes
        self.arch = arch
        self.bitness = bitness

        # Create formatter instance to disassemble code
        self.formatter = iced_x86.Formatter(iced_x86.FormatterSyntax.INTEL)

    def assemble(self, code: str, address_offset: int = 0, split_instuctions: bool = False) -> list[bytes] | bytes:
        """Asseble assembly code to machine code."""

        # Result list
        result = list()

        # Create assembly instance
        assembly = keystone.Ks(assembly_archs[self.arch], assembly_bitnesses[self.bitness])

        # If required to return entine code sequence as bytes
        if not split_instuctions:

            try:
                return assembly.asm(code, address_offset, as_bytes=True)[0] or bytes()

            except Exception:
                raise Exception("Unable to assemble code: `%s`" % code)

        # Iterate per each assembly code instruction and build it one by one
        for index, instucton in enumerate(re.split(r"[\n;]", code.strip()), start=1):

            # If instruction is empty
            if not (instucton := instucton.strip()):
                continue

            try:

                # Assemble instruction and save it into result list
                result.append(assembly.asm(instucton, address_offset, as_bytes=True)[0] or bytes())
                address_offset += len(result[-1])

            except Exception:
                raise Exception("Unable to assemble instuction #%d: `%s`" % (index, instucton))

        return result

    def disassemble(self, code: bytes, address_offset: int = 0, instuction_separator: str = "\n", split_instuctions: bool = False, stop_on_fail: bool = True) -> list[str] | str:
        """Disassemble machine code to assembly code."""

        # Result list
        result = list()

        # Create disassembly instance
        disassembly = iced_x86.Decoder(self.bitness, code, ip=address_offset)

        # Disasseble code until it's possible
        while disassembly.can_decode:

            # Retrieve disassembled instruction
            instruction = disassembly.decode()

            # If code disasseble failed
            if instruction.code == iced_x86.Code.INVALID:

                # If just break without exception raises required
                if stop_on_fail:
                    break

                raise Exception("Unable to disassmble code: `%s`" % code[instruction.ip:instruction.ip + instruction.len].hex(" "))

            # Formatting instuction text
            instruction_text = self.formatter.format(instruction)

            # Save disassebled instuction text into result list
            result.append(instruction_text)

        return instuction_separator.join(result) if not split_instuctions else result
