using UnityEngine;
using System;
using System.Collections.Generic;
using NativeWebSocket; // Import from Package Manager

[Serializable]
public class BroadcastPacket
{
  public string type;
  public PoseDataPacket originalMessage;
}

[Serializable]
public class PoseLandmark
{
  public float x;
  public float y;
  public float z;
  public float visibility;
}

[Serializable]
public class PoseDataPacket
{
  public string type;
  public PoseLandmark[] data;
  public string image; // Optional: Base64 encoded image data
}

public class PoseReceiver : MonoBehaviour
{
  [Header("Main References")]
  public GameObject avatarRoot;
  public Animator characterAnimator; // Reference to the YBot animator
  public float poseScale = 10f; // Scale factor for MediaPipe coordinates

  [Header("Animation Settings")]
  [Range(0f, 1f)]
  public float rotationSmoothing = 0.3f; // Smoothing factor for rotations
  public float minVisibilityThreshold = 0.5f; // Minimum visibility to consider a landmark valid

  [Header("Debug Options")]
  public bool showDebugObjects = true;
  public float debugObjectScale = 0.1f;
  public Material debugMaterial;

  [Header("Connection Settings")]
  public string webSocketUrl = "ws://localhost:8080";

  // MediaPipe provides 33 landmarks for the full body
  private readonly string[] jointNames = {
        "nose", "left_eye_inner", "left_eye", "left_eye_outer",
        "right_eye_inner", "right_eye", "right_eye_outer",
        "left_ear", "right_ear", "mouth_left", "mouth_right",
        "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
        "left_wrist", "right_wrist", "left_pinky", "right_pinky",
        "left_index", "right_index", "left_thumb", "right_thumb",
        "left_hip", "right_hip", "left_knee", "right_knee",
        "left_ankle", "right_ankle", "left_heel", "right_heel",
        "left_foot_index", "right_foot_index"
    };

  private List<GameObject> bodyJoints = new List<GameObject>();
  private WebSocket websocket;
  private Dictionary<int, HumanBodyBones> boneMapping;
  private Dictionary<HumanBodyBones, Quaternion> initialRotations = new Dictionary<HumanBodyBones, Quaternion>();
  private Dictionary<HumanBodyBones, Quaternion> currentRotations = new Dictionary<HumanBodyBones, Quaternion>();

  private Texture2D receivedTexture;
  private Renderer backgroundRenderer;

  // Store reference poses for better calculations
  private Vector3[] refPose;
  private bool initialPoseStored = false;

  async void Start()
  {
    // Setup bone mapping from MediaPipe to Humanoid rig
    SetupBoneMapping();

    // Initialize debug visualization if enabled
    if (showDebugObjects)
    {
      InitializeDebugObjects();
    }

    // Store initial rotations from the animator's pose
    if (characterAnimator != null)
    {
      StoreInitialRotations();
    }

    // Setup background renderer if needed for image data
    GameObject background = GameObject.CreatePrimitive(PrimitiveType.Quad);
    background.transform.parent = avatarRoot.transform;
    background.transform.localPosition = new Vector3(0, 0.75f, 1f);
    background.transform.localScale = new Vector3(4f, 3f, 1f);
    backgroundRenderer = background.GetComponent<Renderer>();
    if (debugMaterial != null)
    {
      backgroundRenderer.material = debugMaterial;
    }
    receivedTexture = new Texture2D(640, 480, TextureFormat.RGBA32, false);

    // Setup WebSocket connection
    websocket = new WebSocket(webSocketUrl);

    websocket.OnOpen += () =>
    {
      Debug.Log("Connection open!");
    };

    websocket.OnError += (e) =>
    {
      Debug.LogError("WebSocket Error! " + e);
    };

    websocket.OnClose += (e) =>
    {
      Debug.Log("Connection closed!");
    };

    websocket.OnMessage += (bytes) =>
    {
      // Getting the message as a string
      var message = System.Text.Encoding.UTF8.GetString(bytes);
      ProcessPoseData(message);
    };

    // Connect to WebSocket server
    try
    {
      await websocket.Connect();
    }
    catch (Exception e)
    {
      Debug.LogError($"Failed to connect: {e.Message}");
    }
  }

  void Update()
  {
#if !UNITY_WEBGL || UNITY_EDITOR
    if (websocket != null)
    {
      websocket.DispatchMessageQueue();
    }
#endif
  }

  // For Unity WebGL direct communication
  public void UpdatePose(string jsonData)
  {
    ProcessPoseData(jsonData);
  }

  private void StoreInitialRotations()
  {
    // Store initial rotations of all bones we care about
    foreach (var mapping in boneMapping)
    {
      HumanBodyBones bone = mapping.Value;
      Transform boneTransform = characterAnimator.GetBoneTransform(bone);

      if (boneTransform != null)
      {
        initialRotations[bone] = boneTransform.localRotation;
        currentRotations[bone] = boneTransform.localRotation;
      }
    }
  }

  private void ProcessPoseData(string jsonData)
  {
    try
    {
      // Try to parse as broadcast packet first
      string fixedJson = fixJson(jsonData);
      BroadcastPacket broadcast = JsonUtility.FromJson<BroadcastPacket>(fixedJson);

      if (broadcast != null && broadcast.originalMessage != null &&
          broadcast.originalMessage.data != null && broadcast.originalMessage.data.Length > 0)
      {
        // Process pose landmarks
        ProcessLandmarks(broadcast.originalMessage.data);

        // Process image data if available
        if (!string.IsNullOrEmpty(broadcast.originalMessage.image))
        {
          ProcessImageData(broadcast.originalMessage.image);
        }
      }
      else
      {
        // Fallback approach for direct data
        CustomPoseDataWrapper wrapper = JsonUtility.FromJson<CustomPoseDataWrapper>(fixedJson);
        if (wrapper != null && wrapper.data != null && wrapper.data.Length > 0)
        {
          ProcessLandmarks(wrapper.data);
        }
      }
    }
    catch (Exception e)
    {
      Debug.LogError($"Error processing pose data: {e.Message}\nJSON: {jsonData}");
    }
  }

  // Helper class for direct JSON parsing
  [Serializable]
  private class CustomPoseDataWrapper
  {
    public PoseLandmark[] data;
  }

  // Helper function to fix JSON string for arrays
  private string fixJson(string json)
  {
    if (json.StartsWith("[") && json.EndsWith("]"))
    {
      json = "{\"data\":" + json + "}";
    }
    else if (!json.StartsWith("{"))
    {
      json = "{\"data\":" + json + "}";
    }
    return json;
  }

  private void ProcessLandmarks(PoseLandmark[] landmarks)
  {
    // Store reference pose if this is the first valid pose
    if (!initialPoseStored && landmarks.Length == jointNames.Length)
    {
      StoreReferencePose(landmarks);
    }

    // Update debug visualization if enabled
    if (showDebugObjects)
    {
      UpdateDebugObjects(landmarks);
    }

    // Apply to character if animator is available
    if (characterAnimator != null)
    {
      ApplyToCharacter(landmarks);
    }
  }

  private void StoreReferencePose(PoseLandmark[] landmarks)
  {
    refPose = new Vector3[landmarks.Length];
    bool allValid = true;

    for (int i = 0; i < landmarks.Length; i++)
    {
      if (landmarks[i].visibility > minVisibilityThreshold)
      {
        refPose[i] = ConvertCoordinates(landmarks[i]);
      }
      else
      {
        allValid = false;
        break;
      }
    }

    if (allValid)
    {
      initialPoseStored = true;
      Debug.Log("Reference pose stored successfully");
    }
  }

  private void UpdateDebugObjects(PoseLandmark[] landmarks)
  {
    // Ensure we have joint objects for visualization
    if (bodyJoints.Count == 0)
    {
      InitializeDebugObjects();
    }

    // Update positions of debug objects
    for (int i = 0; i < landmarks.Length && i < bodyJoints.Count; i++)
    {
      if (landmarks[i].visibility > minVisibilityThreshold)
      {
        // Convert MediaPipe coordinates to Unity
        Vector3 position = ConvertCoordinates(landmarks[i]);
        bodyJoints[i].transform.localPosition = position;
        bodyJoints[i].SetActive(true);
      }
      else
      {
        bodyJoints[i].SetActive(false);
      }
    }
  }

  private void ApplyToCharacter(PoseLandmark[] landmarks)
  {
    // Apply improved bone rotations based on position data
    CalculateImprovedJointRotations(landmarks);
  }

  private Vector3 ConvertCoordinates(PoseLandmark landmark)
  {
    // Convert MediaPipe coordinates (x right, y down, z forward) 
    // to Unity coordinates (x right, y up, z forward)
    return new Vector3(
        landmark.x * poseScale,
        -landmark.y * poseScale, // Invert Y
        landmark.z * poseScale
    );
  }

  private void CalculateImprovedJointRotations(PoseLandmark[] landmarks)
  {
    // More sophisticated rotation calculation - the key improvement

    // Calculate torso orientation first (this serves as our reference)
    CalculateTorsoOrientation(landmarks);

    // Calculate arms
    CalculateArmRotations(landmarks);

    // Calculate legs
    CalculateLegRotations(landmarks);

    // Calculate head orientation
    CalculateHeadOrientation(landmarks);
  }

  private void CalculateTorsoOrientation(PoseLandmark[] landmarks)
  {
    // We need multiple points to establish torso orientation
    if (landmarks[11].visibility > minVisibilityThreshold &&
        landmarks[12].visibility > minVisibilityThreshold &&
        landmarks[23].visibility > minVisibilityThreshold &&
        landmarks[24].visibility > minVisibilityThreshold)
    {
      // Get positions
      Vector3 leftShoulder = ConvertCoordinates(landmarks[11]);
      Vector3 rightShoulder = ConvertCoordinates(landmarks[12]);
      Vector3 leftHip = ConvertCoordinates(landmarks[23]);
      Vector3 rightHip = ConvertCoordinates(landmarks[24]);

      // Calculate spine direction (middle of hips to middle of shoulders)
      Vector3 hipsCenter = (leftHip + rightHip) * 0.5f;
      Vector3 shouldersCenter = (leftShoulder + rightShoulder) * 0.5f;
      Vector3 spineDir = (shouldersCenter - hipsCenter).normalized;

      // Calculate chest forward direction (perpendicular to shoulder line and spine)
      Vector3 shoulderDir = (rightShoulder - leftShoulder).normalized;
      Vector3 chestForward = Vector3.Cross(shoulderDir, spineDir).normalized;

      // Create rotation for spine
      Quaternion spineRotation = Quaternion.LookRotation(chestForward, spineDir);

      // Apply to hips (the root of the character)
      ApplySmoothedRotation(HumanBodyBones.Hips, spineRotation);

      // Calculate and apply chest/spine rotation slightly differently
      Quaternion chestRotation = Quaternion.LookRotation(chestForward, spineDir);
      ApplySmoothedRotation(HumanBodyBones.Spine, chestRotation);
    }
  }

  private void CalculateArmRotations(PoseLandmark[] landmarks)
  {
    // Left arm
    if (landmarks[11].visibility > minVisibilityThreshold &&
        landmarks[13].visibility > minVisibilityThreshold &&
        landmarks[15].visibility > minVisibilityThreshold)
    {
      Vector3 shoulderPos = ConvertCoordinates(landmarks[11]);
      Vector3 elbowPos = ConvertCoordinates(landmarks[13]);
      Vector3 wristPos = ConvertCoordinates(landmarks[15]);

      // Upper arm direction
      Vector3 upperArmDir = (shoulderPos - elbowPos).normalized;

      // We need a perpendicular vector for proper orientation
      // For the left arm, we can use a cross product with the world up
      Vector3 perpVector = Vector3.Cross(upperArmDir, Vector3.up).normalized;

      // Create and apply rotation
      Quaternion upperArmRotation = Quaternion.LookRotation(perpVector, -upperArmDir);
      ApplySmoothedRotation(HumanBodyBones.LeftUpperArm, upperArmRotation);

      // For lower arm (forearm)
      Vector3 lowerArmDir = (elbowPos - wristPos).normalized;
      Vector3 lowerPerpVector = Vector3.Cross(lowerArmDir, upperArmDir).normalized;

      Quaternion lowerArmRotation = Quaternion.LookRotation(lowerPerpVector, -lowerArmDir);
      ApplySmoothedRotation(HumanBodyBones.LeftLowerArm, lowerArmRotation);
    }

    // Right arm - similar approach with adjusted vectors
    if (landmarks[12].visibility > minVisibilityThreshold &&
        landmarks[14].visibility > minVisibilityThreshold &&
        landmarks[16].visibility > minVisibilityThreshold)
    {
      Vector3 shoulderPos = ConvertCoordinates(landmarks[12]);
      Vector3 elbowPos = ConvertCoordinates(landmarks[14]);
      Vector3 wristPos = ConvertCoordinates(landmarks[16]);

      Vector3 upperArmDir = (shoulderPos - elbowPos).normalized;
      Vector3 perpVector = Vector3.Cross(Vector3.up, upperArmDir).normalized;

      Quaternion upperArmRotation = Quaternion.LookRotation(perpVector, -upperArmDir);
      ApplySmoothedRotation(HumanBodyBones.RightUpperArm, upperArmRotation);

      Vector3 lowerArmDir = (elbowPos - wristPos).normalized;
      Vector3 lowerPerpVector = Vector3.Cross(upperArmDir, lowerArmDir).normalized;

      Quaternion lowerArmRotation = Quaternion.LookRotation(lowerPerpVector, -lowerArmDir);
      ApplySmoothedRotation(HumanBodyBones.RightLowerArm, lowerArmRotation);
    }
  }

  private void CalculateLegRotations(PoseLandmark[] landmarks)
  {
    // Left leg
    if (landmarks[23].visibility > minVisibilityThreshold &&
        landmarks[25].visibility > minVisibilityThreshold &&
        landmarks[27].visibility > minVisibilityThreshold)
    {
      Vector3 hipPos = ConvertCoordinates(landmarks[23]);
      Vector3 kneePos = ConvertCoordinates(landmarks[25]);
      Vector3 anklePos = ConvertCoordinates(landmarks[27]);

      // Upper leg direction
      Vector3 upperLegDir = (kneePos - hipPos).normalized;
      Vector3 perpVector = Vector3.Cross(upperLegDir, Vector3.right).normalized;

      Quaternion upperLegRotation = Quaternion.LookRotation(perpVector, upperLegDir);
      ApplySmoothedRotation(HumanBodyBones.LeftUpperLeg, upperLegRotation);

      // Lower leg
      Vector3 lowerLegDir = (anklePos - kneePos).normalized;
      Vector3 lowerPerpVector = Vector3.Cross(lowerLegDir, upperLegDir).normalized;

      Quaternion lowerLegRotation = Quaternion.LookRotation(lowerPerpVector, lowerLegDir);
      ApplySmoothedRotation(HumanBodyBones.LeftLowerLeg, lowerLegRotation);
    }

    // Right leg - similar with adjusted vectors
    if (landmarks[24].visibility > minVisibilityThreshold &&
        landmarks[26].visibility > minVisibilityThreshold &&
        landmarks[28].visibility > minVisibilityThreshold)
    {
      Vector3 hipPos = ConvertCoordinates(landmarks[24]);
      Vector3 kneePos = ConvertCoordinates(landmarks[26]);
      Vector3 anklePos = ConvertCoordinates(landmarks[28]);

      Vector3 upperLegDir = (kneePos - hipPos).normalized;
      Vector3 perpVector = Vector3.Cross(Vector3.right, upperLegDir).normalized;

      Quaternion upperLegRotation = Quaternion.LookRotation(perpVector, upperLegDir);
      ApplySmoothedRotation(HumanBodyBones.RightUpperLeg, upperLegRotation);

      Vector3 lowerLegDir = (anklePos - kneePos).normalized;
      Vector3 lowerPerpVector = Vector3.Cross(upperLegDir, lowerLegDir).normalized;

      Quaternion lowerLegRotation = Quaternion.LookRotation(lowerPerpVector, lowerLegDir);
      ApplySmoothedRotation(HumanBodyBones.RightLowerLeg, lowerLegRotation);
    }
  }

  private void CalculateHeadOrientation(PoseLandmark[] landmarks)
  {
    // Head orientation based on eye and ear positions
    if (landmarks[0].visibility > minVisibilityThreshold &&  // nose
        landmarks[7].visibility > minVisibilityThreshold &&  // left ear
        landmarks[8].visibility > minVisibilityThreshold)    // right ear
    {
      Vector3 nose = ConvertCoordinates(landmarks[0]);
      Vector3 leftEar = ConvertCoordinates(landmarks[7]);
      Vector3 rightEar = ConvertCoordinates(landmarks[8]);

      // Face direction (from ear midpoint to nose)
      Vector3 earMidpoint = (leftEar + rightEar) * 0.5f;
      Vector3 faceDir = (nose - earMidpoint).normalized;

      // Right direction of face
      Vector3 rightDir = (rightEar - leftEar).normalized;

      // Up direction (perpendicular to face direction and right)
      Vector3 upDir = -Vector3.Cross(rightDir, faceDir).normalized;

      // Create and apply head rotation
      Quaternion headRotation = Quaternion.LookRotation(faceDir, upDir);
      ApplySmoothedRotation(HumanBodyBones.Head, headRotation);
    }
  }

  private void ApplySmoothedRotation(HumanBodyBones bone, Quaternion targetRotation)
  {
    Transform boneTransform = characterAnimator.GetBoneTransform(bone);
    if (boneTransform != null)
    {
      // Get initial rotation if available, otherwise use identity
      Quaternion initialRot = Quaternion.identity;
      if (initialRotations.TryGetValue(bone, out Quaternion storedRot))
      {
        initialRot = storedRot;
      }

      // Get current rotation for smoothing
      Quaternion currentRot = Quaternion.identity;
      if (!currentRotations.TryGetValue(bone, out currentRot))
      {
        currentRot = boneTransform.localRotation;
        currentRotations[bone] = currentRot;
      }

      // Calculate new rotation with smoothing
      Quaternion newRotation = Quaternion.Slerp(currentRot, targetRotation, 1f - rotationSmoothing);

      // Store for next frame's smoothing
      currentRotations[bone] = newRotation;

      // Apply rotation
      boneTransform.rotation = newRotation;
    }
  }

  private void InitializeDebugObjects()
  {
    // Clear existing joints
    foreach (var joint in bodyJoints)
    {
      if (joint != null)
      {
        Destroy(joint);
      }
    }
    bodyJoints.Clear();

    // Create visual representation for each joint
    for (int i = 0; i < jointNames.Length; i++)
    {
      GameObject joint = GameObject.CreatePrimitive(PrimitiveType.Sphere);
      joint.name = jointNames[i];
      joint.transform.parent = avatarRoot.transform;
      joint.transform.localScale = Vector3.one * debugObjectScale;

      if (debugMaterial != null)
      {
        joint.GetComponent<Renderer>().material = debugMaterial;
      }

      bodyJoints.Add(joint);
    }
  }

  private void ProcessImageData(string base64Image)
  {
    try
    {
      // Remove the data URL prefix if present
      string base64Data = base64Image;
      if (base64Image.Contains(","))
      {
        base64Data = base64Image.Substring(base64Image.IndexOf(",") + 1);
      }

      // Convert base64 string to byte array
      byte[] imageBytes = Convert.FromBase64String(base64Data);

      // Load image into texture
      receivedTexture.LoadImage(imageBytes);
      receivedTexture.Apply();

      // Apply texture to the background quad
      if (backgroundRenderer != null)
      {
        backgroundRenderer.material.mainTexture = receivedTexture;
      }
    }
    catch (Exception e)
    {
      Debug.LogError($"Error processing image data: {e.Message}");
    }
  }

  private void SetupBoneMapping()
  {
    // Map MediaPipe landmark indices to Humanoid bones
    boneMapping = new Dictionary<int, HumanBodyBones>
        {
            // Head
            {0, HumanBodyBones.Head}, // nose -> head
            
            // Torso
            // {11, HumanBodyBones.LeftShoulder},  // left_shoulder
            // {12, HumanBodyBones.RightShoulder}, // right_shoulder
            // {23, HumanBodyBones.LeftUpperLeg},  // left_hip
            // {24, HumanBodyBones.RightUpperLeg}, // right_hip
            
            // Arms
            {11, HumanBodyBones.LeftShoulder},  // left_shoulder as reference
            {12, HumanBodyBones.RightShoulder}, // right_shoulder as reference
            {13, HumanBodyBones.LeftUpperArm},  // left_elbow
            {14, HumanBodyBones.RightUpperArm}, // right_elbow
            {15, HumanBodyBones.LeftLowerArm},  // left_wrist
            {16, HumanBodyBones.RightLowerArm}, // right_wrist
            
            // Legs
            {23, HumanBodyBones.Hips},          // left_hip as reference
            {24, HumanBodyBones.Hips},          // right_hip as reference
            {25, HumanBodyBones.LeftUpperLeg},  // left_knee
            {26, HumanBodyBones.RightUpperLeg}, // right_knee
            {27, HumanBodyBones.LeftLowerLeg},  // left_ankle
            {28, HumanBodyBones.RightLowerLeg}  // right_ankle
        };
  }

  private async void OnApplicationQuit()
  {
    if (websocket != null)
    {
      await websocket.Close();
    }
  }
}