using Quasar.Common.Video;
using System;
using System.Collections.Generic;
using System.Net;
using System.Net.Sockets;
using System.Threading;

namespace Quasar.Server.Networking
{
    /// <summary>
    /// Listens on a UDP port and reassembles fragmented frames sent by the client.
    /// Fires FrameAssembled when a complete frame is ready.
    /// Thread-safe; call Dispose to stop.
    /// </summary>
    public class UdpFrameReceiver : IDisposable
    {
        /// <summary>Fired on the thread-pool when a complete frame payload is assembled.</summary>
        public event Action<byte[]> FrameAssembled;

        public int Port { get; private set; }

        private UdpClient _udp;
        private Thread _receiveThread;
        private volatile bool _running;

        // Reassembly state for the current in-progress frame
        private readonly object _assemblyLock = new object();
        private uint _currentFrameId = uint.MaxValue;
        private byte[][] _chunks;
        private int _chunksReceived;
        private ushort _totalChunks;
        private ushort _currentFlags;

        public UdpFrameReceiver(int port)
        {
            Port = port;
            _udp = new UdpClient(port);
            _udp.Client.ReceiveBufferSize = 1024 * 1024; // 1 MB receive buffer
        }

        public void Start()
        {
            _running = true;
            _receiveThread = new Thread(ReceiveLoop)
            {
                IsBackground = true,
                Name = "UdpFrameReceiver"
            };
            _receiveThread.Start();
        }

        private void ReceiveLoop()
        {
            IPEndPoint remote = new IPEndPoint(IPAddress.Any, 0);
            while (_running)
            {
                try
                {
                    byte[] data = _udp.Receive(ref remote);
                    if (data.Length < UdpFrameProtocol.HeaderSize) continue;

                    UdpFrameProtocol.ParseHeader(data, out uint frameId, out ushort chunkIndex, out ushort totalChunks, out ushort flags);

                    int payloadLen = data.Length - UdpFrameProtocol.HeaderSize;
                    byte[] payload = new byte[payloadLen];
                    Buffer.BlockCopy(data, UdpFrameProtocol.HeaderSize, payload, 0, payloadLen);

                    lock (_assemblyLock)
                    {
                        // New frame — reset assembly state
                        if (frameId != _currentFrameId)
                        {
                            _currentFrameId = frameId;
                            _totalChunks = totalChunks;
                            _currentFlags = flags;
                            _chunks = new byte[totalChunks][];
                            _chunksReceived = 0;
                        }

                        if (chunkIndex < _totalChunks && _chunks[chunkIndex] == null)
                        {
                            _chunks[chunkIndex] = payload;
                            _chunksReceived++;

                            if (_chunksReceived == _totalChunks)
                            {
                                // Assemble full frame
                                int total = 0;
                                for (int i = 0; i < _totalChunks; i++) total += _chunks[i].Length;

                                byte[] frame = new byte[total];
                                int pos = 0;
                                for (int i = 0; i < _totalChunks; i++)
                                {
                                    Buffer.BlockCopy(_chunks[i], 0, frame, pos, _chunks[i].Length);
                                    pos += _chunks[i].Length;
                                }

                                // Decompress if flagged
                                if ((_currentFlags & UdpFrameProtocol.FlagCompressed) != 0)
                                {
                                    try { frame = UdpFrameProtocol.Decompress(frame); }
                                    catch { continue; }
                                }

                                var assembled = frame;
                                ThreadPool.QueueUserWorkItem(_ => FrameAssembled?.Invoke(assembled));

                                // Reset
                                _chunks = null;
                                _chunksReceived = 0;
                            }
                        }
                    }
                }
                catch (SocketException)
                {
                    // Socket closed — exit loop
                    break;
                }
                catch { /* swallow other errors, keep running */ }
            }
        }

        public void Dispose()
        {
            _running = false;
            try { _udp?.Close(); } catch { }
            _udp = null;
        }
    }
}
