"""
skin_weight_blind_data.py
Save / restore skin weights using per-vertex polygon blind data.

Weights are encoded as comma-separated float strings directly on each vertex,
so they travel with the mesh (export, duplicate, reference) without any
external file.

Usage (Maya Script Editor):
    import skin_weight_blind_data as swbd
    swbd.show_ui()

Script Manager (Run command — after the file is sourced):
    show_ui()

    # Or: import skin_weight_blind_data as swbd; swbd.show_ui()

    # Or headless:
    save_skin_weights()   # uses current selection
    load_skin_weights()   # uses current selection
"""

import inspect
import os
import sys
import types

import maya.cmds as cmds
import maya.api.OpenMaya as om
import maya.api.OpenMayaAnim as oma

# ── Configuration ─────────────────────────────────────────────────────────────
# Change TYPE_ID if it collides with another tool's blind data.
TYPE_ID     = 6789
ATTR_JOINTS = "swJoints"   # comma-separated joint names (same value on every vertex)
ATTR_WEIGHT = "swWeights"  # comma-separated weight floats, one per vertex


# ── Internal helpers ──────────────────────────────────────────────────────────

def _get_shape(node):
    if cmds.nodeType(node) == "mesh":
        return node
    shapes = cmds.listRelatives(node, shapes=True, type="mesh", noIntermediate=True) or []
    return shapes[0] if shapes else None


def _find_skin_cluster(mesh):
    hist = cmds.listHistory(mesh, pruneDagObjects=True) or []
    hits = cmds.ls(hist, type="skinCluster")
    return hits[0] if hits else None


def _dag_path(node_name):
    sel = om.MSelectionList()
    sel.add(node_name)
    return sel.getDagPath(0)


def _get_mesh_fn(mesh):
    """Return MFnMesh for the mesh shape."""
    d = _dag_path(mesh)
    d.extendToShape()
    return om.MFnMesh(d)


def _get_skin_fn(mesh):
    """Return MFnSkinCluster for *mesh*, or None."""
    d = _dag_path(mesh)
    d.extendToShape()
    it = om.MItDependencyGraph(
        d.node(),
        om.MFn.kSkinClusterFilter,
        om.MItDependencyGraph.kUpstream,
    )
    return oma.MFnSkinCluster(it.currentNode()) if not it.isDone() else None


def _full_vertex_component(vtx_count):
    """Build an MObject representing all vertices [0 .. vtx_count-1]."""
    fn  = om.MFnSingleIndexedComponent()
    obj = fn.create(om.MFn.kMeshVertComponent)
    fn.setCompleteData(vtx_count)
    return obj


def _ensure_blind_data_type(mesh_fn):
    """
    Declare the blind data type on the mesh before writing.  setStringBlindData
    does NOT auto-create the type definition, so this must run first.
    Idempotent — skips if the type is already present.
    """
    if mesh_fn.isBlindDataTypeUsed(TYPE_ID):
        return
    # API2 takes a single list of (longName, shortName, formatName) tuples.
    mesh_fn.createBlindDataType(
        TYPE_ID,
        [
            (ATTR_JOINTS, "swJ", "string"),
            (ATTR_WEIGHT, "swW", "string"),
        ],
    )


def _resolve_selection(mesh):
    if mesh is not None:
        return _get_shape(mesh) or mesh
    sel = cmds.ls(sl=True, objectsOnly=True)
    return _get_shape(sel[0]) if sel else None


# ── Public API ────────────────────────────────────────────────────────────────

def save_skin_weights(mesh=None):
    """
    Read the skinCluster weights from *mesh* and store them in per-vertex
    blind data.  Falls back to the active selection when mesh is None.

    Data persists in .ma / .mb.  FBX strips blind data; Alembic preserves it.
    """
    mesh = _resolve_selection(mesh)
    if not mesh or cmds.nodeType(mesh) != "mesh":
        cmds.warning("[SWBD] Select a skinned mesh.")
        return

    sc      = _find_skin_cluster(mesh)
    skin_fn = _get_skin_fn(mesh)
    if not sc or skin_fn is None:
        cmds.warning(f"[SWBD] No skinCluster found on '{mesh}'.")
        return

    # Influence names MUST come from the same source as the weight array so the
    # ordering matches.  getWeights() packs weights by physical influence index,
    # i.e. the order of influenceObjects() — this is NOT guaranteed to match
    # `cmds.skinCluster -q -influence`, which could silently mis-pair weights.
    influences = [p.partialPathName() for p in skin_fn.influenceObjects()]
    n_inf      = len(influences)
    vtx_count  = cmds.polyEvaluate(mesh, vertex=True)

    # ── read all weights in one API call ──────────────────────────────────────
    d = _dag_path(mesh)
    d.extendToShape()
    weights, _ = skin_fn.getWeights(d, _full_vertex_component(vtx_count))
    # flat MDoubleArray: [vtx0_inf0, vtx0_inf1, ..., vtxN_infN]

    # ── serialise ─────────────────────────────────────────────────────────────
    # The joint list is identical for every vertex, so store it once (on vertex
    # 0) instead of duplicating it across the whole mesh.
    joint_str   = ",".join(influences)
    weight_list = [
        ",".join(f"{w:.6f}" for w in weights[i * n_inf : (i + 1) * n_inf])
        for i in range(vtx_count)
    ]

    # ── write via MFnMesh ─────────────────────────────────────────────────────
    mesh_fn  = _get_mesh_fn(mesh)
    _ensure_blind_data_type(mesh_fn)
    comp_ids = om.MIntArray(list(range(vtx_count)))
    # Use the scalar single-component overload for the one-off joint list. The
    # sequence overload silently stores "" when given a single component.
    mesh_fn.setStringBlindData(0, om.MFn.kMeshVertComponent,
                               TYPE_ID, ATTR_JOINTS, joint_str)
    mesh_fn.setStringBlindData(comp_ids, om.MFn.kMeshVertComponent, TYPE_ID,
                               ATTR_WEIGHT, weight_list)
    cmds.dgdirty(mesh)

    print(f"[SWBD] Saved {vtx_count} verts x {n_inf} influences on '{mesh}'.")
    cmds.inViewMessage(
        amg=f"Skin weights <hl>saved</hl> ({vtx_count} verts, {n_inf} joints).",
        pos="topCenter", fade=True)


def load_skin_weights(mesh=None):
    """
    Read per-vertex blind data from *mesh* and apply the stored weights to a
    skinCluster.  A new skinCluster is created automatically when none exists.
    All joints stored in the blind data must already exist in the scene.
    """
    mesh = _resolve_selection(mesh)
    if not mesh or cmds.nodeType(mesh) != "mesh":
        cmds.warning("[SWBD] Select a mesh.")
        return

    vtx_count = cmds.polyEvaluate(mesh, vertex=True)
    mesh_fn   = _get_mesh_fn(mesh)

    # ── read via MFnMesh ──────────────────────────────────────────────────────
    # getStringBlindData returns (MIntArray of compIds, list of str values).
    # compIds may not start at 0 or be contiguous, so we index by vertex ID.
    # Read directly rather than gating on hasBlindData (whose result can be
    # unreliable); an empty/raising read is handled below.
    try:
        jnt_ids, jnt_vals = mesh_fn.getStringBlindData(om.MFn.kMeshVertComponent,
                                                        TYPE_ID, ATTR_JOINTS)
        wt_ids,  wt_vals  = mesh_fn.getStringBlindData(om.MFn.kMeshVertComponent,
                                                        TYPE_ID, ATTR_WEIGHT)
    except Exception:
        jnt_ids = jnt_vals = wt_ids = wt_vals = []

    if not jnt_vals or not wt_vals:
        cmds.warning(f"[SWBD] No skin-weight blind data on '{mesh}'. Run Save first.")
        return

    # ── validate vertex count matches what was stored ─────────────────────────
    if len(wt_vals) != vtx_count:
        cmds.warning(
            f"[SWBD] Vertex-count mismatch on '{mesh}': mesh has {vtx_count}, "
            f"stored data has {len(wt_vals)}. Topology differs — aborting.")
        return

    # Re-order by vertex index
    joint_by_vtx  = {jnt_ids[i]: jnt_vals[i] for i in range(len(jnt_ids))}
    weight_by_vtx = {wt_ids[i]:  wt_vals[i]  for i in range(len(wt_ids))}

    weight_data = [weight_by_vtx.get(v, "") for v in range(vtx_count)]

    # The joint list is stored once; use the first non-empty entry rather than
    # assuming vertex 0 carries it.
    joint_str = next((s for s in joint_by_vtx.values() if s), "")
    if not joint_str:
        cmds.warning(f"[SWBD] Stored joint list is empty on '{mesh}'.")
        return
    influences = joint_str.split(",")
    n_inf      = len(influences)

    # ── validate joints ───────────────────────────────────────────────────────
    missing = [j for j in influences if not cmds.objExists(j)]
    if missing:
        cmds.warning(f"[SWBD] Joints missing from scene: {missing}")
        return

    # ── get or create the skinCluster ─────────────────────────────────────────
    sc        = _find_skin_cluster(mesh)
    transform = (cmds.listRelatives(mesh, parent=True) or [None])[0]

    if not sc:
        sc = cmds.skinCluster(
            influences + [transform],
            toSelectedBones=False,
            normalizeWeights=1,
            maximumInfluences=n_inf,
            obeyMaxInfluences=False,
        )[0]
        print(f"[SWBD] Created skinCluster '{sc}'.")
    else:
        current = set(cmds.skinCluster(sc, q=True, influence=True))
        for jnt in influences:
            if jnt not in current:
                cmds.skinCluster(sc, e=True, addInfluence=jnt, weight=0.0)

    # ── map stored joint names → skinCluster influence indices ────────────────
    skin_fn  = _get_skin_fn(mesh)
    sc_paths = skin_fn.influenceObjects()

    name_to_sc_idx = {}
    for k, p in enumerate(sc_paths):
        name_to_sc_idx[p.partialPathName()] = k
        name_to_sc_idx[p.fullPathName()]    = k

    inf_indices = om.MIntArray()
    unresolved  = []
    for jnt in influences:
        idx = name_to_sc_idx.get(jnt)
        if idx is None:
            # Fall back to a leaf-name match, but never silently guess index 0.
            leaf = jnt.split("|")[-1].split(":")[-1]
            idx = next((v for k, v in name_to_sc_idx.items()
                        if k.split("|")[-1].split(":")[-1] == leaf), None)
        if idx is None:
            unresolved.append(jnt)
        else:
            inf_indices.append(idx)

    if unresolved:
        cmds.warning(f"[SWBD] Could not map joints to skinCluster influences: "
                     f"{unresolved}. Aborting to avoid corrupting weights.")
        return

    # ── build flat weight array ───────────────────────────────────────────────
    flat = []
    for ws in weight_data:
        vals = ws.split(",") if ws else []
        row  = [float(v) for v in vals[:n_inf]]
        row += [0.0] * (n_inf - len(row))
        flat.extend(row)

    # ── apply via API ─────────────────────────────────────────────────────────
    d = _dag_path(mesh)
    d.extendToShape()
    skin_fn.setWeights(d, _full_vertex_component(vtx_count),
                       inf_indices, om.MDoubleArray(flat), False)
    cmds.dgdirty(mesh)

    print(f"[SWBD] Applied {vtx_count} x {n_inf} weights on '{mesh}'.")
    cmds.inViewMessage(
        amg=f"Skin weights <hl>loaded</hl> ({vtx_count} verts, {n_inf} joints).",
        pos="topCenter", fade=True)


# ── Optional: clear blind data ────────────────────────────────────────────────

def clear_blind_data(mesh=None):
    """Remove SWBD blind data from *mesh* (does not affect the skinCluster)."""
    mesh = _resolve_selection(mesh)
    if not mesh:
        cmds.warning("[SWBD] Select a mesh.")
        return
    mesh_fn = _get_mesh_fn(mesh)
    try:
        mesh_fn.clearBlindData(om.MFn.kMeshVertComponent, TYPE_ID, ATTR_JOINTS)
        mesh_fn.clearBlindData(om.MFn.kMeshVertComponent, TYPE_ID, ATTR_WEIGHT)
        cmds.dgdirty(mesh)
        print(f"[SWBD] Cleared blind data from '{mesh}'.")
    except Exception as exc:
        cmds.warning(f"[SWBD] Could not clear blind data: {exc}")


# ── UI ────────────────────────────────────────────────────────────────────────

def show_ui():
    WIN = "swbdWindow"
    if cmds.window(WIN, exists=True):
        cmds.deleteUI(WIN)
    # Maya restores a window's saved geometry from windowPref BEFORE showWindow
    # applies the requested size/position. That restore momentarily draws the
    # window at its old location and then it snaps to the new one — the "flash"
    # of a window opening/closing just before the real one appears. Removing the
    # stale pref makes the window open once, in place.
    if cmds.windowPref(WIN, exists=True):
        cmds.windowPref(WIN, remove=True)

    cmds.window(WIN, title="Skin Weight Blind Data",
                widthHeight=(300, 130), sizeable=False)
    main = cmds.columnLayout(adjustableColumn=True, rowSpacing=4,
                              columnOffset=("both", 10))

    cmds.separator(height=8, style="none")
    cmds.text(label="Select a skinned mesh, then Save or Load.",
              align="center", font="smallPlainLabelFont")
    cmds.separator(height=6)

    cmds.rowLayout(numberOfColumns=2,
                   columnWidth2=(130, 130),
                   columnAttach=[(1, "both", 4), (2, "both", 4)])
    cmds.button(label="Save  >>  Blind Data",
                height=38, backgroundColor=(0.24, 0.44, 0.64),
                command=lambda *_: save_skin_weights())
    cmds.button(label="Load  <<  Blind Data",
                height=38, backgroundColor=(0.28, 0.54, 0.28),
                command=lambda *_: load_skin_weights())
    cmds.setParent(main)

    cmds.separator(height=4)
    cmds.button(label="Clear Blind Data from Mesh",
                height=26, backgroundColor=(0.45, 0.25, 0.25),
                command=lambda *_: clear_blind_data())

    cmds.separator(height=8, style="none")
    cmds.showWindow(WIN)


def _this_script_path():
    path = globals().get("__file__")
    if path:
        return os.path.abspath(path)
    for fr in inspect.stack():
        fn = fr.filename
        if fn and os.path.basename(fn).lower() == "skin_weight_blind_data.py":
            return os.path.abspath(fn)
    return None


def _running_under_script_manager():
    for fr in inspect.stack():
        if fr.function == "_source_py" and fr.filename and "ScriptManager" in fr.filename:
            return True
    return False


def _register_module():
    """Let ``import skin_weight_blind_data`` work when Script Manager sources this file."""
    path = _this_script_path()
    if not path:
        return

    root = os.path.dirname(path)
    if root not in sys.path:
        sys.path.insert(0, root)

    name = "skin_weight_blind_data"
    existing = sys.modules.get(name)
    if existing is not None and getattr(existing, "__file__", None) == path:
        return

    mod = types.ModuleType(name)
    mod.__file__ = path
    mod.__package__ = ""
    for key, value in globals().items():
        if key.startswith("__"):
            continue
        setattr(mod, key, value)
    sys.modules[name] = mod


_register_module()

if __name__ == "__main__" and not _running_under_script_manager():
    show_ui()