import logging

from PySide6.QtWidgets import QWidget, QVBoxLayout, QGraphicsScene, QGraphicsView, QGraphicsItemGroup
from PySide6.QtWidgets import QTabWidget, QPushButton, QAbstractItemView
from PySide6.QtWidgets import  QMessageBox, QTableWidget, QTableWidgetItem
from PySide6.QtGui import QPen, QBrush, QLinearGradient, QColor, QPainter, QImage, QFont
from PySide6.QtCore import Qt, QPoint, QEvent


l = logging.getLogger(name=__name__)


class QTraceViewer(QWidget):
    """
    Load a basic block trace through json and visualize it in the disassembly
    Ref: https://github.com/angr/angr-management/pull/122
    """

    TAG_SPACING = 50
    LEGEND_X = -50
    LEGEND_Y = 0
    LEGEND_WIDTH = 10

    TRACE_FUNC_X = 0
    TRACE_FUNC_Y = 0
    TRACE_FUNC_WIDTH = 50
    TRACE_FUNC_MINHEIGHT = 1000

    TAB_HEADER_SIZE = 40
    MAX_WINDOW_SIZE = 500

    MARK_X = LEGEND_X
    MARK_WIDTH = TRACE_FUNC_X - LEGEND_X + TRACE_FUNC_WIDTH
    MARK_HEIGHT = 1

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

        self.mark = None
        self.legend = None
        self.legend_height = 0
        self.legend_img = None
        self.trace_func_unit_height = 0


        self.trace_func = None
        self.trace_id = None

        self.view = None
        self.traceView = None
        self.traceTab = None
        self.traceScene = None
        self.multiView = None
        self.listView = None
        self.mark = None
        self.curr_position = 0
        self._use_precise_position = False
        self._selected_traces = []

        self._init_widgets()

        self.trace.am_subscribe(self._on_set_trace)
        self.selected_ins.am_subscribe(self._on_select_ins)
        self.traceTab.installEventFilter(self)

    #
    # Forwarding properties
    #

    @property
    def trace(self):
        return self.workspace.main_instance.trace

    @property
    def multi_trace(self):
        return self.workspace.main_instance.multi_trace

    @property
    def selected_ins(self):
        return self.disasm_view.infodock.selected_insns

    def _init_widgets(self):
        self.view = QTabWidget() # QGraphicsView()

        self.traceTab = QWidget()
        tracelayout = QVBoxLayout()

        self.traceView = QGraphicsView()
        self.traceScene = QGraphicsScene()
        self.traceView.setScene(self.traceScene)

        self.listView = QTableWidget(0,2) # row, col
        self.listView.setHorizontalHeaderItem(0, QTableWidgetItem("Trace ID"))
        self.listView.setHorizontalHeaderItem(1, QTableWidgetItem("Input ID"))
        self.listView.setSelectionMode(QAbstractItemView.SingleSelection)
        self.listView.setSelectionBehavior(QAbstractItemView.SelectRows)
        # self.listView.horizontalHeader().setStretchLastSection(True)
        # self.listView.horizontalHeader().setSectionResizeModel(0, QHeaderView.Stretch)
        self.listView.cellClicked.connect(self._switch_current_trace)

        self.traceSeedButton = QPushButton("View Input Seed")
        self.traceSeedButton.clicked.connect(self._view_input_seed)

        tracelayout.addWidget(self.traceView)
        tracelayout.addWidget(self.listView)
        tracelayout.addWidget(self.traceSeedButton)
        self.traceTab.setLayout(tracelayout)

        self.multiView = QWidget()
        multiLayout = QVBoxLayout()
        self.multiTraceList = QTableWidget(0,2) # row, col
        self.multiTraceList.setSelectionMode(QAbstractItemView.MultiSelection)
        self.multiTraceList.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.multiTraceList.setHorizontalScrollMode(QAbstractItemView.ScrollPerPixel)
        self.multiTraceList.setHorizontalHeaderItem(0, QTableWidgetItem("Trace ID"))
        self.multiTraceList.setHorizontalHeaderItem(1, QTableWidgetItem("Input ID"))
        self.selectMultiTrace = QPushButton("Refresh Heatmap")
        self.selectMultiTrace.clicked.connect(self._refresh_heatmap)
        multiLayout.addWidget(self.multiTraceList)
        multiLayout.addWidget(self.selectMultiTrace)
        self.multiView.setLayout(multiLayout)

        self.view.addTab(self.traceTab, "SingleTrace")
        self.view.addTab(self.multiView, "MultiTrace HeatMap")
        self.SINGLE_TRACE = 0
        self.MULTI_TRACE = 1

        self.view.currentChanged.connect(self._on_tab_change)

        self._reset()

        layout = QVBoxLayout()
        layout.addWidget(self.view)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setAlignment(self.view, Qt.AlignLeft)

        self.setLayout(layout)

    def _reset(self):
        self.traceScene.clear() #clear items
        self.listView.clearContents()
        self.multiTraceList.clearContents()
        self.mark = None

        self.legend = None
        self.legend_height = 0

        self.trace_func = QGraphicsItemGroup()
        self.trace_id = QGraphicsItemGroup()
        self.traceScene.addItem(self.trace_func)
        self.hide()

    def _view_input_seed(self):
        current_trace_stats = self.trace.am_obj
        input_id = current_trace_stats.input_id

        inputSeed = self.multi_trace.am_obj.get_input_seed_for_id(input_id)
        msgText = "%s" % inputSeed
        msgDetails = "Input for [%s]" % current_trace_stats.id
        msgbox = QMessageBox()
        msgbox.setWindowTitle("Seed Input")
        msgbox.setDetailedText(msgDetails)
        msgbox.setText(msgText)
        msgbox.setStandardButtons(QMessageBox.Ok)
        msgbox.exec()

    def _switch_current_trace(self, row):
        if self.listView.rowCount() <= 0:
            return

        current_trace = self.trace.am_obj.id
        new_trace = self.multiTraceList.item(row, 0).text()
        if current_trace == new_trace:
            return

        trace_stats = self.multi_trace.am_obj.get_trace_with_id(new_trace)
        if trace_stats:
            self.trace.am_obj = trace_stats
            self._on_set_trace()

    def _on_set_trace(self):
        self._reset()
        if self.trace.am_none or self.trace.count is None:
            return

        l.debug('minheight: %d, count: %d', self.TRACE_FUNC_MINHEIGHT,
                self.trace.count)
        if self.trace.count <= 0:
            l.warning("No valid addresses found in trace to show. Check base address offsets?")
            self.trace.am_obj = None
            self.trace.am_event()
            return
        if self.TRACE_FUNC_MINHEIGHT < self.trace.count * 15:
            self.trace_func_unit_height = 15
            show_func_tag = True
        else:
            self.trace_func_unit_height = self.TRACE_FUNC_MINHEIGHT / self.trace.count
            show_func_tag = True

        self.legend_height = int(self.trace.count * self.trace_func_unit_height)

        self._show_trace_func(show_func_tag)
        self._show_legend()
        self._show_trace_ids()
        self._set_mark_color()
        self._refresh_multi_list()

        boundingSize = self.traceScene.itemsBoundingRect().width()
        windowSize = boundingSize
        if boundingSize > self.MAX_WINDOW_SIZE:
            windowSize = self.MAX_WINDOW_SIZE
        self.traceScene.setSceneRect(self.traceScene.itemsBoundingRect()) #resize
        self.setFixedWidth(windowSize)

        # self.listScene.setSceneRect(self.listScene.itemsBoundingRect()) #resize
        self.multiView.setFixedWidth(windowSize)
        cellWidth = windowSize // 2
        self.listView.setColumnWidth(0, cellWidth)
        self.listView.setColumnWidth(1, cellWidth)
        self.listView.setFixedHeight(self.multiView.height() // 4)
        self.multiTraceList.setColumnWidth(0, cellWidth)
        self.multiTraceList.setColumnWidth(1, cellWidth)
        self.view.setFixedWidth(windowSize)

        self.show()

    def _populate_trace_table(self, view, trace_ids):
        numIDs = len(trace_ids)
        view.clearContents()
        view.setRowCount(numIDs)
        row = 0 #start after label row
        for traceID in trace_ids:
            inputID = self.multi_trace.am_obj.get_input_id_for_trace_id(traceID)
            if inputID is None:
                self.workspace.log("No inputID found for trace %s" % traceID)
            view.setItem(row, 0, QTableWidgetItem(traceID))
            view.setItem(row, 1, QTableWidgetItem(inputID))
            row += 1

    def _refresh_heatmap(self):
        multiTrace = self.multi_trace.am_obj
        multiTrace.clear_heatmap()
        multiTrace.is_active_tab = True

        selected_items = self.multiTraceList.selectedItems()
        self._selected_traces.clear()
        for row in range(self.multiTraceList.rowCount()):
            item = self.multiTraceList.item(row, 0)
            if item in selected_items:
                self._selected_traces.append(item.text())
        multiTrace.reload_heatmap(self._selected_traces)
        self.multi_trace.am_event()

    def _refresh_multi_list(self):
        multiTrace = self.multi_trace.am_obj
        trace_ids = multiTrace.get_all_trace_ids()

        self.multiTraceList.clearContents()
        self._populate_trace_table(self.multiTraceList, trace_ids)
        if self._selected_traces and self.multiTraceList.rowCount() > 0:
            self.multiTraceList.item(0,0).setSelected(True)
            self.multiTraceList.item(0,1).setSelected(True)
        else:
            for row in range(self.multiTraceList.rowCount()):
                item = self.multiTraceList.item(row, 0)
                inputItem = self.multiTraceList.item(row, 1)
                if item.text() in self._selected_traces:
                    item.setSelected(True)
                    inputItem.setSelected(True)
        self.multi_trace.am_event()


    def _on_tab_change(self):
        # self._reset()
        multiTrace = self.multi_trace.am_obj
        if self.view.currentIndex() == self.MULTI_TRACE:
            multiTrace.is_active_tab = True
            self._refresh_multi_list()
        elif self.view.currentIndex() == self.SINGLE_TRACE:
            multiTrace = self.multi_trace.am_obj
            multiTrace.is_active_tab = False
            self._show_trace_ids()

    def _on_select_ins(self, **kwargs): # pylint: disable=unused-argument
        if self.trace.am_none:
            return

        if self.mark is not None:
            for i in self.mark.childItems():
                self.mark.removeFromGroup(i)
                self.traceScene.removeItem(i)
            self.traceScene.removeItem(self.mark)

        self.mark = QGraphicsItemGroup()
        self.traceScene.addItem(self.mark)

        if self.selected_ins:
            addr = next(iter(self.selected_ins))
            positions = self.trace.get_positions(addr)
            if positions: #if addr is in list of positions
                if not self._use_precise_position: #handle case where insn was selected from disas view
                    self.curr_position = positions[0] - self.trace.count
                for p in positions:
                    color = self._get_mark_color(p, self.trace.count)
                    y = self._get_mark_y(p)

                    if p == self.trace.count + self.curr_position: #add thicker line for 'current' mark
                        self.mark.addToGroup(self.traceScene.addRect(self.MARK_X, y, self.MARK_WIDTH,
                                            self.MARK_HEIGHT*4, QPen(QColor('black')), QBrush(color)))
                    else:
                        self.mark.addToGroup(self.traceScene.addRect(self.MARK_X, y, self.MARK_WIDTH,
                                                                self.MARK_HEIGHT, QPen(color), QBrush(color)))

                self.traceScene.update() #force redraw of the traceScene
                self.scroll_to_position(self.curr_position)

    def scroll_to_position(self, position):
        relative_pos = self.trace.count + position
        y_offset = self._get_mark_y(relative_pos)

        scrollValue = 0
        if y_offset > 0.5 * self.traceView.size().height():
            scrollValue = y_offset - 0.5 * self.traceView.size().height()
        scrollValue = min(scrollValue, self.traceView.verticalScrollBar().maximum())
        self.traceView.verticalScrollBar().setValue(scrollValue)
        self._use_precise_position = False

    def jump_next_insn(self):
        if self.curr_position + self.trace.count < self.trace.count - 1: #for some reason indexing is done backwards
            self.curr_position += 1
            self._use_precise_position = True
            bbl_addr = self.trace.get_bbl_from_position(self.curr_position)
            func = self.trace.get_func_from_position(self.curr_position)
            self._jump_bbl(func, bbl_addr)

    def jump_prev_insn(self):
        if self.curr_position + self.trace.count > 0:
            self.curr_position -= 1
            self._use_precise_position = True
            bbl_addr = self.trace.get_bbl_from_position(self.curr_position)
            func = self.trace.get_func_from_position(self.curr_position)
            self._jump_bbl(func, bbl_addr)

    def eventFilter(self, obj, event): #specifically to catch arrow keys #pylint: disable=unused-argument
        # more elegant solution to link w/ self.view's scroll bar keypressevent?
        if event.type() == QEvent.Type.KeyPress:
            if not event.modifiers() & Qt.ShiftModifier: #shift + arrowkeys
                return False
            key = event.key()
            if key in [Qt.Key_Up, Qt.Key_Left]:
                self.jump_prev_insn()
            elif key in [Qt.Key_Down, Qt.Key_Right]:
                self.jump_next_insn()
            return True

        return False  # pass through all other events

    def mousePressEvent(self, event):
        button = event.button()
        pos = self._to_logical_pos(event.pos())
        if button == Qt.LeftButton and self.view.currentIndex() == self.SINGLE_TRACE and self._at_legend(pos):
            func = self._get_func_from_y(pos.y())
            bbl_addr = self._get_bbl_from_y(pos.y())
            self._use_precise_position = True
            self.curr_position = self._get_position(pos.y())
            self._jump_bbl(func, bbl_addr)

    def _jump_bbl(self, func, bbl_addr):
        all_insn_addrs = self.workspace.main_instance.project.factory.block(bbl_addr).instruction_addrs
        # TODO: replace this with am_events perhaps?
        if func is None:
            return
        self.workspace.on_function_selected(func)
        self.selected_ins.clear()
        self.selected_ins.update(all_insn_addrs)
        self.selected_ins.am_event()
        # TODO: this ought to happen automatically as a result of the am_event
        self.disasm_view.current_graph.show_instruction(bbl_addr)

    def _get_mark_color(self, i, total):
        relative_gradient_pos = i * 1000 // total
        return self.legend_img.pixelColor(self.LEGEND_WIDTH // 2,
                                          relative_gradient_pos)

    def _get_mark_y(self, i):
        return self.TRACE_FUNC_Y + self.trace_func_unit_height * i


    def _show_trace_ids(self):
        trace_ids = self.multi_trace.get_all_trace_ids()
        # traceID = self.listScene.addText(id_txt, QFont("Source Code Pro", 7))
        # traceID.setPos(5,5)
        self.listView.clearContents()
        self._populate_trace_table(self.listView, trace_ids)
        if len(self.listView.selectedItems()) <= 0 and not self.trace.am_none:
            for row in range(self.listView.rowCount()):
                item = self.listView.item(row, 0)
                inputItem = self.listView.item(row, 1)
                if self.trace.id in item.text():
                    item.setSelected(True)
                    inputItem.setSelected(True)
                    break



    def _show_trace_func(self, show_func_tag):
        x = self.TRACE_FUNC_X
        y = self.TRACE_FUNC_Y
        prev_name = None
        for position in self.trace.trace_func:
            bbl_addr = position.bbl_addr
            func_name = position.func_name
            l.debug('Draw function %x, %s', bbl_addr, func_name)
            color = self.trace.get_func_color(func_name)
            self.trace_func.addToGroup(self.traceScene.addRect(x, y,
                                                          self.TRACE_FUNC_WIDTH, self.trace_func_unit_height,
                                                          QPen(color), QBrush(color)))
            if show_func_tag is True and func_name != prev_name:
                tag = self.traceScene.addText(func_name, QFont("Source Code Pro", 7))
                tag.setPos(x + self.TRACE_FUNC_WIDTH +
                           self.TAG_SPACING, y -
                           tag.boundingRect().height() // 2)
                self.trace_func.addToGroup(tag)
                anchor = self.traceScene.addLine(
                    self.TRACE_FUNC_X + self.TRACE_FUNC_WIDTH, y,
                    x + self.TRACE_FUNC_WIDTH + self.TAG_SPACING, y)
                self.trace_func.addToGroup(anchor)
                prev_name = func_name
            y += self.trace_func_unit_height

    @staticmethod
    def _make_legend_gradient(x1, y1, x2, y2):
        gradient = QLinearGradient(x1, y1, x2, y2)
        gradient.setColorAt(0.0, Qt.red)
        gradient.setColorAt(0.4, Qt.yellow)
        gradient.setColorAt(0.6, Qt.green)
        gradient.setColorAt(0.8, Qt.blue)
        gradient.setColorAt(1.0, Qt.darkBlue)
        return gradient


    def _show_legend(self):
        pen = QPen(Qt.transparent)

        gradient = self._make_legend_gradient(self.LEGEND_X, self.LEGEND_Y,
                                   self.LEGEND_X, self.LEGEND_Y + self.legend_height)
        brush = QBrush(gradient)
        self.legend = self.traceScene.addRect(self.LEGEND_X, self.LEGEND_Y,
                                         self.LEGEND_WIDTH, self.legend_height, pen, brush)

        reference_gradient = self._make_legend_gradient(0, 0, self.LEGEND_WIDTH, 1000)
        base_img = QImage(self.LEGEND_WIDTH, 1000, QImage.Format.Format_ARGB32)
        p = QPainter(base_img)
        p.fillRect(base_img.rect(),reference_gradient)
        self.legend_img = base_img #reference shade


    def _set_mark_color(self):
        for p in range(self.trace.count):
            color = self._get_mark_color(p, self.trace.count)
            self.trace.set_mark_color(p, color)

    def _at_legend(self, pos):
        x = pos.x()
        y = pos.y()
        return self.TRACE_FUNC_X + self.LEGEND_X < x < self.traceView.width() and \
           self.TRACE_FUNC_Y < y < self.TRACE_FUNC_Y + self.legend_height

    def _to_logical_pos(self, pos):
        x_offset = self.traceView.horizontalScrollBar().value()
        y_offset = self.traceView.verticalScrollBar().value()
        return QPoint(pos.x() + x_offset, pos.y() + y_offset)

    def _get_position(self, y):
        y_relative = y - self.legend_height - self.TAB_HEADER_SIZE

        return int(y_relative // self.trace_func_unit_height)

    def _get_bbl_from_y(self, y):
        position = self._get_position(y)
        return self.trace.get_bbl_from_position(position)

    def _get_func_from_y(self, y):
        position = self._get_position(y)
        func = self.trace.get_func_from_position(position)
        return func
