from pyqtgraph.parametertree import ParameterTree
from pyqtgraph.parametertree.parameterTypes import Parameter, GroupParameter
from collections import deque, OrderedDict, defaultdict
from queue import Queue
from qtpy.QtGui import QStandardItem, QStandardItemModel
from qtpy.QtWidgets import QProgressBar, QWidget, QVBoxLayout
from qtpy.QtCore import QItemSelectionModel, Signal, Qt, QTimer
import sys
import uuid
import datetime
import numpy as np
import warnings
from typing import Iterable, Sequence
from databroker.core import BlueskyRun
from xicam.core import msg, threads
from xicam.gui.patches.PyQtGraph import CounterGroupParameter, LazyGroupParameter, from_dict


# TODO: map list to groupparameter

# TODO: suggest integration of type mapping into pyqtgraph


class MetadataWidgetBase(ParameterTree):
    def __init__(self, *args, **kwargs):
        self.excludedkeys = kwargs.get("excludedkeys")
        if "excludedkeys" in kwargs:
            del kwargs["excludedkeys"]

        super(MetadataWidgetBase, self).__init__(*args, **kwargs)
        LazyGroupParameter.itemClass.initialize_treewidget(self)
        self.kwargs = kwargs

    def insert(self, doctype: str, child, groups: dict):
        group = groups[doctype]
        group.addChildren([child])

    def make_child(self, doctype: str, document: dict, groups: dict):
        if doctype not in groups:
            text = f"doctype '{doctype}' is not supported by MetadataWidget (not in header parameters); skipping"
            msg.logError(KeyError(text))
            return


        # TODO: add responsive design to uid display
        # Attach meaningful names to group name if document has a name (e.g. descriptors)
        try:
            name = MetadataView._lookup_name_for_group(doctype, document)
        except KeyError as ex:
            msg.logError(ex)
            msg.logMessage(f"Insert failed for {doctype} document")
            return
        else:

            return LazyGroupParameter(name=name, children=document, expanded=False,
                                  readonly=True)

    @staticmethod
    def _lookup_name_for_group(doctype, document):
        """Gets the "ID" (name) to be used for creating a GroupParameter for the given doctype."""
        if doctype in ['start', 'event', 'resource', 'stop']:
            return document['uid'][:6]
        elif doctype in ['datum']:
            return f"{document['datum_id'][:6]}...{document['datum_id'][-2:]}"
        elif doctype in ['datum_page']:
            return f"{document['datum_id'][0][:6]}...{document['datum_id'][0][-2:]}"
        elif doctype in ['descriptor']:
            return f"{document['name']} {document['uid'][:6]}"
        elif doctype in ['event_page']:
            return f"{document['uid'][0][:6]}"
        else:
            raise KeyError(f"Cannot find document type '{doctype}' in supported header parameters")


class MetadataWidget(QWidget):
    def __init__(self, *args, **kwargs):
        super(MetadataWidget, self).__init__()
        self.header = None  # type: HeaderParameter
        self.groups = dict()
        self.insert_queue = Queue()
        self.insert_timer = QTimer()
        self.tasks_done = 0

        self.insert_timer.setInterval(100)
        self.insert_timer.timeout.connect(self.insert)

        self.setLayout(QVBoxLayout())
        self.metadata_base = MetadataWidgetBase(*args, **kwargs)
        self.progress_bar = QProgressBar()
        self.layout().addWidget(self.metadata_base)
        self.layout().addWidget(self.progress_bar)
        self.progress_bar.hide()
        self.progress_bar.setRange(0, 0)

        self.reset()

    def make_child(self, name, doc):
        if name == "start":
            self.header.setName(doc["uid"])
        child = self.metadata_base.make_child(name, doc, self.groups)
        return child

    def doc_consumer(self, name, doc):
        child = self.make_child(name, doc)
        self.insert_queue.put((name, child))
        if not self.insert_timer.isActive():
            self.insert_timer.start()
        if not self.progress_bar.isVisible():
            self.progress_bar.show()

    def show_catalog(self, catalog, reset=True):
        if reset:
            self.reset()
        self._show_catalog(catalog)  # must be done as an event to make sure that it happens after reset is complete
        self.progress_bar.show()

    def _show_catalog(self, catalog):
        for name, doc in catalog.documents(fill='no'):
            child = self.make_child(name, doc)
            if child:
                self.insert_queue.put((name, child))
        self.progress_bar.setRange(0, len(self.insert_queue.queue))

    def reset(self):
        with self.insert_queue.mutex:
            self.insert_queue.queue.clear()
        self.header = HeaderParameter(name=" ")
        self.groups = {group.name(): group for group in self.header.children()}
        self.metadata_base.setParameters(self.header)
        self.insert_timer.start()
        self.tasks_done = 0

    def insert(self):
        i = 0
        while not self.insert_queue.empty() and i < 100:
            doctype, child = self.insert_queue.get()
            self.metadata_base.insert(doctype, child, self.groups)
            self.insert_queue.task_done()
            i += 1
            self.tasks_done += 1

        if self.progress_bar.maximum():
            self.progress_bar.setValue(self.tasks_done)

        if self.insert_queue.empty():
            self.insert_timer.stop()
            self.progress_bar.hide()
            self.progress_bar.setRange(0, 0)


class MetadataView(MetadataWidgetBase):
    sigUpdate = Signal()

    def __init__(self, headermodel: QStandardItemModel, selectionmodel: QItemSelectionModel, *args, **kwargs):
        super(MetadataView, self).__init__(*args, **kwargs)
        self._seen = set()
        self._last_uid = None
        self._thread_id = "MetadataView" + str(uuid.uuid4())
        self.headermodel = headermodel
        self.selectionmodel = selectionmodel
        self.selectionmodel.currentChanged.connect(self.update)
        self.sigUpdate.connect(self.update)
        self.update()

    def update(self):
        index = self.selectionmodel.currentIndex()
        if not index.isValid():
            return

        catalog: BlueskyRun = self.headermodel.itemFromIndex(index).data(Qt.UserRole)

        if catalog.name != self._last_uid:
            self._seen = set()

        param = HeaderParameter(name=catalog.name)
        groups = param.groups

        # TODO: make compatible with actively streaming header

        # filter out documents already emitted in the stream
        for doctype, document in catalog.documents(fill='no'):
            uid = document["uid"]

            # for event-page's, make the uid value hashable by joining
            if isinstance(document["uid"], list):
                uid = ', '.join(uid)

            if uid in self._seen:
                continue

            self._seen.add(uid)
            self.insert(doctype, self.make_child(doctype, document, groups), groups)

        if catalog.name != self._last_uid:
            self._last_uid = catalog.name
            # try:
            self.setParameters(param, showTop=False)


class HeaderParameter(GroupParameter):
    def __init__(self, *args, **kwargs):
        super(HeaderParameter, self).__init__(*args, **kwargs)

        self.groups = {
            "start": CounterGroupParameter(name="start", title="Start", expanded=False),
            "descriptor": CounterGroupParameter(name="descriptor", title="Descriptors", expanded=False),
            "event": CounterGroupParameter(name="event", title="Events", expanded=False),
            "event_page": CounterGroupParameter(name="event_page", title="Event Pages", expanded=False),
            "resource": CounterGroupParameter(name="resource", title="Resources", expanded=False),
            "datum": CounterGroupParameter(name="datum", title="Data", expanded=False),
            "datum_page": CounterGroupParameter(name="datum_page", title="Datum Pages", expanded=False),
            "stop": CounterGroupParameter(name="stop", title="Stop", expanded=False),
        }

        self.addChildren(self.groups.values())


class HeaderBuffer:
    def __init__(self, uid, param: GroupParameter, docs=None):
        self.uid = uid
        self.buf = deque()
        self.param = param

        for name, doc in docs:
            self.buf.append((name, doc))

    def stream(self, *args):
        yield from self.buf


class MDVConusumer(MetadataView):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._buffers = {}
        self._descriptor_map = {}

    def doc_consumer(self, name, doc):
        if name == "start":
            uid = doc["uid"]
            item = QStandardItem(doc["uid"])
            item.header = self._buffers[uid] = HeaderBuffer(uid, HeaderParameter(name="uid"))
            self.headermodel.appendRow(item)
            self.selectionmodel.setCurrentIndex(self.headermodel.indexFromItem(item), QItemSelectionModel.ClearAndSelect)
        elif name == "descriptor":
            uid = self._descriptor_map[doc["uid"]] = doc["run_start"]
        elif name == "event":
            uid = self._descriptor_map[doc["descriptor"]]
        elif name == "stop":
            uid = doc["run_start"]
        else:
            raise ValueError

        buffer = self._buffers[uid].buf  # type: deque
        buffer.append((name, doc))
        self.sigUpdate.emit()

    def show_row_n(self, n):
        item = self.headermodel.item(n)
        if item is None:
            return
        self.selectionmodel.setCurrentIndex(self.headermodel.indexFromItem(item), QItemSelectionModel.ClearAndSelect)
        self.update()


if __name__ == "__main__":
    import os

    os.environ["OPHYD_CONTROL_LAYER"] = "caproto"
    from qtpy import QtWidgets
    from qtpy import QtCore
    from qtpy import QtGui
    from mily.runengine import spawn_RE
    from mily.widgets import ControlGui, Count, Scan1D, MotorSelector, DetectorSelector, MISpin
    import bluesky.plans as bp
    from ophyd.sim import hw
    import matplotlib

    matplotlib.interactive(True)

    class MDVWithButtons(QtWidgets.QWidget):
        def __init__(self, mdv, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.mdv = mdv

            w_layout = QtWidgets.QVBoxLayout()
            w_layout.addWidget(mdv)

            button_layout = QtWidgets.QHBoxLayout()
            w_layout.addLayout(button_layout)

            self.spinner = spinner = MISpin("run")
            w_layout.addWidget(spinner)
            spinner.setRange(0, 100)

            def on_change(n):
                self.mdv.show_row_n(n)

            spinner.valueChanged.connect(on_change)
            self.setLayout(w_layout)

        def doc_consumer(self, name, doc):
            self.mdv.doc_consumer(name, doc)
            self.spinner.setRange(0, self.mdv.headermodel.rowCount() - 1)

    app = QtWidgets.QApplication.instance()
    if app is None:
        app = QtWidgets.QApplication([b"Live data demo"])
        app.lastWindowClosed.connect(app.quit)

    headermodel = QtGui.QStandardItemModel()
    selectionmodel = QtCore.QItemSelectionModel(headermodel)

    metadataview = MDVConusumer(headermodel, selectionmodel)
    view_box = MDVWithButtons(metadataview)

    hw = hw()
    hw.motor.set(15)
    hw.motor.delay = 0.1
    hw.motor1.delay = 0.2
    hw.motor2.delay = 0.3

    hw.det.kind = "hinted"
    hw.det1.kind = "hinted"
    hw.det2.kind = "hinted"

    RE, queue, thread, teleport = spawn_RE()

    cg = ControlGui(
        queue,
        teleport,
        Count("Count", bp.count, DetectorSelector(detectors=[hw.det, hw.det1, hw.det2])),
        Scan1D(
            "1D absolute scan",
            bp.scan,
            MotorSelector([hw.motor, hw.motor1, hw.motor2]),
            DetectorSelector(detectors=[hw.det, hw.det1, hw.det2]),
        ),
        Scan1D(
            "1D relative scan",
            bp.rel_scan,
            MotorSelector([hw.motor, hw.motor1, hw.motor2]),
            DetectorSelector(detectors=[hw.det, hw.det1, hw.det2]),
        ),
        live_widget=view_box,
    )

    cg.show()
    sys.exit(app.exec_())
