/*
 * NeonAvatarDriver.cs — Face Tracking Driver for Unity
 * =====================================================
 * Receives OpenSeeFace tracking data from Z-Ray Facetracker (UDP 11573)
 * and maps it to a standard SkinnedMeshRenderer with ARKit blendshapes.
 * 
 * No VRM needed — works with any avatar that has ARKit-compatible blendshapes.
 * 
 * Setup:
 * 1. Import avatar FBX with blendshapes into Unity
 * 2. Add this script to a GameObject
 * 3. Assign the SkinnedMeshRenderer of the avatar's head
 * 4. Make sure Z-Ray Facetracker is running (flatpak run de.z_ray.Facetracker)
 */

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

public class NeonAvatarDriver : MonoBehaviour
{
    [Header("References")]
    [Tooltip("The OpenSeeFace receiver component")]
    public OpenSee openSeeReceiver;
    
    [Tooltip("SkinnedMeshRenderer of the avatar head with blendshapes")]
    public SkinnedMeshRenderer headMesh;
    
    [Tooltip("SkinnedMeshRenderer of the avatar body (optional, for body movement)")]
    public SkinnedMeshRenderer bodyMesh;
    
    [Header("Head Movement")]
    [Tooltip("How much the avatar head rotates with the webcam tracking")]
    public float headRotationScale = 1.0f;
    
    [Tooltip("Smooth factor for head rotation (higher = smoother)")]
    public float headSmoothing = 5.0f;
    
    [Tooltip("Transform of the head bone (for rotation)")]
    public Transform headBone;
    
    [Header("Eye Tracking")]
    public bool enableEyeTracking = true;
    public float eyeRotationScale = 1.0f;
    
    [Header("Blendshape Settings")]
    [Tooltip("Prefix for blendshape names (empty if using ARKit names directly)")]
    public string blendshapePrefix = "";
    
    [Tooltip("Multiplier for mouth blendshapes")]
    public float mouthScale = 1.5f;
    
    [Tooltip("Multiplier for brow blendshapes")]
    public float browScale = 1.2f;
    
    [Tooltip("Multiplier for eye blendshapes")]
    public float eyeScale = 1.0f;
    
    [Header("Idle Animation")]
    [Tooltip("Enable idle breathing animation")]
    public bool enableBreathing = true;
    
    [Tooltip("Breathing speed")]
    public float breathingSpeed = 0.3f;
    
    [Tooltip("Breathing intensity")]
    public float breathingIntensity = 0.015f;
    
    [Header("Auto Blink")]
    public bool enableAutoBlink = true;
    public float minBlinkInterval = 3.0f;
    public float maxBlinkInterval = 8.0f;
    public float blinkSpeed = 0.15f;
    
    // Private state
    private Dictionary<string, int> blendshapeMap = new Dictionary<string, int>();
    private Quaternion smoothedHeadRotation = Quaternion.identity;
    private float breathingTimer = 0f;
    private float nextBlinkTime = 0f;
    private float blinkState = 0f;
    private bool blinking = false;
    
    // ARKit blendshape names (standard 52)
    private static readonly string[] ARKIT_SHAPES = {
        "EyeBlinkLeft", "EyeLookInLeft", "EyeLookOutLeft", "EyeLookUpLeft", "EyeLookDownLeft",
        "EyeBlinkRight", "EyeLookInRight", "EyeLookOutRight", "EyeLookUpRight", "EyeLookDownRight",
        "EyeSquintLeft", "EyeSquintRight", "EyeWideLeft", "EyeWideRight",
        "JawOpen", "JawLeft", "JawRight", "JawForward",
        "MouthClose", "MouthFunnel", "MouthPucker", "MouthLeft", "MouthRight",
        "MouthSmileLeft", "MouthSmileRight", "MouthFrownLeft", "MouthFrownRight",
        "MouthDimpleLeft", "MouthDimpleRight", "MouthStretchLeft", "MouthStretchRight",
        "MouthRollLower", "MouthRollUpper", "MouthShrugLower", "MouthShrugUpper",
        "MouthPressLeft", "MouthPressRight", "MouthLowerOuterLeft", "MouthLowerOuterRight",
        "MouthUpperOuterLeft", "MouthUpperOuterRight",
        "CheekSquintLeft", "CheekSquintRight", "CheekPuff",
        "NoseSneerLeft", "NoseSneerRight",
        "BrowInnerUp", "BrowInnerDown",
        "BrowOuterUpLeft", "BrowOuterUpRight", "BrowOuterDownLeft", "BrowOuterDownRight",
    };
    
    void Start()
    {
        // Build blendshape name → index map
        if (headMesh != null && headMesh.sharedMesh != null)
        {
            var mesh = headMesh.sharedMesh;
            for (int i = 0; i < mesh.blendShapeCount; i++)
            {
                string name = mesh.GetBlendShapeName(i);
                blendshapeMap[name] = i;
            }
            
            Debug.Log($"[NeonAvatar] Found {blendshapeMap.Count} blendshapes on {headMesh.name}");
            
            // Log which ARKit shapes are present
            int found = 0;
            foreach (var shape in ARKIT_SHAPES)
            {
                if (blendshapeMap.ContainsKey(blendshapePrefix + shape))
                    found++;
            }
            Debug.Log($"[NeonAvatar] {found}/{ARKIT_SHAPES.Length} ARKit blendshapes found");
        }
        
        // Set next blink time
        ScheduleNextBlink();
    }
    
    void Update()
    {
        // Breathing
        if (enableBreathing && bodyMesh != null)
        {
            breathingTimer += Time.deltaTime * breathingSpeed;
            float breath = Mathf.Sin(breathingTimer * 2f * Mathf.PI) * breathingIntensity;
            // Scale chest slightly
            bodyMesh.transform.localScale = new Vector3(
                bodyMesh.transform.localScale.x,
                bodyMesh.transform.localScale.y + breath * 0.01f,
                bodyMesh.transform.localScale.z
            );
        }
        
        // Auto blink
        if (enableAutoBlink)
        {
            HandleAutoBlink();
        }
        
        // Process face tracking data
        if (openSeeReceiver == null || openSeeReceiver.trackingData == null || openSeeReceiver.trackingData.Length == 0)
            return;
        
        var data = openSeeReceiver.trackingData[0];
        if (data == null)
            return;
        
        // --- Head Rotation ---
        if (headBone != null)
        {
            // OpenSeeFace provides head rotation as quaternion
            Quaternion targetRot = data.rotation;
            targetRot = Quaternion.Euler(
                -targetRot.x * headRotationScale,
                targetRot.y * headRotationScale,
                -targetRot.z * headRotationScale
            );
            
            smoothedHeadRotation = Quaternion.Slerp(smoothedHeadRotation, targetRot, Time.deltaTime * headSmoothing);
            headBone.localRotation = smoothedHeadRotation;
        }
        
        // --- Eye Open/Close ---
        if (enableEyeTracking)
        {
            SetBlendshape("EyeBlinkLeft", (1f - data.leftEyeOpen) * eyeScale);
            SetBlendshape("EyeBlinkRight", (1f - data.rightEyeOpen) * eyeScale);
        }
        
        // --- Mouth ---
        // OpenSeeFace provides mouth-related data via landmarks
        // We need to compute blendshape values from the facial landmarks
        if (data.landmarks != null && data.landmarks.Length >= 70)
        {
            ComputeMouthBlendshapes(data.landmarks);
            ComputeBrowBlendshapes(data.landmarks);
        }
        
        // --- Eye Gaze ---
        if (enableEyeTracking && data.landmarks != null)
        {
            SetBlendshape("EyeLookInLeft", 0f);
            SetBlendshape("EyeLookOutLeft", 0f);
            SetBlendshape("EyeLookUpLeft", 0f);
            SetBlendshape("EyeLookDownLeft", 0f);
            SetBlendshape("EyeLookInRight", 0f);
            SetBlendshape("EyeLookOutRight", 0f);
            SetBlendshape("EyeLookUpRight", 0f);
            SetBlendshape("EyeLookDownRight", 0f);
        }
    }
    
    void ComputeMouthBlendshapes(Vector2[] landmarks)
    {
        if (landmarks.Length < 70)
            return;
        
        // OpenSeeFace 68-point landmark indices:
        // 48-67 = mouth
        // 0-16 = jaw
        // 17-21 = left eyebrow
        // 22-26 = right eyebrow
        // 27-30 = nose
        // 31-35 = nostrils
        // 36-41 = left eye
        // 42-47 = right eye
        
        // Mouth corners (48 = left, 54 = right)
        Vector2 mouthLeft = landmarks[48];
        Vector2 mouthRight = landmarks[54];
        Vector2 mouthTop = landmarks[51];
        Vector2 mouthBottom = landmarks[57];
        Vector2 mouthCenter = (mouthLeft + mouthRight) * 0.5f;
        
        // Mouth width (smile detection)
        float mouthWidth = Vector2.Distance(mouthLeft, mouthRight);
        float mouthHeight = Vector2.Distance(mouthTop, mouthBottom);
        
        // Jaw/face width for normalization
        float faceWidth = Vector2.Distance(landmarks[0], landmarks[16]);
        if (faceWidth < 0.01f) faceWidth = 0.01f;
        
        // Normalize
        float normalizedWidth = mouthWidth / faceWidth;
        float normalizedHeight = mouthHeight / faceWidth;
        
        // Mouth open (JawOpen)
        float jawOpen = Mathf.Clamp01((normalizedHeight - 0.02f) * 10f) * mouthScale;
        SetBlendshape("JawOpen", jawOpen);
        
        // Smile detection — corners higher than center
        float leftCornerHeight = mouthCenter.y - mouthLeft.y;
        float rightCornerHeight = mouthCenter.y - mouthRight.y;
        
        float smileLeft = Mathf.Clamp01(leftCornerHeight * 5f) * mouthScale;
        float smileRight = Mathf.Clamp01(rightCornerHeight * 5f) * mouthScale;
        SetBlendshape("MouthSmileLeft", smileLeft);
        SetBlendshape("MouthSmileRight", smileRight);
        
        // Frown — corners lower than center
        float frownLeft = Mathf.Clamp01(-leftCornerHeight * 5f) * mouthScale;
        float frownRight = Mathf.Clamp01(-rightCornerHeight * 5f) * mouthScale;
        SetBlendshape("MouthFrownLeft", frownLeft);
        SetBlendshape("MouthFrownRight", frownRight);
        
        // Mouth pucker
        float pucker = Mathf.Clamp01(0.35f - normalizedWidth) * 3f * mouthScale;
        SetBlendshape("MouthPucker", pucker);
        
        // Mouth close (when mouth is slightly open but lips together)
        if (normalizedHeight < 0.03f)
        {
            SetBlendshape("MouthClose", 0.5f);
        }
        else
        {
            SetBlendshape("MouthClose", 0f);
        }
    }
    
    void ComputeBrowBlendshapes(Vector2[] landmarks)
    {
        // Left eyebrow: 17-21, Right eyebrow: 22-26
        // Eye center for reference
        Vector2 leftEyeCenter = (landmarks[36] + landmarks[39]) * 0.5f;
        Vector2 rightEyeCenter = (landmarks[42] + landmarks[45]) * 0.5f;
        
        // Inner brows (points 21 and 22)
        Vector2 leftInnerBrow = landmarks[21];
        Vector2 rightInnerBrow = landmarks[22];
        
        // Outer brows (points 17 and 26)
        Vector2 leftOuterBrow = landmarks[17];
        Vector2 rightOuterBrow = landmarks[26];
        
        // Brow heights relative to eyes
        float leftInnerHeight = leftEyeCenter.y - leftInnerBrow.y;
        float rightInnerHeight = rightEyeCenter.y - rightInnerBrow.y;
        float leftOuterHeight = leftEyeCenter.y - leftOuterBrow.y;
        float rightOuterHeight = rightEyeCenter.y - rightOuterBrow.y;
        
        // Normalize by face width
        float faceWidth = Vector2.Distance(landmarks[0], landmarks[16]);
        if (faceWidth < 0.01f) faceWidth = 0.01f;
        
        float normLeftInner = leftInnerHeight / faceWidth;
        float normRightInner = rightInnerHeight / faceWidth;
        float normLeftOuter = leftOuterHeight / faceWidth;
        float normRightOuter = rightOuterHeight / faceWidth;
        
        // Brow up
        SetBlendshape("BrowInnerUp", Mathf.Clamp01(normLeftInner * 8f) * browScale);
        SetBlendshape("BrowOuterUpLeft", Mathf.Clamp01(normLeftOuter * 8f) * browScale);
        SetBlendshape("BrowOuterUpRight", Mathf.Clamp01(normRightOuter * 8f) * browScale);
        
        // Brow down (when brow is lower than neutral)
        SetBlendshape("BrowInnerDown", Mathf.Clamp01(-normLeftInner * 8f + 0.3f) * browScale);
        SetBlendshape("BrowOuterDownLeft", Mathf.Clamp01(-normLeftOuter * 8f + 0.3f) * browScale);
        SetBlendshape("BrowOuterDownRight", Mathf.Clamp01(-normRightOuter * 8f + 0.3f) * browScale);
    }
    
    void HandleAutoBlink()
    {
        if (blinking)
        {
            blinkState += Time.deltaTime / blinkSpeed;
            float blinkVal = 0f;
            
            if (blinkState < 0.5f)
                blinkVal = blinkState * 2f; // Close
            else if (blinkState < 1f)
                blinkVal = (1f - blinkState) * 2f; // Open
            else
            {
                blinkVal = 0f;
                blinking = false;
                blinkState = 0f;
                ScheduleNextBlink();
            }
            
            // Only apply if tracking isn't already closing the eyes
            if (openSeeReceiver == null || openSeeReceiver.trackingData == null || openSeeReceiver.trackingData.Length == 0)
            {
                SetBlendshape("EyeBlinkLeft", blinkVal);
                SetBlendshape("EyeBlinkRight", blinkVal);
            }
        }
        else if (Time.time > nextBlinkTime)
        {
            blinking = true;
            blinkState = 0f;
        }
    }
    
    void ScheduleNextBlink()
    {
        nextBlinkTime = Time.time + Random.Range(minBlinkInterval, maxBlinkInterval);
    }
    
    void SetBlendshape(string name, float value)
    {
        string fullName = blendshapePrefix + name;
        if (blendshapeMap.TryGetValue(fullName, out int index))
        {
            headMesh.SetBlendShapeWeight(index, Mathf.Clamp01(value) * 100f);
        }
    }
    
    void OnGUI()
    {
        // Debug info
        if (openSeeReceiver != null && openSeeReceiver.trackingData != null && openSeeReceiver.trackingData.Length > 0)
        {
            GUI.Label(new Rect(10, 10, 300, 20), $"Tracking: {openSeeReceiver.receivedPackets} packets");
        }
    }
}