using UnityEngine;
using System;
using System.Collections.Generic;
using NativeWebSocket; // You'll need to import this package via 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("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<int, int> parentJoints;

  private Texture2D receivedTexture;
  private Renderer backgroundRenderer;

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

    // Setup parent-child relationships for calculating rotations
    SetupParentJoints();

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

    // 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, .75f, 0f);
    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);
      // Debug.Log("message: " + message);
      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 ProcessPoseData(string jsonData)
  {
    try
    {
      // Try to parse as data packet first
      BroadcastPacket broadcast = JsonUtility.FromJson<BroadcastPacket>(fixJson(jsonData));
      PoseDataPacket packet = broadcast.originalMessage;

      string fixedJson = fixJson(jsonData);
      Debug.Log("Fixed JSON: " + fixedJson);

      Debug.Log("fixedJson" + fixedJson);
      Debug.Log("Image" + packet.image);
      Debug.Log("Image bool" + string.IsNullOrEmpty(packet.image));
      if (packet != null && packet.data != null && packet.data.Length > 0)
      {
        // Process pose landmarks
        ProcessLandmarks(packet.data);

        // Process image data if available
        if (!string.IsNullOrEmpty(packet.image))
        {
          ProcessImageData(packet.image);
        }
      }
      else
      {
        // Fallback: try to parse as direct landmark array
        PoseLandmark[] landmarks = JsonUtility.FromJson<PoseLandmark[]>(fixJson(jsonData));
        if (landmarks != null && landmarks.Length > 0)
        {
          ProcessLandmarks(landmarks);
        }
      }
    }
    catch (Exception e)
    {
      Debug.LogError($"Error processing pose data: {e.Message}");
    }
  }

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

  private void ProcessLandmarks(PoseLandmark[] landmarks)
  {
    // Update debug visualization if enabled
    if (showDebugObjects)
    {
      UpdateDebugObjects(landmarks);
    }

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

  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 > 0.5f)
      {
        // Convert MediaPipe coordinates to Unity
        Vector3 position = ConvertCoordinates(landmarks[i]);
        bodyJoints[i].transform.localPosition = position;
      }
    }
  }

  private void ApplyToCharacter(PoseLandmark[] landmarks)
  {
    // Apply positions to character bones based on mapping
    foreach (var mapping in boneMapping)
    {
      int jointIndex = mapping.Key;
      HumanBodyBones boneType = mapping.Value;

      if (jointIndex < landmarks.Length && landmarks[jointIndex].visibility > 0.5f)
      {
        Transform bone = characterAnimator.GetBoneTransform(boneType);
        if (bone != null)
        {
          // Instead of setting position directly, we'll calculate bone rotations
          // The direct position approach isn't ideal for character animation
        }
      }
    }

    // Calculate bone rotations based on position data
    CalculateJointRotations(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 CalculateJointRotations(PoseLandmark[] landmarks)
  {
    // Specific rotations for YBot rig based on MediaPipe landmarks
    // This is a simple version - you may need to refine this for better results

    // Calculate key bone rotations
    foreach (var parentEntry in parentJoints)
    {
      int childIndex = parentEntry.Key;
      int parentIndex = parentEntry.Value;

      // Skip if landmarks aren't visible enough
      if (childIndex >= landmarks.Length || parentIndex >= landmarks.Length ||
          landmarks[childIndex].visibility < 0.5f || landmarks[parentIndex].visibility < 0.5f)
      {
        continue;
      }

      // Get positions
      Vector3 childPos = ConvertCoordinates(landmarks[childIndex]);
      Vector3 parentPos = ConvertCoordinates(landmarks[parentIndex]);

      // Calculate direction vector
      Vector3 direction = childPos - parentPos;

      if (direction.magnitude > 0.01f)
      {
        // Find the corresponding bone in the humanoid rig
        if (boneMapping.TryGetValue(parentIndex, out HumanBodyBones bone))
        {
          Transform boneTransform = characterAnimator.GetBoneTransform(bone);
          if (boneTransform != null)
          {
            // Calculate rotation to point the bone toward the child joint
            // This is a simplified approach - in production you'd want a more
            // sophisticated IK solution
            Quaternion targetRotation = Quaternion.LookRotation(direction);

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

  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.Hips},          // left_hip
            {24, HumanBodyBones.Hips},          // right_hip
            
            // Arms
            {13, HumanBodyBones.LeftUpperArm},  // left_elbow
            {14, HumanBodyBones.RightUpperArm}, // right_elbow
            {15, HumanBodyBones.LeftLowerArm},  // left_wrist
            {16, HumanBodyBones.RightLowerArm}, // right_wrist
            
            // Legs
            {25, HumanBodyBones.LeftUpperLeg},  // left_knee
            {26, HumanBodyBones.RightUpperLeg}, // right_knee
            {27, HumanBodyBones.LeftLowerLeg},  // left_ankle
            {28, HumanBodyBones.RightLowerLeg}  // right_ankle
        };
  }

  private void SetupParentJoints()
  {
    // Define parent-child relationships for joint rotation calculation
    parentJoints = new Dictionary<int, int>
        {
            // Head
            {1, 0},  // left_eye_inner -> nose
            {2, 1},  // left_eye -> left_eye_inner
            {3, 2},  // left_eye_outer -> left_eye
            {4, 0},  // right_eye_inner -> nose
            {5, 4},  // right_eye -> right_eye_inner
            {6, 5},  // right_eye_outer -> right_eye
            {7, 3},  // left_ear -> left_eye_outer
            {8, 6},  // right_ear -> right_eye_outer
            {9, 0},  // mouth_left -> nose
            {10, 0}, // mouth_right -> nose
            
            // Arms
            {13, 11}, // left_elbow -> left_shoulder
            {14, 12}, // right_elbow -> right_shoulder
            {15, 13}, // left_wrist -> left_elbow
            {16, 14}, // right_wrist -> right_elbow
            {17, 15}, // left_pinky -> left_wrist
            {18, 16}, // right_pinky -> right_wrist
            {19, 15}, // left_index -> left_wrist
            {20, 16}, // right_index -> right_wrist
            {21, 15}, // left_thumb -> left_wrist
            {22, 16}, // right_thumb -> right_wrist
            
            // Legs
            {25, 23}, // left_knee -> left_hip
            {26, 24}, // right_knee -> right_hip
            {27, 25}, // left_ankle -> left_knee
            {28, 26}, // right_ankle -> right_knee
            {29, 27}, // left_heel -> left_ankle
            {30, 28}, // right_heel -> right_ankle
            {31, 29}, // left_foot_index -> left_heel
            {32, 30}  // right_foot_index -> right_heel
        };
  }

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