"""PySide2 tool for mapping Maya joints onto the 137 OpenPose_full keypoint
slots of a selected ``openposeCharacter`` node.

Mapping edits go through ``cmds.connectAttr``/``disconnectAttr`` (not raw
API calls) so every edit is a normal, undoable Maya scene operation.
"""

from functools import partial

import maya.cmds as cmds
import maya.OpenMayaUI as omui

try:
    from PySide2 import QtCore, QtWidgets
    from shiboken2 import wrapInstance
except ImportError:
    from PySide6 import QtCore, QtWidgets
    from shiboken6 import wrapInstance

from .. import constants
from .. import face_blendshapes
from .. import mapping_node

_PART_LABELS = {
    constants.PART_BODY: "Body",
    constants.PART_HAND_LEFT: "Left Hand",
    constants.PART_HAND_RIGHT: "Right Hand",
    constants.PART_FACE: "Face",
}

_window_instance = None


def _maya_main_window():
    ptr = omui.MQtUtil.mainWindow()
    return wrapInstance(int(ptr), QtWidgets.QWidget)


def _target_plug(character, index):
    return "{0}.targetJoint[{1}]".format(character, index)


class MappingEditor(QtWidgets.QDialog):

    def __init__(self, parent=None):
        super(MappingEditor, self).__init__(parent or _maya_main_window())
        self.setWindowTitle("OpenPose Mapping Editor")
        self.setWindowFlags(self.windowFlags() & ~QtCore.Qt.WindowContextHelpButtonHint)
        self.resize(620, 720)

        self._row_items = {}  # index -> QTreeWidgetItem
        self._build_ui()
        self.refresh_characters()

    def _build_ui(self):
        layout = QtWidgets.QVBoxLayout(self)

        char_row = QtWidgets.QHBoxLayout()
        char_row.addWidget(QtWidgets.QLabel("Character:"))
        self.character_combo = QtWidgets.QComboBox()
        self.character_combo.currentIndexChanged.connect(self.refresh_mapping)
        char_row.addWidget(self.character_combo, 1)

        refresh_btn = QtWidgets.QPushButton("Refresh")
        refresh_btn.clicked.connect(self.refresh_characters)
        char_row.addWidget(refresh_btn)

        create_btn = QtWidgets.QPushButton("Create Character")
        create_btn.clicked.connect(self._on_create_character)
        char_row.addWidget(create_btn)

        char_row.addWidget(QtWidgets.QLabel("Body Model:"))
        self.body_model_combo = QtWidgets.QComboBox()
        self.body_model_combo.addItems(constants.BODY_MODELS)
        self.body_model_combo.currentIndexChanged.connect(self._on_body_model_changed)
        char_row.addWidget(self.body_model_combo)

        layout.addLayout(char_row)

        face_row = QtWidgets.QHBoxLayout()
        face_row.addWidget(QtWidgets.QLabel("Face Source:"))
        self.face_source_combo = QtWidgets.QComboBox()
        self.face_source_combo.addItems(
            [mapping_node.FACE_SOURCE_JOINTS, mapping_node.FACE_SOURCE_ARKIT_BLENDSHAPES]
        )
        self.face_source_combo.currentIndexChanged.connect(self._on_face_source_changed)
        face_row.addWidget(self.face_source_combo)

        face_row.addWidget(QtWidgets.QLabel("Face Scale:"))
        self.face_scale_spin = QtWidgets.QDoubleSpinBox()
        self.face_scale_spin.setRange(0.0, 100000.0)
        self.face_scale_spin.setValue(mapping_node.DEFAULT_FACE_SCALE)
        self.face_scale_spin.valueChanged.connect(self._on_face_scale_changed)
        face_row.addWidget(self.face_scale_spin)

        autoconnect_btn = QtWidgets.QPushButton("Auto-Connect ARKit Blendshapes from Selected")
        autoconnect_btn.setToolTip(
            "Connects the selected node's attributes to this character's 52 blendshapeWeight\n"
            "slots wherever an attribute matching the canonical ARKit name (e.g. jawOpen,\n"
            "mouthSmileLeft) exists on it -- typically a blendShape node whose targets are\n"
            "named per the ARKit convention, or an ARKit-mocap-driven control."
        )
        autoconnect_btn.clicked.connect(self._on_autoconnect_blendshapes)
        face_row.addWidget(autoconnect_btn)

        self.blendshape_status_label = QtWidgets.QLabel("")
        face_row.addWidget(self.blendshape_status_label)
        face_row.addStretch(1)

        layout.addLayout(face_row)

        self.tree = QtWidgets.QTreeWidget()
        self.tree.setColumnCount(3)
        self.tree.setHeaderLabels(["Keypoint", "Mapped Joint", ""])
        self.tree.setColumnWidth(0, 220)
        self.tree.setColumnWidth(1, 180)
        layout.addWidget(self.tree, 1)

        self._build_rows()

    def _build_rows(self):
        self.tree.clear()
        self._row_items = {}

        part_items = {}
        for part_id, _start, _count in constants.PARTS:
            part_item = QtWidgets.QTreeWidgetItem([_PART_LABELS[part_id], "", ""])
            self.tree.addTopLevelItem(part_item)
            part_items[part_id] = part_item

        for index in range(constants.TOTAL_KEYPOINTS):
            part_id = constants.part_of(index)
            row = QtWidgets.QTreeWidgetItem([constants.KEYPOINT_NAMES[index], "<unmapped>", ""])
            part_items[part_id].addChild(row)

            button_widget = QtWidgets.QWidget()
            button_layout = QtWidgets.QHBoxLayout(button_widget)
            button_layout.setContentsMargins(0, 0, 0, 0)
            set_btn = QtWidgets.QPushButton("Set from Selected")
            set_btn.clicked.connect(partial(self._on_set_joint, index))
            clear_btn = QtWidgets.QPushButton("Clear")
            clear_btn.clicked.connect(partial(self._on_clear_joint, index))
            button_layout.addWidget(set_btn)
            button_layout.addWidget(clear_btn)
            self.tree.setItemWidget(row, 2, button_widget)

            self._row_items[index] = row

        self.tree.expandAll()

    def _current_character(self):
        return self.character_combo.currentText() or None

    def refresh_characters(self):
        current = self._current_character()
        self.character_combo.blockSignals(True)
        self.character_combo.clear()
        characters = sorted(cmds.ls(type="openposeCharacter") or [])
        self.character_combo.addItems(characters)
        if current in characters:
            self.character_combo.setCurrentText(current)
        self.character_combo.blockSignals(False)
        self.refresh_mapping()

    def refresh_mapping(self):
        character = self._current_character()
        for index, row in self._row_items.items():
            label = "<unmapped>"
            if character:
                plug = _target_plug(character, index)
                sources = cmds.listConnections(plug, source=True, destination=False) or []
                if sources:
                    label = sources[0]
            row.setText(1, label)

        self.body_model_combo.blockSignals(True)
        if character:
            self.body_model_combo.setEnabled(True)
            self.body_model_combo.setCurrentIndex(cmds.getAttr(character + ".bodyModel"))
        else:
            self.body_model_combo.setEnabled(False)
        self.body_model_combo.blockSignals(False)

        for widget in (self.face_source_combo, self.face_scale_spin):
            widget.blockSignals(True)
        if character:
            self.face_source_combo.setEnabled(True)
            self.face_scale_spin.setEnabled(True)
            self.face_source_combo.setCurrentIndex(cmds.getAttr(character + ".faceSource"))
            self.face_scale_spin.setValue(cmds.getAttr(character + ".faceScale"))
            connected = sum(
                1 for i in range(face_blendshapes.BLENDSHAPE_COUNT)
                if cmds.listConnections("{0}.blendshapeWeight[{1}]".format(character, i), source=True, destination=False)
            )
            self.blendshape_status_label.setText(
                "{0} / {1} blendshapes connected".format(connected, face_blendshapes.BLENDSHAPE_COUNT)
            )
        else:
            self.face_source_combo.setEnabled(False)
            self.face_scale_spin.setEnabled(False)
            self.blendshape_status_label.setText("")
        for widget in (self.face_source_combo, self.face_scale_spin):
            widget.blockSignals(False)

    def _on_body_model_changed(self, index):
        character = self._current_character()
        if character:
            cmds.setAttr(character + ".bodyModel", index)

    def _on_face_source_changed(self, index):
        character = self._current_character()
        if character:
            cmds.setAttr(character + ".faceSource", index)

    def _on_face_scale_changed(self, value):
        character = self._current_character()
        if character:
            cmds.setAttr(character + ".faceScale", value)

    def _on_autoconnect_blendshapes(self):
        character = self._current_character()
        if not character:
            QtWidgets.QMessageBox.warning(self, "OpenPose", "Create or select a character first.")
            return
        selected = cmds.ls(selection=True)
        if not selected:
            QtWidgets.QMessageBox.warning(
                self, "OpenPose", "Select a blendShape node (or control) with ARKit-named attributes first."
            )
            return
        source_node = selected[0]

        connected_count = 0
        for index, name in enumerate(face_blendshapes.ARKIT_BLENDSHAPE_NAMES):
            if not cmds.attributeQuery(name, node=source_node, exists=True):
                continue
            plug = "{0}.blendshapeWeight[{1}]".format(character, index)
            existing = cmds.listConnections(plug, source=True, destination=False, plugs=True) or []
            for src in existing:
                cmds.disconnectAttr(src, plug)
            cmds.connectAttr("{0}.{1}".format(source_node, name), plug, force=True)
            connected_count += 1

        self.refresh_mapping()
        QtWidgets.QMessageBox.information(
            self, "OpenPose",
            "Connected {0} / {1} ARKit blendshapes from '{2}'.".format(
                connected_count, face_blendshapes.BLENDSHAPE_COUNT, source_node
            ),
        )

    def _on_create_character(self):
        name, ok = QtWidgets.QInputDialog.getText(
            self, "Create Character", "Name:", text="openposeCharacter1"
        )
        if not ok:
            return
        created = cmds.openposeCreateCharacter(name=name) if name else cmds.openposeCreateCharacter()
        if isinstance(created, list):  # MPxCommand string results come back as a 1-item list
            created = created[0]
        self.refresh_characters()
        combo_index = self.character_combo.findText(created)
        if combo_index >= 0:
            self.character_combo.setCurrentIndex(combo_index)

    def _on_set_joint(self, index):
        character = self._current_character()
        if not character:
            QtWidgets.QMessageBox.warning(self, "OpenPose", "Create or select a character first.")
            return
        selected = cmds.ls(selection=True, type="joint")
        if not selected:
            QtWidgets.QMessageBox.warning(self, "OpenPose", "Select a joint first.")
            return
        joint = selected[0]
        plug = _target_plug(character, index)
        existing = cmds.listConnections(plug, source=True, destination=False, plugs=True) or []
        for src in existing:
            cmds.disconnectAttr(src, plug)
        cmds.connectAttr(joint + ".message", plug, force=True)
        self.refresh_mapping()

    def _on_clear_joint(self, index):
        character = self._current_character()
        if not character:
            return
        plug = _target_plug(character, index)
        existing = cmds.listConnections(plug, source=True, destination=False, plugs=True) or []
        for src in existing:
            cmds.disconnectAttr(src, plug)
        self.refresh_mapping()


def show():
    global _window_instance
    if _window_instance is not None:
        try:
            _window_instance.close()
            _window_instance.deleteLater()
        except RuntimeError:
            pass
    _window_instance = MappingEditor()
    _window_instance.show()
    return _window_instance
