import os
import sys

from PyQt5 import QtCore, QtWidgets, QtGui

import idaapi

from . import exceptions


def capture_widget(widget, path=None):
    """Grab an image of a Qt widget

    Args:
        widget: The Qt Widget to capture
        path (optional): The path to save to. If not provided - will return image data.

    Returns:
        If a path is provided, the image will be saved to it.
        If not, the PNG buffer will be returned.
    """
    pixmap = widget.grab()

    if path:
        pixmap.save(path)

    else:
        image_buffer = QtCore.QBuffer()
        image_buffer.open(QtCore.QIODevice.ReadWrite)

        pixmap.save(image_buffer, "PNG")

        return image_buffer.data().data()


def get_widget(title):
    """Get the Qt widget of the IDA window with the given title."""
    tform = idaapi.find_widget(title)
    if not tform:
        raise exceptions.FormNotFound("No form titled {!r} found.".format(title))

    return idaapi.PluginForm.FormToPyQtWidget(tform)


def resize_widget(widget, width, height):
    """Resize a Qt widget."""
    widget.setGeometry(0, 0, width, height)


def get_window():
    """Get IDA's top level window."""
    tform = idaapi.get_current_widget()

    # Required sometimes when closing IDBs and not IDA.
    if not tform:
        tform = idaapi.find_widget("Output window")

    widget = idaapi.PluginForm.FormToPyQtWidget(tform)
    window = widget.window()
    return window


class MenuManager(object):
    """IDA Menu Manipulation

    Use this class to add your own top-level menus.
    While this is discouraged by the SDK:

    > You should not change top level menu, or the Edit,Plugins submenus

    (documentation for `attach_action_to_menu`, kernwin.hpp)

    Adding top-level menus is useful sometimes.
    Nonetheless, you should be careful and make sure to remove all your menus
    when you are done. Leaving them handing would force users to restart IDA
    to remove them.

    Usage of this class should be as follows:

    >>> # Use the manager to add top-level menus
    >>> menu_manager = MenuManager()
    >>> menu_manager.add_menu("My Menu")
    >>> # Use the standard API to add menu items
    >>> idaapi.attach_action_to_menu("My Menu/", ":My-Action:", idaapi.SETMENU_APP)
    >>> # When a menu is not needed, remove it
    >>> menu_manager.remove_menu("My Menu")
    >>> # When you are done with the manager (and want to remove all menus you added,)
    >>> # clear it before deleting.
    >>> menu_manager.clear()
    """

    def __init__(self):
        super(MenuManager, self).__init__()

        self._window = get_window()
        self._menu = self._window.findChild(QtWidgets.QMenuBar)

        self._menus = {}

    def add_menu(self, name):
        """Add a top-level menu.

        The menu manager only allows one menu of the same name. However, it does
        not make sure that there are no pre-existing menus of that name.
        """
        if name in self._menus:
            raise exceptions.MenuAlreadyExists("Menu name {!r} already exists.".format(name))
        menu = self._menu.addMenu(name)
        self._menus[name] = menu

    def remove_menu(self, name):
        """Remove a top-level menu.

        Only removes menus created by the same menu manager.
        """
        if name not in self._menus:
            raise exceptions.MenuNotFound(
                "Menu {!r} was not found. It might be deleted, or belong to another menu manager.".format(name))

        self._menu.removeAction(self._menus[name].menuAction())
        del self._menus[name]

    def clear(self):
        """Clear all menus created by this manager."""
        for menu in self._menus.values():
            self._menu.removeAction(menu.menuAction())
        self._menus = {}