#!/usr/bin/env python

from pygrap import NodeInfo, node_t

try:
    from idaapi import is_code, get_flags, print_insn_mnem, generate_disasm_line, create_insn
    from idc import print_operand, get_bytes
except:
    from idaapi import isCode
    from idc import GetDisasm, GetFlags, GetMnem, GetOpnd, GetManyBytes, MakeCode
from idagrap.error.Exceptions import CodeException
import capstone


class Node(node_t):
    """Representation of a pygrap node, with useful methods.

    Args:
        ea (ea_t): Effective Address of the node.

    Attributs:
        node (node_t): Pygrap structure for a node.

    Raises:
        CodeException: If the ``ea`` is not a code instruction.
    """

    def __init__(self, ea, info, cs):
        """Initialization function."""
        # Init the node structure
        node_t.__init__(self)

        # Check if it's a code instruction
        try:
            is_c = is_code(get_flags(ea))
        except:
            is_c = isCode(GetFlags(ea))
        if not is_c:
            raise CodeException

        #
        # fill node_t struct
        #

        # NodeInfo
        self.info = NodeInfo()
        inst_elements = []

        try:
            size = create_insn(ea)
            bytes = get_bytes(ea, size)
        except:
            size = MakeCode(ea)
            bytes = GetManyBytes(ea, size)

        (address, size, mnemonic, op_str) = next(cs.disasm_lite(bytes, ea, count=1))
        self.info.opcode = mnemonic

        self.info.inst_str = self.info.opcode + " " + op_str

        splitted = op_str.split(", ")
        self.info.nargs = 0

        if len(splitted) >= 1:
            self.info.arg1 = splitted[0]
            self.info.nargs += 1
            if len(splitted) >= 2:
                self.info.arg2 = splitted[1]
                self.info.nargs += 1
                if len(splitted) >= 3:
                    self.info.arg3 = splitted[2]
                    self.info.nargs += 1

        # No node will be root but this is acceptable for CFGs
        self.info.is_root = False

        self.info.address = ea
        self.info.has_address = True

        # node_t
        self.node_id = self._genid()

    def getid(self):
        """Get the node id.

        Returns:
            vsize_t: The id of the node.
        """
        return self.node_id

    def _genid(self):
        """Generate a unique ID for the node.

        Returns:
            vsize_t: The id for the node
        """
        return self.info.address