using UnityEngine;
using System;
using System.Collections.Generic;
using NativeWebSocket; // You'll need to import this package

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

public class PoseReceiver : MonoBehaviour
{
  public GameObject avatarRoot;
  public List<GameObject> bodyJoints = new List<GameObject>();
  private WebSocket websocket;

  // 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"
    };

  async void Start()
  {
    // Initialize body joint GameObjects if not set in inspector
    if (bodyJoints.Count == 0)
    {
      InitializeBodyJoints();
    }

    // For WebSocket implementation
    websocket = new WebSocket("ws://localhost:8080");

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

    websocket.OnError += (e) =>
    {
      Debug.Log("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
    await websocket.Connect();
  }

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

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

  private void ProcessPoseData(string jsonPoseData)
  {
    try
    {
      // Parse the JSON data
      PoseLandmark[] landmarks = JsonUtility.FromJson<PoseLandmark[]>(jsonPoseData);

      if (landmarks != null && landmarks.Length > 0)
      {
        // Update the position of each joint
        for (int i = 0; i < landmarks.Length && i < bodyJoints.Count; i++)
        {
          if (landmarks[i].visibility > 0.5f) // Only update if visible enough
          {
            // Convert MediaPipe coordinates to Unity coordinates
            // MediaPipe: (x right, y down, z forward) 
            // Unity: (x right, y up, z forward)
            Vector3 position = new Vector3(
                landmarks[i].x,
                -landmarks[i].y, // Invert Y axis
                landmarks[i].z
            );

            // Apply scale factor if needed
            position *= 10f; // Adjust as needed for your scene scale

            // Set the position
            bodyJoints[i].transform.localPosition = position;
          }
        }

        // If needed, calculate and set rotations between joints
        CalculateJointRotations();
      }
    }
    catch (Exception e)
    {
      Debug.LogError($"Error processing pose data: {e.Message}");
    }
  }

  private void InitializeBodyJoints()
  {
    // Create joints if they don't exist
    for (int i = 0; i < jointNames.Length; i++)
    {
      GameObject joint = new GameObject(jointNames[i]);
      joint.transform.parent = avatarRoot.transform;
      bodyJoints.Add(joint);
    }
  }

  private void CalculateJointRotations()
  {
    // Calculate rotations between key joints
    // This is a simplified example - you'll need to adapt this
    // to your specific 3D model and requirements

    // Example: Orient shoulder-to-elbow
    if (bodyJoints.Count > 14) // Make sure we have enough joints
    {
      // Left arm
      CalculateRotation(11, 13, 15); // left_shoulder, left_elbow, left_wrist

      // Right arm
      CalculateRotation(12, 14, 16); // right_shoulder, right_elbow, right_wrist

      // More joints can be calculated similarly
    }
  }

  private void CalculateRotation(int joint1, int joint2, int joint3)
  {
    if (joint1 < bodyJoints.Count && joint2 < bodyJoints.Count && joint3 < bodyJoints.Count)
    {
      Vector3 dir = bodyJoints[joint3].transform.position - bodyJoints[joint2].transform.position;
      bodyJoints[joint2].transform.rotation = Quaternion.LookRotation(dir);
    }
  }

  private async void OnApplicationQuit()
  {
    await websocket.Close();
  }
}