from collections import defaultdict
from typing import List

from PySide6.QtGui import QCursor, QContextMenuEvent
from PySide6.QtWidgets import QTreeWidget, QTreeWidgetItem, QMenu, QMessageBox, QInputDialog, QAbstractItemView
from PySide6.QtCore import Qt

from angr import SimState
from inspect import isfunction


class SimgrViewerAbstractTreeItem(QTreeWidgetItem):
    def handle_context_menu_event(self, event: QContextMenuEvent):
        """Handles right-click actions on the specific QTreeWidgetItem that was clicked in the view"""
        raise NotImplementedError

class StashTreeItem(SimgrViewerAbstractTreeItem):

    def __init__(self, stash_name, simgr_viewer):
        self.simgr_viewer = simgr_viewer
        self.stash_name = stash_name
        super().__init__(simgr_viewer)
        self.setFlags(self.flags() & ~Qt.ItemIsSelectable)
        self.refresh()

    def __iter__(self):
        for i in range(self.childCount()):
            yield self.child(i)

    @property
    def states(self):
        return self.simgr_viewer.simgr.stashes[self.stash_name]

    def refresh(self):
        self.takeChildren()
        for state in self.simgr_viewer.simgr.stashes[self.stash_name]:
            if self.stash_name == "errored" and getattr(state, "state", None):
                state = state.state
            self.addChild(StateTreeItem(state, self.simgr_viewer))
        self.setText(0, "%s (%d)" % (self.stash_name, len(self.states)))

    def handle_context_menu_event(self, event):
        menu = QMenu()
        menu.addAction("Copy states", self.copy_states)
        menu.addAction("Cut states", self.cut_states)
        if self.simgr_viewer.state_clipboard:
            plural = ""
            if len(self.simgr_viewer.state_clipboard) > 1:
                plural += "s"
            menu.addAction("Paste state" + plural, self.paste_states)
        menu.addAction("Delete stash", self.delete_stash)
        if self.stash_name != "active":
            menu.addAction("Move states to here", self.move_states)
        menu.exec_(QCursor.pos())

    def copy_states(self):
        self.simgr_viewer.state_clipboard = [s.state for s in self]
        self.refresh()

    def cut_states(self):
        self.simgr_viewer.state_clipboard = [s.state for s in self]
        self.simgr_viewer.simgr.drop(stash=self.stash_name, filter_func=lambda state: state in self.states)
        self.refresh()

    def delete_stash(self, *args, **kwargs): #pylint: disable=unused-argument
        self.simgr_viewer.simgr._stashes.pop(self.stash_name)
        self.simgr_viewer.refresh()

    def paste_states(self, *args, **kwargs): #pylint: disable=unused-argument
        self.simgr_viewer.paste_from_clipboard(self.stash_name)
        self.refresh()

    def move_states(self, *args, **kwargs): #pylint: disable=unused-argument
        self.simgr_viewer.move_to_stash(self.stash_name)
        self.refresh()


class StateTreeItem(SimgrViewerAbstractTreeItem):
    def __init__(self, state, simgr_viewer):
        self.state = state
        self.simgr_viewer: QSimulationManagerViewer = simgr_viewer
        super().__init__([str(state)])
        self.setData(0, 1, state)

    @property
    def stash_name(self):
        return self.parent().stash_name

    def handle_context_menu_event(self, event):
        menu = QMenu()
        self.add_menu_action(menu, "Copy state", self.copy_states)
        self.add_menu_action(menu, "Cut state", self.cut_states)
        self.add_menu_action(menu, "Delete state", self.delete_states)
        if self.simgr_viewer.state_clipboard:
            self.add_menu_action(menu, "Paste state", self.paste_states)
        menu.exec_(QCursor.pos())

    def add_menu_action(self, menu, string, action):
        plural = ""
        if len(self.simgr_viewer.selectedItems()) > 1:
            plural = "s"
        menu.addAction(string + plural, action)

    def copy_states(self, *args, **kwargs): #pylint: disable=unused-argument
        self.simgr_viewer.copy_selected_to_clipboard()

    def cut_states(self, *args, **kwargs): #pylint: disable=unused-argument
        self.simgr_viewer.cut_selected_to_clipboard()

    def delete_states(self, *args, **kwargs): #pylint: disable=unused-argument
        self.simgr_viewer.delete_selected_states()

    def paste_states(self, *args, **kwargs): #pylint: disable=unused-argument
        self.simgr_viewer.paste_from_clipboard(self.stash_name)


class QSimulationManagerViewer(QTreeWidget):
    state_clipboard: List[SimState]

    def __init__(self, simgr, parent=None):
        super().__init__(parent)

        self.setColumnCount(1)
        self.setHeaderHidden(True)

        self.simgr = simgr
        self.state_clipboard = []

        self._init_widgets()

        self.simgr.am_subscribe(self.refresh)
        self.setSelectionMode(QAbstractItemView.ExtendedSelection)

    def _stash_to_selected_states(self):
        stash_to_states = defaultdict(list)
        for state_tree_item in self.selectedItems():
            stash_to_states[state_tree_item.stash_name].append(state_tree_item.state)
        return stash_to_states

    def copy_selected_to_clipboard(self):
        self.state_clipboard = [item.state.copy() for item in self.selectedItems()]

    def cut_selected_to_clipboard(self):
        self.copy_selected_to_clipboard()
        self.delete_selected_states()

    def delete_selected_states(self):
        stash_to_states = self._stash_to_selected_states()
        for stash_name, states in stash_to_states.items():
            self.simgr.drop(stash=stash_name, filter_func=lambda state, state_set=states: state in state_set)
            self.get_stash_tree_item(stash_name).refresh()

    def paste_from_clipboard(self, stash_name):
        self.simgr.populate(stash_name, self.state_clipboard)
        self.get_stash_tree_item(stash_name).refresh()

    def move_to_stash(self, stash_name):
        lambda_str = ""
        lambda_func = None
        while True:
            lambda_str, accepted = QInputDialog.getText(self, "Move state from active to here", "Condition lambda",
                                                        text=lambda_str)
            if not accepted:
                return
            try:
                lambda_func = eval(lambda_str) #pylint: disable=eval-used
                if not isfunction(lambda_func):
                    raise ValueError()
            except Exception as e:  # pylint: disable=broad-except
                QMessageBox.critical(self, "Exception!", str(e))
                continue
            break
        self.simgr.move(from_stash="active", to_stash=stash_name, filter_func=lambda_func)

    def contextMenuEvent(self, event: QContextMenuEvent):
        item = self.itemAt(event.pos())
        if item is not None:
            item.handle_context_menu_event(event)
        else:
            menu = QMenu()
            menu.addAction("Create new stash", self._create_new_stash)
            menu.exec_(QCursor.pos())

    def _create_new_stash(self, *args, **kwargs): #pylint: disable=unused-argument

        stash_name, accepted = QInputDialog.getText(self, "Stash name", "Blah")

        if not accepted or stash_name.strip() == "":
            # The user didn't provide a stash name
            return

        if stash_name in self.simgr.stashes:
            QMessageBox.critical(None, 'Duplicate stash name',
                                 f"A stash with the name {stash_name} already exists in the current simulation manager.")
            return
        self.simgr._stashes[stash_name] = []
        self.refresh()

    def refresh(self, **kwargs):
        if kwargs.get('src',"") != 'simgr_viewer':
            self._init_widgets()

    def current_state(self):
        item = self.currentItem()
        if item is None:
            return None
        return item.data(0, 1)

    def select_state(self, state):
        if state is None:
            self.setCurrentItem(None)
        else:
            for i in range(self.topLevelItemCount()):
                item = self.topLevelItem(i)
                for j in range(item.childCount()):
                    subitem = item.child(j)
                    if subitem.data(0, 1) == state:
                        self.setCurrentItem(subitem)
                        break
                else:
                    continue
                break

    def get_stash_tree_item(self, stash_name):
        return self.stash_tree_items[stash_name]

    def _init_widgets(self):
        # save expanded state
        expended_stash = {}
        selected_state = set()
        for topidx in range(self.topLevelItemCount()):
            top = self.topLevelItem(topidx)
            expended_stash[top.stash_name] = top.isExpanded()
            for idx in range(top.childCount()):
                item = top.child(idx)
                if item.isSelected():
                    selected_state.add(item.state)

        self.clear()

        if self.simgr.am_none:
            return

        self.stash_tree_items = {}
        for stash_name, stash in self.simgr.stashes.items(): #pylint: disable=unused-variable
            # if not stash and stash_name not in ('active', 'deadended', 'avoided'):
            #     continue
            item = StashTreeItem(stash_name, simgr_viewer=self)
            self.stash_tree_items[stash_name] = item
            self.addTopLevelItem(item)
            item.setExpanded(expended_stash.get(stash_name, False))
