from __future__ import annotations

from typing import Dict, List, Optional, Union

from boa3.compiler.codegenerator.methodtokencollection import MethodTokenCollection
from boa3.compiler.codegenerator.vmcodemap import VMCodeMap
from boa3.compiler.compileroutput import CompilerOutput
from boa3.model.builtin.method import IBuiltinMethod
from boa3.neo.vm.VMCode import VMCode
from boa3.neo.vm.opcode import OpcodeHelper
from boa3.neo.vm.opcode.OpcodeInformation import OpcodeInformation
from boa3.neo3.contracts.contracttypes import CallFlags


class VMCodeMapping:
    """
    This class is responsible for managing the Neo VM instruction during the bytecode generation.
    """
    _instance: VMCodeMapping = None

    @classmethod
    def instance(cls):
        """
        :return: the singleton instance
        """
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    def __init__(self):
        self._code_map: VMCodeMap = VMCodeMap()
        self._method_tokens: MethodTokenCollection = MethodTokenCollection()

    @classmethod
    def reset(cls):
        """
        Resets the map to the first state
        """
        if cls._instance is not None:
            cls._instance._code_map.clear()
            cls._instance._method_tokens.clear()

    def add_method_token(self, method: IBuiltinMethod, call_flag: CallFlags) -> Optional[int]:
        """
        Creates a new method token if the method call another contract and return its id.
        Otherwise, returns None
        """
        if hasattr(method, 'contract_script_hash'):
            return self._method_tokens.append(method, call_flag)
        return None

    @property
    def codes(self) -> List[VMCode]:
        """
        Gets a list with the included vm codes

        :return: a list of vm codes ordered by its address in the bytecode
        """
        return self._code_map.get_code_list()

    @property
    def code_map(self) -> Dict[int, VMCode]:
        """
        Gets a dictionary that maps each vm code with its address.

        :return: a dictionary that maps each instruction with its address. The keys are ordered by the address.
        """
        return self._code_map.get_code_map()

    def targeted_address(self) -> Dict[int, List[int]]:
        """
        Gets a dictionary that maps each address to the opcodes that targets it

        :return: a dictionary that maps the targeted instructions to its source.
        """
        target_maps = {}
        for code in self._code_map.get_code_with_target_list():
            if code.target is not None and code.target is not code:
                address = self.get_start_address(code)
                target = self.get_start_address(code.target)
                if target not in target_maps:
                    target_maps[target] = [address]
                else:
                    target_maps[target].append(address)
        return target_maps

    def bytecode(self) -> bytes:
        """
        Gets the bytecode of the translated code

        :return: the generated bytecode
        """
        self._remove_empty_targets()
        self._update_larger_codes()

        bytecode = bytearray()
        for code in self.codes:
            bytecode += code.opcode
            if code.data is not None:
                bytecode += code.data
        return bytes(bytecode)

    def result(self) -> CompilerOutput:
        """
        Gets the complete output of the translated code
        """
        bytecode = self.bytecode()
        return CompilerOutput(bytecode, self._method_tokens.to_list())

    @property
    def bytecode_size(self) -> int:
        return self._code_map.get_bytecode_size()

    def insert_code(self, vm_code: VMCode):
        return self._code_map.insert_code(vm_code, has_target=OpcodeHelper.has_target(vm_code.opcode))

    def get_code(self, address: int) -> Optional[VMCode]:
        """
        Gets the VM Opcode at the given position

        :param address: the position of the opcode
        :return: the opcode if it exists. None otherwise
        :rtype: VMCode or None
        """
        return self._code_map.get_code(address)

    def get_addresses(self, start_address: int, end_address: int) -> List[int]:
        return self._code_map.get_addresses(start_address, end_address)

    def get_start_address(self, vm_code: VMCode) -> int:
        """
        Gets the vm code's first byte address

        :param vm_code: the instruction to get the address
        :return: the vm code's address if it's in the map. Otherwise, return's zero.
        """
        return self._code_map.get_start_address(vm_code)

    def get_end_address(self, vm_code: VMCode) -> int:
        """
        Gets the vm code's last byte address

        :param vm_code: the instruction to get the address
        :return: the vm code's last address if it's in the map. Otherwise, return's zero.
        """
        return self._code_map.get_end_address(vm_code)

    def get_opcodes(self, addresses: List[int]) -> List[VMCode]:
        return self._code_map.get_opcodes(addresses)

    def update_vm_code(self, vm_code: VMCode, opcode: OpcodeInformation, data: bytes = bytes()):
        """
        Updates the information from an inserted code

        :param vm_code: code to be updated
        :param opcode: updated opcode information
        :param data: updated opcode data
        """
        code_size = vm_code.size
        vm_code._info = opcode
        vm_code._data = data
        if vm_code.size != code_size:
            self._update_addresses(self.get_start_address(vm_code))

    def _update_addresses(self, start_address: int = 0):
        """
        Updates the instruction map's keys when a opcode is changed

        :param start_address: the address from the changed opcode
        """
        return self._code_map.update_addresses(start_address)

    def _update_targets(self):
        from boa3.neo.vm.type.Integer import Integer
        for code in self._code_map.get_code_with_target_list():
            if code.target is None:
                relative = Integer.from_bytes(code.data)
                absolute = self._code_map.get_start_address(code) + relative
                if absolute in self.code_map:
                    code.set_target(self.code_map[absolute])

    def _update_larger_codes(self):
        """
        Checks if each instruction data fits in its opcode maximum size and updates the opcode from those that don't
        """
        # gets a list with all instructions which its opcode has a larger equivalent, ordered by its address
        instr_with_small_codes = [code for code in self._code_map.get_code_list() if OpcodeHelper.has_larger_opcode(code.opcode)]
        instr_with_small_codes.sort(key=lambda code: self.get_start_address(code), reverse=True)

        from boa3.neo.vm.opcode.OpcodeInfo import OpcodeInfo
        # total_len is initialized with zero because the loop must run at least once
        total_len = 0
        current_size = self.bytecode_size

        # if any instruction is updated, the following instruction addresses and the total size will change as well
        # with the change, previous instruction data may have overflowed the opcode maximum value
        # to make sure, it must check the opcodes that haven't changed again
        while total_len != current_size:
            total_len = current_size

            # verifies each instruction data length
            for code in instr_with_small_codes.copy():  # it's a copy because the list may change during the iteration
                if len(code.raw_data) > code.info.max_data_len:
                    # gets the shortest opcode equivalent that fits the instruction data
                    info = OpcodeInfo.get_info(OpcodeHelper.get_larger_opcode(code.opcode))
                    while len(code.raw_data) > info.max_data_len and OpcodeHelper.has_larger_opcode(info.opcode):
                        info = OpcodeInfo.get_info(OpcodeHelper.get_larger_opcode(code.opcode))

                    self.update_vm_code(code, info)
                    if info.opcode == OpcodeHelper.get_larger_opcode(info.opcode):
                        # if it's the largest equivalent, it won't be updated anymore
                        instr_with_small_codes.remove(code)
            current_size = self.bytecode_size

    def _validate_targets(self, code_or_address: Union[int, VMCode]):
        if isinstance(code_or_address, int):
            address = code_or_address
            code = self.get_code(address)
        else:
            code = code_or_address
            address = self.get_start_address(code)

        targeted_addresses = self.targeted_address()

        if address in targeted_addresses:
            next_address = self.get_end_address(code) + 1
            if next_address < self.bytecode_size:
                next_code = self._code_map.get_code(next_address)
                for source in targeted_addresses[address]:
                    self._code_map.get_code(source).set_target(next_code)

    def move_to_end(self, first_code_address: int, last_code_address: int) -> int:
        """
        Moves a set of instructions to the end of the current bytecode

        :param first_code_address: first instruction start address
        :param last_code_address: last instruction end address
        """
        result = self._code_map.move_to_end(first_code_address, last_code_address)
        if not isinstance(result, int):
            return self.bytecode_size

        self._update_targets()
        return result

    def remove_opcodes(self, first_code_address: int, last_code_address: int = None):
        if not isinstance(last_code_address, int):
            last_code_address = self.bytecode_size
        addresses_to_remove = self._code_map.get_addresses(first_code_address, last_code_address)
        for address in addresses_to_remove:
            self._validate_targets(address)
        return self._code_map.remove_opcodes_by_addresses(addresses_to_remove)

    def remove_opcodes_by_code(self, codes: List[VMCode]):
        addresses_to_remove = self._code_map.get_addresses_from_codes(codes)
        for address in addresses_to_remove:
            self._validate_targets(address)
        return self._code_map.remove_opcodes_by_addresses(addresses_to_remove)

    def _remove_empty_targets(self):
        """
        Checks if each instruction that requires a target has one set and remove those that don't
        """
        addresses_to_remove = []
        for code in self._code_map.get_code_with_target_list():
            if code.target is None or code.target is code:
                address = self.get_start_address(code)
                self._validate_targets(address)
                addresses_to_remove.append(address)

        if len(addresses_to_remove) > 0:
            self._code_map.remove_opcodes_by_addresses(addresses_to_remove)
