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

public class UDPAvatarBroadcaster : MonoBehaviour
{
    [Header("Network Configuration")]
    [SerializeField] private bool enableBroadcast = true;
    
    [Header("Avatar Configuration")]
    [SerializeField] private Transform avatarRoot;
    [SerializeField] private byte playerID = 1; // Unique ID for this player (1-255)
    [SerializeField] private float updateRate = 30f; // Updates per second
    
    [Header("Data Selection Mode")]
    [SerializeField] private bool fullDataMode = false; // Send ALL bones and blend shapes
    [SerializeField] private bool showFullModeWarning = true;
    
    [Header("Optimized Mode - Core Bones")]
    [SerializeField] private string[] priorityBones = {
        "Hips", "Spine", "Spine1", "Spine2", "Neck", "Head",
        "LeftShoulder", "LeftArm", "LeftForeArm", "LeftHand",
        "RightShoulder", "RightArm", "RightForeArm", "RightHand",
        "LeftUpLeg", "LeftLeg", "LeftFoot",
        "RightUpLeg", "RightLeg", "RightFoot"
    };
    
    [Header("Facial Animation")]
    [SerializeField] private bool includeFacialData = true;
    [SerializeField] private int maxBlendShapes = 10; // Most important blend shapes only (optimized mode)
    
    [Header("Debug")]
    [SerializeField] private bool showDebugInfo = false;
    [SerializeField] private bool logPacketSize = false;
    
    private UdpClient udpClient;
    private IPEndPoint broadcastEndPoint;
    private string broadcastAddress;
    private int broadcastPort;
    private Dictionary<string, Transform> boneCache;
    private List<Transform> allBones; // For full data mode
    private List<SkinnedMeshRenderer> allMeshes; // For full blend shapes
    private SkinnedMeshRenderer headMesh;
    private float lastUpdateTime;
    private uint sequenceNumber = 0;
    private int currentBoneCount;
    private int currentBlendShapeCount;
    
    // Magic number for packet identification
    private const uint AVATAR_MAGIC = 0xC0A0;
    
    // Avatar data structure - supports both optimized and full data modes
    private struct CompactAvatarData
    {
        public byte playerID;           // 1 byte
        public uint sequenceNumber;     // 4 bytes
        public uint timestamp;          // 4 bytes (Unity time as uint milliseconds)
        public bool isFullDataMode;     // 1 byte - indicates if this is full or optimized data
        
        // Root transform (40 bytes)
        public Vector3 rootPosition;    // 12 bytes
        public Quaternion rootRotation; // 16 bytes
        public Vector3 rootScale;       // 12 bytes
        
        // Variable-length bone data
        // Optimized: 20 bones * 28 bytes = 560 bytes
        // Full mode: All bones * 28 bytes (can be 80+ bones = 2240+ bytes)
        public BoneTransform[] bones;   
        
        // Variable-length blend shape data
        // Optimized: 10 floats = 40 bytes
        // Full mode: All blend shapes = up to 200+ floats = 800+ bytes
        public float[] blendShapes;
        
        // Total optimized: ~650 bytes, Full mode: potentially 3000+ bytes
    }
    
    private struct BoneTransform
    {
        public Vector3 position;    // 12 bytes
        public Quaternion rotation; // 16 bytes
    }
    
    void Start()
    {
        // Get network config from global instance
        var cfg = NetworkConfig.Instance;
        if (cfg != null)
        {
            broadcastAddress = cfg.ipAddress;
            broadcastPort = cfg.port;
            
            // Subscribe to peer discovery if enabled
            if (cfg.useAutoDiscovery && Convai.Scripts.Runtime.Multiplayer.UDPPeerDiscovery.Instance != null)
            {
                Convai.Scripts.Runtime.Multiplayer.UDPPeerDiscovery.Instance.OnPeerDiscovered += HandlePeerDiscovered;
                Convai.Scripts.Runtime.Multiplayer.UDPPeerDiscovery.Instance.OnPeerLost += HandlePeerLost;
                Debug.Log("Avatar broadcaster subscribed to peer discovery");
            }
        }
        else
        {
            Debug.LogError("NetworkConfig not found! Please ensure NetworkConfig.asset exists in Resources folder.");
            broadcastAddress = "255.255.255.255";
            broadcastPort = 1221;
        }
        
        InitializeNetworking();
        CacheAvatarComponents();
        
        if (showDebugInfo)
        {
            Debug.Log($"UDP Avatar Broadcaster initialized - Player {playerID}");
            Debug.Log($"Broadcasting to {broadcastAddress}:{broadcastPort}");
            Debug.Log($"Priority bones: {priorityBones.Length}");
        }
    }
    
    void InitializeNetworking()
    {
        try
        {
            udpClient = new UdpClient();
            udpClient.EnableBroadcast = true;
            broadcastEndPoint = new IPEndPoint(IPAddress.Parse(broadcastAddress), broadcastPort);
            
            if (showDebugInfo)
                Debug.Log("UDP client initialized successfully");
        }
        catch (Exception e)
        {
            Debug.LogError($"Failed to initialize UDP client: {e.Message}");
            enableBroadcast = false;
        }
    }
    
    void CacheAvatarComponents()
    {
        if (avatarRoot == null)
            avatarRoot = transform;
            
        boneCache = new Dictionary<string, Transform>();
        allBones = new List<Transform>();
        allMeshes = new List<SkinnedMeshRenderer>();
        
        Transform[] allChildren = avatarRoot.GetComponentsInChildren<Transform>();
        
        if (fullDataMode)
        {
            // Full data mode: Cache ALL bones with SkinnedMeshRenderer references
            SkinnedMeshRenderer[] meshRenderers = avatarRoot.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);
            currentBoneCount = allBones.Count;
            
            if (showDebugInfo)
                Debug.Log($"Full data mode: Cached {currentBoneCount} bones");
                
            // Show warning about packet size
            if (showFullModeWarning)
            {
                int estimatedSize = CalculateFullModePacketSize();
                if (estimatedSize > 1400) // Close to UDP limit
                {
                    Debug.LogWarning($"Full data mode packet size: ~{estimatedSize} bytes. Consider optimized mode for better network performance.");
                }
            }
        }
        else
        {
            // Optimized mode: Cache priority bones only
            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;
                    }
                }
            }
            currentBoneCount = priorityBones.Length;
        }
        
        // Cache facial animation data
        if (includeFacialData)
        {
            SkinnedMeshRenderer[] meshes = avatarRoot.GetComponentsInChildren<SkinnedMeshRenderer>();
            
            if (fullDataMode)
            {
                // Full mode: collect all meshes with blend shapes
                foreach (var mesh in meshes)
                {
                    if (mesh.sharedMesh != null && mesh.sharedMesh.blendShapeCount > 0)
                    {
                        allMeshes.Add(mesh);
                    }
                }
                
                // Calculate total blend shape count
                currentBlendShapeCount = 0;
                foreach (var mesh in allMeshes)
                {
                    currentBlendShapeCount += mesh.sharedMesh.blendShapeCount;
                }
                
                if (showDebugInfo)
                    Debug.Log($"Full data mode: Found {allMeshes.Count} meshes with {currentBlendShapeCount} total blend shapes");
            }
            else
            {
                // Optimized mode: find head mesh only
                foreach (var mesh in meshes)
                {
                    if (mesh.name.ToLower().Contains("head") || mesh.name.ToLower().Contains("face"))
                    {
                        headMesh = mesh;
                        if (showDebugInfo)
                            Debug.Log($"Found head mesh: {mesh.name} with {mesh.sharedMesh.blendShapeCount} blend shapes");
                        break;
                    }
                }
                currentBlendShapeCount = maxBlendShapes;
            }
        }
        
        string modeStr = fullDataMode ? "FULL" : "OPTIMIZED";
        Debug.Log($"{modeStr} mode: {currentBoneCount} bones, {currentBlendShapeCount} blend shapes");
    }
    
    int CalculateFullModePacketSize()
    {
        int headerSize = 1 + 4 + 4 + 1; // playerID + sequence + timestamp + isFullDataMode
        int rootTransformSize = (3 + 4 + 3) * 4; // position + rotation + scale
        int bonesSize = 4 + (currentBoneCount * 7 * 4); // bone count + bones * (pos + rot) * float size
        int blendShapesSize = 4 + (currentBlendShapeCount * 4); // blend shape count + weights * float size
        return headerSize + rootTransformSize + bonesSize + blendShapesSize;
    }
    
    void Update()
    {
        if (!enableBroadcast || udpClient == null)
            return;
            
        // Rate-limited updates
        if (Time.time - lastUpdateTime >= 1f / updateRate)
        {
            BroadcastAvatarData();
            lastUpdateTime = Time.time;
        }
    }
    
    void BroadcastAvatarData()
    {
        try
        {
            CompactAvatarData data = new CompactAvatarData
            {
                playerID = playerID,
                sequenceNumber = sequenceNumber++,
                timestamp = (uint)(Time.time * 1000), // Convert to milliseconds
                isFullDataMode = fullDataMode,
                
                rootPosition = avatarRoot.position,
                rootRotation = avatarRoot.rotation,
                rootScale = avatarRoot.localScale,
                
                bones = new BoneTransform[currentBoneCount],
                blendShapes = new float[currentBlendShapeCount]
            };
            
            // Pack bone data based on mode
            if (fullDataMode)
            {
                // Full mode: pack all cached bones
                for (int i = 0; i < allBones.Count; i++)
                {
                    Transform bone = allBones[i];
                    data.bones[i] = new BoneTransform
                    {
                        position = bone.localPosition,
                        rotation = bone.localRotation
                    };
                }
            }
            else
            {
                // Optimized mode: pack priority bones
                for (int i = 0; i < priorityBones.Length; i++)
                {
                    if (boneCache.TryGetValue(priorityBones[i], out Transform bone))
                    {
                        data.bones[i] = new BoneTransform
                        {
                            position = bone.localPosition,
                            rotation = bone.localRotation
                        };
                    }
                    else
                    {
                        // Default values for missing bones
                        data.bones[i] = new BoneTransform
                        {
                            position = Vector3.zero,
                            rotation = Quaternion.identity
                        };
                    }
                }
            }
            
            // Pack facial blend shape data based on mode
            if (includeFacialData)
            {
                if (fullDataMode)
                {
                    // Full mode: collect all blend shapes from 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)
                                {
                                    data.blendShapes[blendShapeIndex] = mesh.GetBlendShapeWeight(i) / 100f; // Normalize to 0-1
                                    blendShapeIndex++;
                                }
                            }
                        }
                    }
                }
                else
                {
                    // Optimized mode: use head mesh only
                    if (headMesh != null && headMesh.sharedMesh != null)
                    {
                        int blendShapeCount = Mathf.Min(maxBlendShapes, headMesh.sharedMesh.blendShapeCount);
                        for (int i = 0; i < blendShapeCount; i++)
                        {
                            data.blendShapes[i] = headMesh.GetBlendShapeWeight(i) / 100f; // Normalize to 0-1
                        }
                    }
                }
            }
            
            // Serialize to bytes
            byte[] packetData = SerializeCompactData(data);
            
            // Send UDP packet
            udpClient.Send(packetData, packetData.Length, broadcastEndPoint);
            
            if (logPacketSize && sequenceNumber % 30 == 0) // Log every 30 packets (about once per second)
            {
                string modeStr = fullDataMode ? "FULL" : "OPT";
                Debug.Log($"Sent {modeStr} packet #{sequenceNumber}, size: {packetData.Length} bytes, bones: {currentBoneCount}, blend shapes: {currentBlendShapeCount}");
            }
        }
        catch (Exception e)
        {
            Debug.LogError($"Failed to broadcast avatar data: {e.Message}");
        }
    }
    
    byte[] SerializeCompactData(CompactAvatarData data)
    {
        using (MemoryStream stream = new MemoryStream())
        using (BinaryWriter writer = new BinaryWriter(stream))
        {
            // Magic number for packet identification
            writer.Write(AVATAR_MAGIC);
            
            // Header
            writer.Write(data.playerID);
            writer.Write(data.sequenceNumber);
            writer.Write(data.timestamp);
            writer.Write(data.isFullDataMode);
            
            // Root transform
            WriteVector3(writer, data.rootPosition);
            WriteQuaternion(writer, data.rootRotation);
            WriteVector3(writer, data.rootScale);
            
            // Variable-length bones
            writer.Write(data.bones.Length); // Write bone count first
            for (int i = 0; i < data.bones.Length; i++)
            {
                WriteVector3(writer, data.bones[i].position);
                WriteQuaternion(writer, data.bones[i].rotation);
            }
            
            // Variable-length blend shapes
            writer.Write(data.blendShapes.Length); // Write blend shape count first
            for (int i = 0; i < data.blendShapes.Length; i++)
            {
                writer.Write(data.blendShapes[i]);
            }
            
            return stream.ToArray();
        }
    }
    
    void WriteVector3(BinaryWriter writer, Vector3 vector)
    {
        writer.Write(vector.x);
        writer.Write(vector.y);
        writer.Write(vector.z);
    }
    
    void WriteQuaternion(BinaryWriter writer, Quaternion quaternion)
    {
        writer.Write(quaternion.x);
        writer.Write(quaternion.y);
        writer.Write(quaternion.z);
        writer.Write(quaternion.w);
    }
    
    public void SetPlayerID(byte newPlayerID)
    {
        playerID = newPlayerID;
    }
    
    public void SetUpdateRate(float newRate)
    {
        updateRate = Mathf.Clamp(newRate, 1f, 120f);
    }
    
    void OnDestroy()
    {
        // Unsubscribe from peer discovery
        if (Convai.Scripts.Runtime.Multiplayer.UDPPeerDiscovery.Instance != null)
        {
            Convai.Scripts.Runtime.Multiplayer.UDPPeerDiscovery.Instance.OnPeerDiscovered -= HandlePeerDiscovered;
            Convai.Scripts.Runtime.Multiplayer.UDPPeerDiscovery.Instance.OnPeerLost -= HandlePeerLost;
        }
        
        if (udpClient != null)
        {
            udpClient.Close();
            udpClient.Dispose();
        }
    }
    
    private void HandlePeerDiscovered(string peerIP)
    {
        broadcastAddress = peerIP;
        broadcastEndPoint = new IPEndPoint(IPAddress.Parse(peerIP), broadcastPort);
        Debug.Log($"👤 Avatar broadcaster now targeting peer at {peerIP}:{broadcastPort}");
    }
    
    private void HandlePeerLost()
    {
        var cfg = NetworkConfig.Instance;
        if (cfg != null)
        {
            broadcastAddress = cfg.fallbackBroadcastIP;
            broadcastEndPoint = new IPEndPoint(IPAddress.Parse(broadcastAddress), broadcastPort);
            Debug.LogWarning($"👤 Avatar broadcaster falling back to broadcast: {broadcastAddress}");
        }
    }
    
    void OnApplicationPause(bool pauseStatus)
    {
        enableBroadcast = !pauseStatus;
    }
    
    void OnGUI()
    {
        if (!showDebugInfo) return;
        
        GUILayout.BeginArea(new Rect(10, 10, 300, 240));
        GUILayout.Label($"UDP Broadcaster - Player {playerID}");
        GUILayout.Label($"Mode: {(fullDataMode ? "FULL DATA" : "OPTIMIZED")}");
        GUILayout.Label($"Status: {(enableBroadcast ? "Broadcasting" : "Stopped")}");
        GUILayout.Label($"Sequence: {sequenceNumber}");
        GUILayout.Label($"Update Rate: {updateRate:F1} Hz");
        GUILayout.Label($"Bones: {currentBoneCount}");
        GUILayout.Label($"Blend Shapes: {currentBlendShapeCount}");
        
        if (fullDataMode)
        {
            int estimatedSize = CalculateFullModePacketSize();
            GUILayout.Label($"Est. Packet Size: {estimatedSize} bytes");
            if (estimatedSize > 1400)
            {
                GUI.color = Color.yellow;
                GUILayout.Label("⚠️ Large packet size!");
                GUI.color = Color.white;
            }
        }
        
        if (GUILayout.Button(enableBroadcast ? "Stop Broadcasting" : "Start Broadcasting"))
        {
            enableBroadcast = !enableBroadcast;
        }
        
        GUILayout.EndArea();
    }
}