import os

from pymol.Qt import QtGui, QtCore, QtWidgets
from pymol.Qt.utils import UpdateLock, PopupOnException
import pymol
Qt = QtCore.Qt

from pymol.setting import name_dict

# Functions to convert atom property values to strings for display in the GUI
strfunctions = {
    'color': lambda c: hex(c) if (c >= 0x40000000) else str(c),
    'reps': bin,
    'flags': bin,
}

class UneditableDelegate(QtWidgets.QStyledItemDelegate):
    def createEditor(self, parent, option, index):
        return None

class FunctionSuspender:
    def __init__(self, func):
        self.func = func
    def __enter__(self):
        self.func.suspended = True
    def __exit__(self, exc_type, exc_value, traceback):
        self.func.suspended = False

def suspendable(func):
    def wrapper(*args, **kwargs):
        if not func.suspended:
            return func(*args, **kwargs)
    func.suspended = False
    wrapper.suspend = FunctionSuspender(func)
    return wrapper

def get_object_names(_self):
    # was _self.get_names('public_objects') in PyMOL 2.1
    # but that throws exceptions for groups/isosurfaces/etc.
    names = _self.get_object_list()
    return names

class PropsDialog(QtWidgets.QWidget):
    def __init__(self, parent):
        QtWidgets.QWidget.__init__(self)
        self.cmd = parent.cmd
        self.form = parent.load_form('props')
        self.setup_tree_widget()
        self.setup_behavior()
        self.item_changed_skip = False

    def make_entry(self, parent, label):
        item = QtWidgets.QTreeWidgetItem(parent)
        item.setText(0, str(label))
        item.setFlags(QtCore.Qt.ItemFlag.ItemIsEditable |
                      QtCore.Qt.ItemFlag.ItemIsEnabled |
                      QtCore.Qt.ItemFlag.ItemIsSelectable )
        return item

    def make_cat(self, parent, label):
        item = QtWidgets.QTreeWidgetItem(parent)
        item.setText(0, str(label))
        item.setFirstColumnSpanned(True)
        item.setExpanded(True)
        item.setChildIndicatorPolicy(
            QtWidgets.QTreeWidgetItem.ChildIndicatorPolicy.ShowIndicator)
        return item

    def setup_tree_widget(self):
        self.form.treeWidget.setItemDelegateForColumn(0, UneditableDelegate())

        self.item_object = self.make_cat(self.form.treeWidget, "Object-Level")
        self.item_object_ttt = self.make_entry(self.item_object, "TTT Matrix")
        self.item_object_settings = self.make_cat(self.item_object, "Settings")

        self.item_ostate = self.make_cat(self.form.treeWidget, "Object-State-Level")
        self.item_ostate_title = self.make_entry(self.item_ostate, "Title")
        self.item_ostate_matrix = self.make_entry(self.item_ostate, "State Matrix")
        self.item_ostate_settings = self.make_cat(self.item_ostate, "Settings")
        self.item_ostate_properties = Ellipsis  # Incentive PyMOL only

        self.item_atom = self.make_cat(self.form.treeWidget, "Atom-Level")
        self.item_atom_identifiers = self.make_cat(self.item_atom, "Identifiers")
        self.item_atom_builtins = self.make_cat(self.item_atom, "Properties (built-in)")
        self.item_atom_settings = self.make_cat(self.item_atom, "Settings")
        self.item_atom_properties = Ellipsis  # Incentive PyMOL only

        self.item_astate = self.make_cat(self.form.treeWidget, "Atom-State-Level")
        self.item_astate_builtins = self.make_cat(self.item_astate, "Properties (built-in)")
        self.item_astate_settings = self.make_cat(self.item_astate, "Settings")

        self.keys_atom_identifiers = ['model', 'index', 'segi', 'chain', 'resi',
                                'resn',
                                'oneletter',  # read-only
                                'name', 'alt', 'ID', 'rank']
        self.keys_atom_builtins = ['elem', 'q', 'b', 'type', 'formal_charge',
                            'partial_charge', 'numeric_type', 'text_type',
                            # avoid stereo auto-assignment errors
                            # 'stereo',
                            'vdw', 'ss', 'color', 'reps',
                            'flags',
                            'label', 'cartoon',
                            'protons', 'geom', 'valence', 'elec_radius']
        self.keys_astate_builtins = ['state', 'x', 'y', 'z']

        self.items = {}
        for key in self.keys_atom_identifiers:
            self.items[key] = self.make_entry(self.item_atom_identifiers, key)
        for key in self.keys_atom_builtins:
            self.items[key] = self.make_entry(self.item_atom_builtins, key)
        for key in self.keys_astate_builtins:
            self.items[key] = self.make_entry(self.item_astate_builtins, key)

        self.items['model'].setDisabled(True)
        self.items['index'].setDisabled(True)
        self.items['state'].setDisabled(True)
        self.items['oneletter'].setDisabled(True)

    def setup_behavior(self):
        # init input fields
        self.form.input_model.addItems(get_object_names(self.cmd))
        self.form.input_state.setValue(self.cmd.get_state())

        # select pk1 atom if available
        self.update_from_pk1()

        # Install an eventfilter
        self.form.treeWidget.installEventFilter(self)

        # hook up events
        self.form.input_model.currentIndexChanged.connect(self.update_treewidget_model)
        self.form.input_state.valueChanged.connect(self.update_treewidget_state)
        self.form.input_index.valueChanged.connect(self.update_treewidget)
        self.form.button_refresh.clicked.connect(self.update_model_list)

        # themed icons only available by default on X11
        if self.form.button_refresh.icon().isNull():
            self.form.button_refresh.setIcon(QtGui.QIcon(
                os.path.expandvars('$PYMOL_DATA/pmg_qt/icons/refresh.svg')))

        # update and show
        self.update_treewidget_model()
        self.form.treeWidget.setColumnWidth(0, 200)

        self.form.treeWidget.itemChanged.connect(self.item_changed)

    def get_dialog(self):
        return self.form._dialog

    def item_changed(self, item, column):
        """
        Edits current item.
        """

        if self.item_changed_skip:
            return

        model = self.form.input_model.currentText()
        state = self.form.input_state.value()
        key = item.text(0)
        new_value = item.text(1)
        parent = item.parent()
        result = False

        if not result:
            if item is self.item_object_ttt:
                try:
                    if new_value:
                        result = self.cmd.set_object_ttt(model, new_value)
                except (ValueError, IndexError):
                    result = False
            elif item is self.item_ostate_title:
                result = self.cmd.set_title(model, state, new_value)
            elif item is self.item_ostate_matrix:
                self.cmd.matrix_reset(model, state)
                try:
                    new_value = self.cmd.safe_eval(new_value)
                    result = self.cmd.transform_object(model, new_value, state)
                except: # CmdTransformObject-DEBUG: bad matrix
                    result = False
            elif parent is self.item_object_settings:
                with PopupOnException():
                    self.cmd.set(key, new_value, model, quiet=0)
            elif parent is self.item_ostate_settings:
                with PopupOnException():
                    self.cmd.set(key, new_value, model, state, quiet=0)
            elif parent is self.item_ostate_properties:
                self.cmd.set_property(key, new_value, model, state, quiet=0)
            else:
                is_state = False

                if parent is self.item_atom_properties:
                    key = 'p.' + key
                elif parent is self.item_atom_settings:
                    key = 's.' + key
                elif parent is self.item_astate_settings:
                    key = 's.' + key
                    is_state = True
                elif key in self.keys_astate_builtins:
                    is_state = True

                def get_new_value(old_value):
                    if isinstance(old_value, (tuple, list, bool)):
                        return self.cmd.safe_eval(new_value)
                    try:
                        # cast to old type (required for e.g. 'resv = "3"')
                        if isinstance(old_value, int):
                            return int(new_value, 0)
                        return type(old_value)(new_value)
                    except ValueError:
                        # return str and let PyMOL handle it (e.g. change
                        # type of user property)
                        return new_value.encode('utf-8')

                alter_args = ('pk1', key + '= get_new_value(' + key + ')', 0,
                        {'get_new_value': get_new_value})

                with PopupOnException():
                    if is_state:
                        result = self.cmd.alter_state(state, *alter_args)
                    else:
                        result = self.cmd.alter(*alter_args)

        if not result:
            self.update_treewidget_model()

        self.item_changed_skip = False

    def unset_item(self, item):
        '''
        Calls unset on item if applicable. Matrices will be reset.
        '''
        model = self.form.input_model.currentText()
        state = self.form.input_state.value()
        key = item.text(0)
        old_value = item.text(1)
        old_value_type = type(old_value)
        parent = item.parent()

        unset_result = True

        if old_value == "":
            unset_result = False
            return unset_result

        if item is self.item_object_ttt:
            self.cmd.matrix_reset(model, mode=1) # mode=1 for TTT matrix reset
        elif item is self.item_ostate_title:
            self.cmd.set_title(model, state, '')
        elif item is self.item_ostate_matrix:
            self.cmd.matrix_reset(model, state, mode=2) # mode=2 for state matrix reset
        elif parent is self.item_object_settings:
            self.cmd.unset(key, model, quiet=0)
        elif parent is self.item_ostate_settings:
            self.cmd.unset(key, model, state, quiet=0)
        elif parent is self.item_ostate_properties:
            self.cmd.set_property(key, None, model, state, quiet=0)
        elif parent is self.item_atom_properties:
            key = 'p.' + key
            alter_args = ('pk1', key + '= None', 0)
            unset_result = self.cmd.alter(*alter_args)
        elif parent is self.item_atom_settings:
            key = 's.' + key
            alter_args = ('pk1', key + '= None', 0)
            unset_result = self.cmd.alter(*alter_args)
        else:
            unset_result = False

        if unset_result:
            self.update_treewidget_model()

        return unset_result

    def unset_caller(self):
        item_list_selected = self.form.treeWidget.selectedItems()
        self.unset_item(item_list_selected[0])

    def eventFilter(self, source, event):
        '''
        Event filter for creating new shortcuts. Processes the key event before passing it on.
        '''
        if (event.type() == QtCore.QEvent.Type.KeyPress and source is self.form.treeWidget):
            if (event.key() == QtCore.Qt.Key.Key_Delete):
                self.unset_caller()
                return 0
        return super().eventFilter(source, event)

    def update_object_settings(self, parent, model, state):
        parent.takeChildren()
        for sitem in (self.cmd.get_object_settings(model, state) or []):
            key = name_dict.get(sitem[0], sitem[0])
            item = self.make_entry(parent, key)
            item.setText(1, str(sitem[2]))

    def update_atom_settings(self, wrapper, parent):
        parent.takeChildren()
        for key in wrapper:
            item = self.make_entry(parent, name_dict.get(key, key))
            value = wrapper[key]
            item.setText(1, str(value))

    def update_atom_properties(self, wrapper, parent):
        parent.takeChildren()
        for key in wrapper:
            item = self.make_entry(parent, key)
            value = wrapper[key]
            item.setText(1, str(value))

    def update_atom_fields(self, ns):
        for key in self.keys_atom_identifiers + self.keys_atom_builtins:
            try:
                value = ns[key]
            except Exception as e:
                value = 'ERROR: ' + str(e)
            self.items[key].setText(1, strfunctions.get(key, str)(value))
        self.update_atom_settings(ns['s'], self.item_atom_settings)

    def update_astate_fields(self, ns):
        for key in self.keys_astate_builtins:
            value = ns[key]
            self.items[key].setText(1, str(value))
        self.update_atom_settings(ns['s'], self.item_astate_settings)

    def update_from_pk1(self):
        pk1_atom = []
        if self.cmd.iterate('?pk1', 'pk1_atom[:] = [model, index]', space=locals()) > 0:
            self.form.input_model.setCurrentIndex(self.form.input_model.findText(pk1_atom[0]))
            self.form.input_index.setValue(pk1_atom[1])

    def update_pk1(self):
        model = self.form.input_model.currentText()
        index = self.form.input_index.value()

        if model and index:
            try:
                self.cmd.edit((model, index))
                return True
            except pymol.CmdException:
                pass

        return False

    def update_treewidget(self, *args):
        if not self.update_pk1():
            return

        state = self.form.input_state.value()

        self.item_changed_skip = True
        space = {'update_atom_fields': self.update_atom_fields,'update_astate_fields': self.update_astate_fields,'locals': locals}
        count = self.cmd.iterate(
            'pk1', 'update_atom_fields(locals())', space=space)
        self.item_atom.setDisabled(not count)
        if count:
            count = self.cmd.iterate_state(
                state, 'pk1',
                'update_astate_fields(locals())', space=space)
        self.item_astate.setDisabled(not count)
        self.item_changed_skip = False

    def update_treewidget_state(self, *args):
        model = self.form.input_model.currentText()
        state = self.form.input_state.value()
        if not (model and state):
            return

        self.item_changed_skip = True
        self.item_ostate_title.setText(
            1, str(self.cmd.get_title(model, state) or ''))
        self.item_ostate_matrix.setText(
            1, str(
                self.cmd.get_object_matrix(
                    model, state, 0) or ''))

        self.update_object_settings(self.item_ostate_settings, model, state)

        self.update_treewidget()
        self.item_changed_skip = False

    @suspendable
    def update_treewidget_model(self, *args):
        self.item_changed_skip = True
        model = self.form.input_model.currentText()

        if not model:
            return

        self.item_object_ttt.setText(1, str(self.cmd.get_object_ttt(model) or ''))
        self.update_object_settings(self.item_object_settings, model, 0)

        natoms = self.cmd.count_atoms('?' + model)
        nstates = self.cmd.count_states('?' + model)

        self.form.input_state.setMinimum(1)
        self.form.input_state.setMaximum(nstates)
        self.form.input_index.setMinimum(1)
        self.form.input_index.setMaximum(natoms)

        self.item_atom.setHidden(natoms == 0)
        self.item_astate.setHidden(natoms == 0)
        self.item_ostate.setHidden(nstates == 0)

        self.update_treewidget_state()
        self.item_changed_skip = False

    def update_model_list(self, *args):
        if 'pk1' not in self.cmd.get_names('selections'):
            self.update_pk1()

        with self.update_treewidget_model.suspend:
            self.form.input_model.clear()
            self.form.input_model.addItems(get_object_names(self.cmd))
            self.update_from_pk1()

        self.update_treewidget_model()
