//======================================================================================================
// OptitrackFaceDevice
//
// Subscribes to one or more analog devices that the iFacialMocapPlugin (Motive plugin)
// publishes via NatNet, and applies blendshape weights + (optionally) head/eye pose to
// SkinnedMeshRenderers / Transforms.
//
// The plugin splits a face stream when total channels exceed NatNet's 32-per-device limit:
//   - <baseName>            (e.g. "iFacialMocap_40001")  - single device, <= 32 channels
//   - <baseName>_A + _B                                  - two devices when > 32 channels
//
// Set DeviceBaseName to the base ("iFacialMocap_40001"); this component automatically
// finds and merges _A and _B if they exist.
//
// Pose channel names produced by the plugin:
//   head_eulerX/Y/Z, head_posX/Y/Z, leftEye_eulerX/Y/Z, rightEye_eulerX/Y/Z
//======================================================================================================

using System;
using System.Collections.Generic;
using UnityEngine;

[DefaultExecutionOrder(-50)]
public class OptitrackFaceDevice : MonoBehaviour
{
    [Tooltip("OptitrackStreamingClient instance to read device data from. Falls back to FindObjectOfType<>() at Start if null.")]
    public OptitrackStreamingClient streamingClient;

    [Tooltip("Motive device base name (e.g. \"iFacialMocap_40001\"). _A/_B suffixes are auto-detected if the device was split.")]
    public string deviceBaseName = "iFacialMocap_40001";

    [Header("Output: Blendshapes")]
    [Tooltip("SkinnedMeshRenderers whose blendshape weights will be driven by incoming channel values. Channel names are matched against blendshape names (case-sensitive, with _L/_R or Left/Right tolerance).")]
    public SkinnedMeshRenderer[] faceMeshRenderers;

    [Tooltip("Multiply incoming blendshape values (0-100 typically) by this scalar before applying.")]
    [Range(0f, 3f)]
    public float blendshapeScale = 1.0f;

    [Header("Output: Head Pose (optional)")]
    [Tooltip("If set, applies head_eulerXYZ to this Transform's localEulerAngles. Leave null to skip.")]
    public Transform headBone;

    [Tooltip("If set, applies leftEye_eulerXYZ to this Transform.")]
    public Transform leftEyeBone;

    [Tooltip("If set, applies rightEye_eulerXYZ to this Transform.")]
    public Transform rightEyeBone;

    [Header("Diagnostics (read-only)")]
    [SerializeField] private bool _connected;
    [SerializeField] private int _resolvedDeviceCount;
    [SerializeField] private int _totalChannelCount;
    [SerializeField] private int _appliedBlendshapeCount;
    [SerializeField] private float _lastJawOpen;
    [SerializeField] private float _lastEyeBlinkLeft;

    // Pre-resolved channel name -> (renderer, blendshape index) for O(1) apply.
    private struct BlendShapeMapping
    {
        public SkinnedMeshRenderer renderer;
        public int index;
    }
    private Dictionary<string, List<BlendShapeMapping>> _blendshapeCache;
    private bool _cacheBuilt;

    // Resolved device names (1 or 2: base, base_A+_B)
    private List<string> _resolvedDeviceNames = new List<string>(2);
    private float _lastResolveAttempt = -10f;
    private const float k_ResolveRetrySeconds = 1.0f;

    void Start()
    {
        if (streamingClient == null)
            streamingClient = FindObjectOfType<OptitrackStreamingClient>();

        if (streamingClient != null && !streamingClient.ReceiveDevices)
        {
            Debug.LogWarning($"[OptitrackFaceDevice] {streamingClient.name}.ReceiveDevices is OFF. " +
                             "Enable it on the streaming client to get face data.", this);
        }

        BuildBlendshapeCache();
    }

    void BuildBlendshapeCache()
    {
        _blendshapeCache = new Dictionary<string, List<BlendShapeMapping>>(StringComparer.OrdinalIgnoreCase);
        if (faceMeshRenderers == null) return;

        foreach (var smr in faceMeshRenderers)
        {
            if (smr == null || smr.sharedMesh == null) continue;
            for (int i = 0; i < smr.sharedMesh.blendShapeCount; ++i)
            {
                string name = smr.sharedMesh.GetBlendShapeName(i);
                AddToCache(name, smr, i);

                // Also accept _L/_R shorthand <-> Left/Right canonical
                if (name.EndsWith("_L"))
                    AddToCache(name.Substring(0, name.Length - 2) + "Left", smr, i);
                else if (name.EndsWith("_R"))
                    AddToCache(name.Substring(0, name.Length - 2) + "Right", smr, i);
                else if (name.EndsWith("Left"))
                    AddToCache(name.Substring(0, name.Length - 4) + "_L", smr, i);
                else if (name.EndsWith("Right"))
                    AddToCache(name.Substring(0, name.Length - 5) + "_R", smr, i);
            }
        }
        _cacheBuilt = true;
    }

    void AddToCache(string key, SkinnedMeshRenderer smr, int idx)
    {
        if (!_blendshapeCache.TryGetValue(key, out var list))
        {
            list = new List<BlendShapeMapping>(1);
            _blendshapeCache[key] = list;
        }
        list.Add(new BlendShapeMapping { renderer = smr, index = idx });
    }

    void TryResolveDevices()
    {
        // Probe for base, base_A, base_B. Devices appear dynamically as the
        // plugin sees its first packet, so retry every k_ResolveRetrySeconds.
        if (Time.unscaledTime - _lastResolveAttempt < k_ResolveRetrySeconds) return;
        _lastResolveAttempt = Time.unscaledTime;

        _resolvedDeviceNames.Clear();
        _totalChannelCount = 0;

        var defs = streamingClient.GetDeviceDefinitions();
        foreach (var d in defs)
        {
            if (d.Name == deviceBaseName ||
                d.Name == deviceBaseName + "_A" ||
                d.Name == deviceBaseName + "_B")
            {
                _resolvedDeviceNames.Add(d.Name);
                _totalChannelCount += d.ChannelCount;
            }
        }
        _resolvedDeviceCount = _resolvedDeviceNames.Count;
        _connected = _resolvedDeviceCount > 0;
    }

    void Update()
    {
        if (streamingClient == null) { _connected = false; return; }
        if (!_cacheBuilt) BuildBlendshapeCache();

        if (_resolvedDeviceNames.Count == 0)
        {
            TryResolveDevices();
            if (_resolvedDeviceNames.Count == 0) return;
        }

        // Pose accumulators (last device wins if both _A and _B somehow have pose
        // — in practice pose lives entirely in one device).
        Vector3? headEuler = null, headPos = null, leftEyeEuler = null, rightEyeEuler = null;

        int appliedBs = 0;
        for (int i = 0; i < _resolvedDeviceNames.Count; ++i)
        {
            var state = streamingClient.GetLatestDeviceState(_resolvedDeviceNames[i]);
            if (state == null) continue;

            // Snapshot to a local list to avoid holding the lock during apply.
            // ChannelValues is a Dictionary updated by the NatNet thread; iterating
            // it while it mutates would throw. The streaming client's GetLatestDeviceState
            // returns the live reference, so we iterate quickly under no lock — accept
            // occasional torn reads (one frame stale value) over the perf cost of
            // copying the dict every frame.
            foreach (var kv in state.ChannelValues)
            {
                string name = kv.Key;
                float v = kv.Value;

                // Pose channels?
                switch (name)
                {
                    case "head_eulerX": headEuler = With(headEuler, 0, v); continue;
                    case "head_eulerY": headEuler = With(headEuler, 1, v); continue;
                    case "head_eulerZ": headEuler = With(headEuler, 2, v); continue;
                    case "head_posX":   headPos = With(headPos, 0, v); continue;
                    case "head_posY":   headPos = With(headPos, 1, v); continue;
                    case "head_posZ":   headPos = With(headPos, 2, v); continue;
                    case "leftEye_eulerX": leftEyeEuler = With(leftEyeEuler, 0, v); continue;
                    case "leftEye_eulerY": leftEyeEuler = With(leftEyeEuler, 1, v); continue;
                    case "leftEye_eulerZ": leftEyeEuler = With(leftEyeEuler, 2, v); continue;
                    case "rightEye_eulerX": rightEyeEuler = With(rightEyeEuler, 0, v); continue;
                    case "rightEye_eulerY": rightEyeEuler = With(rightEyeEuler, 1, v); continue;
                    case "rightEye_eulerZ": rightEyeEuler = With(rightEyeEuler, 2, v); continue;
                }

                // Blendshape — apply scaled value to all matching renderers.
                if (_blendshapeCache != null && _blendshapeCache.TryGetValue(name, out var maps))
                {
                    float weight = Mathf.Clamp(v * blendshapeScale, 0f, 100f);
                    foreach (var m in maps)
                    {
                        if (m.renderer != null)
                            m.renderer.SetBlendShapeWeight(m.index, weight);
                    }
                    appliedBs++;
                }

                // Diagnostics
                if (name == "jawOpen") _lastJawOpen = v;
                if (name == "eyeBlink_L" || name == "eyeBlinkLeft") _lastEyeBlinkLeft = v;
            }
        }
        _appliedBlendshapeCount = appliedBs;

        // Apply pose
        if (headBone != null && headEuler.HasValue)
            headBone.localEulerAngles = headEuler.Value;
        if (leftEyeBone != null && leftEyeEuler.HasValue)
            leftEyeBone.localEulerAngles = leftEyeEuler.Value;
        if (rightEyeBone != null && rightEyeEuler.HasValue)
            rightEyeBone.localEulerAngles = rightEyeEuler.Value;
    }

    // Helper: set component i of a Vector3?, initializing to zero if null.
    private static Vector3 With(Vector3? cur, int axis, float v)
    {
        Vector3 r = cur ?? Vector3.zero;
        if (axis == 0) r.x = v;
        else if (axis == 1) r.y = v;
        else r.z = v;
        return r;
    }
}