import re
import typing
import keystone
import capstone


# ==-------------------------------------------------------------------== #
# 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,
}

disassembly_archs = {
    "x86": capstone.CS_ARCH_X86
}

disassembly_bitnesses = {
    64: capstone.CS_MODE_64,
    32: capstone.CS_MODE_32,
    16: capstone.CS_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 assembly and disassembly instances
        self.assembly = keystone.Ks(assembly_archs[self.arch], assembly_bitnesses[self.bitness])
        self.disassembly = capstone.Cs(disassembly_archs[self.arch], disassembly_bitnesses[self.bitness])

        # Set up assembly and disassembly parameters
        self.disassembly.skipdata = True

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

        # Result list
        result = list()

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

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

            try:

                # Assemble instuction
                assembled_instruction = self.assembly.asm(instucton, address_offset, as_bytes=True)[0]

                # Save assembled instruction to result list
                result.append({
                    "address": address_offset,
                    "instruction": instucton,
                    "size": len(assembled_instruction),
                    "bytes": assembled_instruction,
                })

                # Move address offset
                address_offset += result[-1]["size"]

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

        return result

    def disassemble(self, code: bytes, address_offset: int = 0, raise_on_fail: bool = False) -> list[dict[str]]:
        """Disassemble machine code to assembly code."""

        # Result list
        result = list()

        # Disasseble code until it's possible
        for instruction in self.disassembly.disasm(code, address_offset):

            # If code disasseble failed
            if instruction.mnemonic == ".byte":

                # If exception raising doesn't required
                if not raise_on_fail:
                    break

                raise Exception("Unable to disassmble code: `%s`" % instruction.op_str[2:])

            # Save disassembled instruction to result list
            result.append({
                "size": instruction.size,
                "bytes": instruction.bytes,
                "address": instruction.address,
                "mnemonic": instruction.mnemonic,
                "operation": instruction.op_str,
                "code": f"{instruction.mnemonic} {instruction.op_str}".strip(),
            })

        return result
