using System;
using System.Collections.Generic;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using UnityEngine;

/// <summary>
/// UDPAvatarReceiverAgent - Adapted for CC_Base_ bone naming convention (Ready Player Me / Reallusion Character Creator)
/// </summary>
public class UDPAvatarReceiverAgent : MonoBehaviour
{
    [Header("Network Configuration")]
    [SerializeField] private bool enableReceiver = true;
    [SerializeField] private byte targetPlayerID = 5; // Which player to receive data from (0 = any)
    [SerializeField] private bool allowPortSharing = true; // For local testing with multiple components
    
    [Header("Avatar Configuration")]
    [SerializeField] private Transform targetAvatarRoot;
    [SerializeField] private bool smoothTransitions = true;
    [SerializeField] private float transitionSpeed = 10f;
    
    [Header("Data Selection - Core Bones (CC_Base_ naming)")]
    [SerializeField] private string[] priorityBones = {
        "CC_Base_Hip", "CC_Base_Pelvis", "CC_Base_Spine01", "CC_Base_Spine02", "CC_Base_Neck", "CC_Base_Head",
        "CC_Base_L_Clavicle", "CC_Base_L_Upperarm", "CC_Base_L_Forearm", "CC_Base_L_Hand",
        "CC_Base_R_Clavicle", "CC_Base_R_Upperarm", "CC_Base_R_Forearm", "CC_Base_R_Hand",
        "CC_Base_L_Thigh", "CC_Base_L_Calf", "CC_Base_L_Foot",
        "CC_Base_R_Thigh", "CC_Base_R_Calf", "CC_Base_R_Foot"
    };
    
    [Header("Synchronization Options")]
    [SerializeField] private bool syncRootPosition = true;
    [SerializeField] private bool syncRootRotation = true;
    [SerializeField] private bool syncRootScale = false;
    [SerializeField] private bool syncFacialData = true;
    
    [Header("Debug")]
    [SerializeField] private bool showDebugInfo = false;
    [SerializeField] private bool logReceivedPackets = false;
    
    private UdpClient udpClient;
    private Thread udpListenerThread;
    private bool threadRunning = false;
    private int listenPort;
    private Dictionary<string, Transform> boneCache;
    private List<Transform> allBones; // For full data mode
    private List<SkinnedMeshRenderer> allMeshes; // For full blend shapes
    private SkinnedMeshRenderer headMesh;
    
    // Latest received data
    private CompactAvatarData lastReceivedData;
    private bool hasNewData = false;
    private object dataLock = new object();
    
    // Stats
    private uint lastSequenceNumber = 0;
    private int packetsReceived = 0;
    private int packetsDropped = 0;
    private float lastPacketTime = 0f;
    
    // Magic number for packet identification
    private const uint AVATAR_MAGIC = 0xC0A0;
    
    // Avatar data structure - supports both optimized and full data modes (matches broadcaster)
    private struct CompactAvatarData
    {
        public byte playerID;
        public uint sequenceNumber;
        public uint timestamp;
        public bool isFullDataMode;     // indicates if this is full or optimized data
        
        public Vector3 rootPosition;
        public Quaternion rootRotation;
        public Vector3 rootScale;
        
        public BoneTransform[] bones;   // Variable-length
        public float[] blendShapes;     // Variable-length
    }
    
    private struct BoneTransform
    {
        public Vector3 position;
        public Quaternion rotation;
    }
    
    void Start()
    {
        if (targetAvatarRoot == null)
            targetAvatarRoot = transform;
            
        // Get network config from global instance
        var cfg = NetworkConfig.Instance;
        if (cfg != null)
        {
            listenPort = cfg.port;
        }
        else
        {
            Debug.LogError("NetworkConfig not found! Please ensure NetworkConfig.asset exists in Resources folder.");
            listenPort = 1221;
        }

        CacheAvatarComponents();
        
        if (enableReceiver)
        {
            StartUDPListener();
        }
        
        if (showDebugInfo)
        {
            Debug.Log($"UDP Avatar Receiver Agent initialized");
            Debug.Log($"Listening on port {listenPort}");
            Debug.Log($"Target Player ID: {(targetPlayerID == 0 ? "Any" : targetPlayerID.ToString())}");
        }
    }
    
    void CacheAvatarComponents()
    {
        boneCache = new Dictionary<string, Transform>();
        allBones = new List<Transform>();
        allMeshes = new List<SkinnedMeshRenderer>();
        
        Transform[] allChildren = targetAvatarRoot.GetComponentsInChildren<Transform>();
        
        // Cache priority bones for optimized mode
        foreach (Transform child in allChildren)
        {
            foreach (string boneName in priorityBones)
            {
                if (child.name.Equals(boneName, StringComparison.OrdinalIgnoreCase))
                {
                    if (!boneCache.ContainsKey(boneName))
                    {
                        boneCache[boneName] = child;
                        if (showDebugInfo)
                            Debug.Log($"Cached priority bone: {boneName}");
                    }
                    break;
                }
            }
        }
        
        // Cache ALL bones for full data mode
        SkinnedMeshRenderer[] meshRenderers = targetAvatarRoot.GetComponentsInChildren<SkinnedMeshRenderer>();
        HashSet<Transform> uniqueBones = new HashSet<Transform>();
        
        foreach (SkinnedMeshRenderer smr in meshRenderers)
        {
            if (smr.bones != null)
            {
                foreach (Transform bone in smr.bones)
                {
                    if (bone != null)
                    {
                        uniqueBones.Add(bone);
                    }
                }
            }
        }
        
        allBones.AddRange(uniqueBones);
        if (showDebugInfo)
            Debug.Log($"Cached {allBones.Count} total bones for full data mode");
        
        // Cache facial animation data
        if (syncFacialData)
        {
            SkinnedMeshRenderer[] meshes = targetAvatarRoot.GetComponentsInChildren<SkinnedMeshRenderer>();
            
            // Find head mesh for optimized mode
            foreach (var mesh in meshes)
            {
                if (mesh.name.ToLower().Contains("head") || mesh.name.ToLower().Contains("face") || mesh.name.ToLower().Contains("wolf3d_head"))
                {
                    headMesh = mesh;
                    if (showDebugInfo)
                        Debug.Log($"Found head mesh: {mesh.name} with {mesh.sharedMesh.blendShapeCount} blend shapes");
                    break;
                }
            }
            
            // Cache all meshes with blend shapes for full data mode
            foreach (var mesh in meshes)
            {
                if (mesh.sharedMesh != null && mesh.sharedMesh.blendShapeCount > 0)
                {
                    allMeshes.Add(mesh);
                }
            }
            
            if (showDebugInfo)
            {
                int totalBlendShapes = 0;
                foreach (var mesh in allMeshes)
                    totalBlendShapes += mesh.sharedMesh.blendShapeCount;
                Debug.Log($"Found {allMeshes.Count} meshes with {totalBlendShapes} total blend shapes for full data mode");
            }
        }
        
        Debug.Log($"Receiver Agent ready: {boneCache.Count}/{priorityBones.Length} priority bones, {allBones.Count} total bones");
    }
    
    void StartUDPListener()
    {
        try
        {
            if (allowPortSharing)
            {
                // Create UDP client with port reuse for local testing
                udpClient = new UdpClient();
                udpClient.Client.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
                udpClient.Client.Bind(new IPEndPoint(IPAddress.Any, listenPort));
            }
            else
            {
                // Standard UDP client binding
                udpClient = new UdpClient(listenPort);
            }
            
            threadRunning = true;
            udpListenerThread = new Thread(new ThreadStart(UDPListenerThread));
            udpListenerThread.IsBackground = true;
            udpListenerThread.Start();
            
            if (showDebugInfo)
                Debug.Log($"UDP Avatar Receiver Agent started on port {listenPort} (filtering binary avatar data only, Port sharing: {allowPortSharing})");
        }
        catch (Exception e)
        {
            if (allowPortSharing)
            {
                Debug.LogWarning($"Failed to start UDP listener with port sharing: {e.Message}");
                Debug.LogWarning("Trying with different port...");
                TryAlternativePort();
            }
            else
            {
                Debug.LogError($"Failed to start UDP listener: {e.Message}");
                enableReceiver = false;
            }
        }
    }
    
    void TryAlternativePort()
    {
        // Try a few alternative ports for local testing
        int[] alternativePorts = { 8081, 8082, 8083, 8084, 8085 };
        
        foreach (int port in alternativePorts)
        {
            try
            {
                udpClient = new UdpClient(port);
                threadRunning = true;
                udpListenerThread = new Thread(new ThreadStart(UDPListenerThread));
                udpListenerThread.IsBackground = true;
                udpListenerThread.Start();
                
                Debug.Log($"UDP listener started on alternative port {port}");
                return;
            }
            catch (Exception)
            {
                // Try next port
                continue;
            }
        }
        
        Debug.LogError("Failed to start UDP listener on any available port");
        enableReceiver = false;
    }
    
    void UDPListenerThread()
    {
        IPEndPoint remoteEndPoint = new IPEndPoint(IPAddress.Any, 0);
        
        while (threadRunning)
        {
            try
            {
                byte[] data = udpClient.Receive(ref remoteEndPoint);
                
                if (data.Length > 0 && IsAvatarData(data))
                {
                    CompactAvatarData avatarData = DeserializeCompactData(data);
                    
                    // Check if this is from the target player (0 means accept from any player)
                    if (targetPlayerID == 0 || avatarData.playerID == targetPlayerID)
                    {
                        // Check for packet loss
                        if (avatarData.sequenceNumber > lastSequenceNumber + 1)
                        {
                            packetsDropped += (int)(avatarData.sequenceNumber - lastSequenceNumber - 1);
                        }
                        
                        lastSequenceNumber = avatarData.sequenceNumber;
                        packetsReceived++;
                        
                        // Store the new data (thread-safe)
                        lock (dataLock)
                        {
                            lastReceivedData = avatarData;
                            hasNewData = true;
                        }
                        
                        if (logReceivedPackets && packetsReceived % 30 == 0)
                        {
                            string modeStr = avatarData.isFullDataMode ? "FULL" : "OPT";
                            Debug.Log($"Received {modeStr} packet #{avatarData.sequenceNumber} from player {avatarData.playerID}, size: {data.Length} bytes, bones: {avatarData.bones?.Length ?? 0}, blend shapes: {avatarData.blendShapes?.Length ?? 0}");
                        }
                    }
                }
            }
            catch (Exception e)
            {
                if (threadRunning) // Only log errors if we're supposed to be running
                {
                    Debug.LogError($"UDP receive error: {e.Message}");
                }
            }
        }
    }
    
    bool IsAvatarData(byte[] data)
    {
        try
        {
            // Check minimum size for magic number
            if (data.Length < 4)
                return false;
            
            // Check for avatar magic number (0xC0A0)
            uint magic = BitConverter.ToUInt32(data, 0);
            return magic == AVATAR_MAGIC;
        }
        catch
        {
            return false;
        }
    }
    
    CompactAvatarData DeserializeCompactData(byte[] data)
    {
        using (MemoryStream stream = new MemoryStream(data))
        using (BinaryReader reader = new BinaryReader(stream))
        {
            // Read and validate magic number
            uint magic = reader.ReadUInt32();
            if (magic != AVATAR_MAGIC)
            {
                Debug.LogWarning($"Invalid avatar packet magic number: 0x{magic:X} (expected 0x{AVATAR_MAGIC:X})");
                return default;
            }
            
            CompactAvatarData result = new CompactAvatarData
            {
                // Header
                playerID = reader.ReadByte(),
                sequenceNumber = reader.ReadUInt32(),
                timestamp = reader.ReadUInt32(),
                isFullDataMode = reader.ReadBoolean(),
                
                // Root transform
                rootPosition = ReadVector3(reader),
                rootRotation = ReadQuaternion(reader),
                rootScale = ReadVector3(reader)
            };
            
            // Variable-length bones
            int boneCount = reader.ReadInt32();
            result.bones = new BoneTransform[boneCount];
            for (int i = 0; i < boneCount; i++)
            {
                result.bones[i] = new BoneTransform
                {
                    position = ReadVector3(reader),
                    rotation = ReadQuaternion(reader)
                };
            }
            
            // Variable-length blend shapes
            int blendShapeCount = reader.ReadInt32();
            result.blendShapes = new float[blendShapeCount];
            for (int i = 0; i < blendShapeCount; i++)
            {
                result.blendShapes[i] = reader.ReadSingle();
            }
            
            return result;
        }
    }
    
    Vector3 ReadVector3(BinaryReader reader)
    {
        return new Vector3(
            reader.ReadSingle(),
            reader.ReadSingle(),
            reader.ReadSingle()
        );
    }
    
    Quaternion ReadQuaternion(BinaryReader reader)
    {
        return new Quaternion(
            reader.ReadSingle(),
            reader.ReadSingle(),
            reader.ReadSingle(),
            reader.ReadSingle()
        );
    }
    
    void Update()
    {
        if (!enableReceiver)
            return;
            
        // Apply new data if available
        lock (dataLock)
        {
            if (hasNewData)
            {
                ApplyAvatarData(lastReceivedData);
                hasNewData = false;
                lastPacketTime = Time.time;
            }
        }
    }
    
    void ApplyAvatarData(CompactAvatarData data)
    {
        // Apply root transform
        if (syncRootPosition)
        {
            if (smoothTransitions)
            {
                targetAvatarRoot.position = Vector3.Lerp(
                    targetAvatarRoot.position,
                    data.rootPosition,
                    Time.deltaTime * transitionSpeed
                );
            }
            else
            {
                targetAvatarRoot.position = data.rootPosition;
            }
        }
        
        if (syncRootRotation)
        {
            if (smoothTransitions)
            {
                targetAvatarRoot.rotation = Quaternion.Lerp(
                    targetAvatarRoot.rotation,
                    data.rootRotation,
                    Time.deltaTime * transitionSpeed
                );
            }
            else
            {
                targetAvatarRoot.rotation = data.rootRotation;
            }
        }
        
        if (syncRootScale)
        {
            if (smoothTransitions)
            {
                targetAvatarRoot.localScale = Vector3.Lerp(
                    targetAvatarRoot.localScale,
                    data.rootScale,
                    Time.deltaTime * transitionSpeed
                );
            }
            else
            {
                targetAvatarRoot.localScale = data.rootScale;
            }
        }
        
        // Apply bone transforms based on data mode
        if (data.isFullDataMode)
        {
            // Full data mode: apply to all bones based on order
            for (int i = 0; i < Mathf.Min(data.bones.Length, allBones.Count); i++)
            {
                Transform bone = allBones[i];
                if (bone != null)
                {
                    if (smoothTransitions)
                    {
                        bone.localPosition = Vector3.Lerp(
                            bone.localPosition,
                            data.bones[i].position,
                            Time.deltaTime * transitionSpeed
                        );
                        
                        bone.localRotation = Quaternion.Lerp(
                            bone.localRotation,
                            data.bones[i].rotation,
                            Time.deltaTime * transitionSpeed
                        );
                    }
                    else
                    {
                        bone.localPosition = data.bones[i].position;
                        bone.localRotation = data.bones[i].rotation;
                    }
                }
            }
        }
        else
        {
            // Optimized mode: apply to priority bones
            for (int i = 0; i < Mathf.Min(data.bones.Length, priorityBones.Length); i++)
            {
                if (boneCache.TryGetValue(priorityBones[i], out Transform bone))
                {
                    if (smoothTransitions)
                    {
                        bone.localPosition = Vector3.Lerp(
                            bone.localPosition,
                            data.bones[i].position,
                            Time.deltaTime * transitionSpeed
                        );
                        
                        bone.localRotation = Quaternion.Lerp(
                            bone.localRotation,
                            data.bones[i].rotation,
                            Time.deltaTime * transitionSpeed
                        );
                    }
                    else
                    {
                        bone.localPosition = data.bones[i].position;
                        bone.localRotation = data.bones[i].rotation;
                    }
                }
            }
        }
        
        // Apply facial blend shapes based on data mode
        if (syncFacialData)
        {
            if (data.isFullDataMode)
            {
                // Full data mode: apply all blend shapes to all meshes
                int blendShapeIndex = 0;
                foreach (var mesh in allMeshes)
                {
                    if (mesh != null && mesh.sharedMesh != null)
                    {
                        for (int i = 0; i < mesh.sharedMesh.blendShapeCount; i++)
                        {
                            if (blendShapeIndex < data.blendShapes.Length)
                            {
                                float targetWeight = data.blendShapes[blendShapeIndex] * 100f; // Convert back to 0-100 range
                                
                                if (smoothTransitions)
                                {
                                    float currentWeight = mesh.GetBlendShapeWeight(i);
                                    float newWeight = Mathf.Lerp(currentWeight, targetWeight, Time.deltaTime * transitionSpeed);
                                    mesh.SetBlendShapeWeight(i, newWeight);
                                }
                                else
                                {
                                    mesh.SetBlendShapeWeight(i, targetWeight);
                                }
                                blendShapeIndex++;
                            }
                        }
                    }
                }
            }
            else
            {
                // Optimized mode: apply to head mesh only
                if (headMesh != null && headMesh.sharedMesh != null)
                {
                    int blendShapeCount = Mathf.Min(data.blendShapes.Length, headMesh.sharedMesh.blendShapeCount);
                    for (int i = 0; i < blendShapeCount; i++)
                    {
                        float targetWeight = data.blendShapes[i] * 100f; // Convert back to 0-100 range
                        
                        if (smoothTransitions)
                        {
                            float currentWeight = headMesh.GetBlendShapeWeight(i);
                            float newWeight = Mathf.Lerp(currentWeight, targetWeight, Time.deltaTime * transitionSpeed);
                            headMesh.SetBlendShapeWeight(i, newWeight);
                        }
                        else
                        {
                            headMesh.SetBlendShapeWeight(i, targetWeight);
                        }
                    }
                }
            }
        }
    }
    
    public void SetTargetPlayerID(byte playerID)
    {
        targetPlayerID = playerID;
    }
    
    public bool IsReceivingData()
    {
        return Time.time - lastPacketTime < 1f; // Consider "receiving" if we got data in the last second
    }
    
    public float GetPacketLossRate()
    {
        int totalExpected = packetsReceived + packetsDropped;
        return totalExpected > 0 ? (float)packetsDropped / totalExpected : 0f;
    }
    
    void StopUDPListener()
    {
        threadRunning = false;
        
        if (udpClient != null)
        {
            udpClient.Close();
            udpClient.Dispose();
            udpClient = null;
        }
        
        if (udpListenerThread != null)
        {
            udpListenerThread.Join(1000); // Wait up to 1 second for thread to finish
            udpListenerThread = null;
        }
    }
    
    void OnDestroy()
    {
        StopUDPListener();
    }
    
    void OnApplicationPause(bool pauseStatus)
    {
        if (pauseStatus)
        {
            StopUDPListener();
        }
        else if (enableReceiver)
        {
            StartUDPListener();
        }
    }
    
    void OnGUI()
    {
        if (!showDebugInfo) return;
        
        GUILayout.BeginArea(new Rect(320, 10, 300, 260));
        GUILayout.Label($"UDP Avatar Receiver Agent (Binary Only)");
        GUILayout.Label($"Status: {(IsReceivingData() ? "Receiving" : "No Data")}");
        
        // Show last received data mode
        if (hasNewData)
        {
            string modeStr = lastReceivedData.isFullDataMode ? "FULL DATA" : "OPTIMIZED";
            GUILayout.Label($"Last Data Mode: {modeStr}");
            GUILayout.Label($"Bones Received: {lastReceivedData.bones?.Length ?? 0}");
            GUILayout.Label($"Blend Shapes: {lastReceivedData.blendShapes?.Length ?? 0}");
        }
        
        GUILayout.Label($"Listen Port: {GetActualListenPort()} (shared)");
        GUILayout.Label($"Target Player: {(targetPlayerID == 0 ? "Any" : targetPlayerID.ToString())}");
        GUILayout.Label($"Packets: {packetsReceived} received, {packetsDropped} dropped");
        GUILayout.Label($"Packet Loss: {GetPacketLossRate():P1}");
        GUILayout.Label($"Last Sequence: {lastSequenceNumber}");
        GUILayout.Label($"Cached: {boneCache.Count} priority, {allBones.Count} total bones");
        GUILayout.Label($"Port Sharing: {allowPortSharing}");
        
        if (GUILayout.Button(enableReceiver ? "Stop Receiver" : "Start Receiver"))
        {
            enableReceiver = !enableReceiver;
            if (enableReceiver)
                StartUDPListener();
            else
                StopUDPListener();
        }
        
        GUILayout.EndArea();
    }
    
    int GetActualListenPort()
    {
        if (udpClient?.Client?.LocalEndPoint != null)
        {
            return ((IPEndPoint)udpClient.Client.LocalEndPoint).Port;
        }
        return listenPort;
    }
}