using Sick.SafetyScanners.DataStructures;

namespace Sick.SafetyScanners.DataProcessing;

/// <summary>
/// Merges UDP packets that may be fragmented across multiple UDP packets
/// Thread-safe implementation - matches C++ logic
/// </summary>
public sealed class UdpPacketMerger : IDisposable
{
    private readonly Dictionary<uint, List<ParsedPacketBuffer>> _parsedPacketBufferMap = [];
    private bool _isComplete = false;
    private PacketBuffer? _deployedBuffer;
    private bool _disposed;
    private readonly Lock _lock = new();
    
    /// <summary>
    /// Indicates whether a data packet is complete
    /// </summary>
    public bool IsComplete
    {
        get
        {
            lock (_lock)
            {
                return _isComplete;
            }
        }
    }
    
    /// <summary>
    /// Indicates whether the buffer is empty
    /// </summary>
    public bool IsEmpty
    {
        get
        {
            lock (_lock)
            {
                return _parsedPacketBufferMap.Count == 0;
            }
        }
    }
    
    /// <summary>
    /// Adds a UDP packet to the merger (stores reference, doesn't copy data yet)
    /// Returns true if the packet is complete
    /// </summary>
    public bool AddUdpPacket(PacketBuffer buffer)
    {
        ObjectDisposedException.ThrowIf(_disposed, nameof(UdpPacketMerger));
        ArgumentNullException.ThrowIfNull(buffer);

        lock (_lock)
        {
            if (_isComplete)
            {
                _isComplete = false;
                _deployedBuffer = null;
            }
            
            // Parse datagram header
            var headerParser = new ParseDatagramHeader();
            var datagramHeader = headerParser.ParseUdpSequence(buffer);
            
            // Add to map
            AddToMap(buffer, datagramHeader);
            
            // Check if complete and deploy if so
            DeployPacketIfComplete(datagramHeader);
            
            return _isComplete;
        }
    }
    
    /// <summary>
    /// Gets the latest complete data packet
    /// </summary>
    public PacketBuffer GetDeployedPacketBuffer()
    {
        ObjectDisposedException.ThrowIf(_disposed, nameof(UdpPacketMerger));
        
        lock (_lock)
        {
            if (!_isComplete || _deployedBuffer == null)
                throw new InvalidOperationException("No complete packet available");
            
            _isComplete = false;
            var result = _deployedBuffer;
            _deployedBuffer = null;
            return result;
        }
    }
    
    /// <summary>
    /// Resets the merger (clears all buffers)
    /// </summary>
    public void Reset()
    {
        lock (_lock)
        {
            _parsedPacketBufferMap.Clear();
            _isComplete = false;
            _deployedBuffer = null;
        }
    }
    
    private void AddToMap(PacketBuffer buffer, DatagramHeader header)
    {
        var parsedBuffer = new ParsedPacketBuffer(buffer, header);
        
        if (_parsedPacketBufferMap.TryGetValue(header.Identification, out var list))
        {
            list.Add(parsedBuffer);
        }
        else
        {
            _parsedPacketBufferMap[header.Identification] = [parsedBuffer];
        }
    }
    
    private void DeployPacketIfComplete(DatagramHeader header)
    {
        if (!_parsedPacketBufferMap.TryGetValue(header.Identification, out var list))
            return;
        
        if (!CheckIfComplete(header, list))
            return;
        
        // Sort by fragment offset
        var sortedList = list.OrderBy(p => p.DatagramHeader.FragmentOffset).ToList();
        
        // Remove headers and merge data
        var mergedData = RemoveHeaderFromParsedPacketBuffers(sortedList);
        
        _deployedBuffer = new PacketBuffer(mergedData);
        _parsedPacketBufferMap.Remove(header.Identification);
        _isComplete = true;
    }
    
    private static bool CheckIfComplete(DatagramHeader header, List<ParsedPacketBuffer> list)
    {
        var totalLength = header.TotalLength;
        var currentLength = CalculateCurrentLength(list);
        
        if (currentLength != totalLength)
            return false;
        
        return true;
    }
    
    private static uint CalculateCurrentLength(List<ParsedPacketBuffer> list)
    {
        uint currentLength = 0;
        
        foreach (var parsedBuffer in list)
        {
            var packetBuffer = parsedBuffer.PacketBuffer;
            currentLength += (uint)(packetBuffer.GetBuffer().Length - DatagramHeader.HeaderSize);
        }
        
        return currentLength;
    }
    
    private static byte[] RemoveHeaderFromParsedPacketBuffers(List<ParsedPacketBuffer> sortedList)
    {
        var result = new List<byte>();
        
        foreach (var parsedBuffer in sortedList)
        {
            var packetBuffer = parsedBuffer.PacketBuffer;
            var bufferData = packetBuffer.GetBuffer();
            
            // Skip header (first HeaderSize bytes) and add rest
            if (bufferData.Length > DatagramHeader.HeaderSize)
            {
                var dataWithoutHeader = bufferData.Span[DatagramHeader.HeaderSize..];
                result.AddRange(dataWithoutHeader.ToArray());
            }
        }
        
        return [.. result];
    }
    
    public void Dispose()
    {
        if (_disposed)
            return;
        
        lock (_lock)
        {
            _parsedPacketBufferMap.Clear();
            _deployedBuffer?.Dispose();
            _deployedBuffer = null;
        }
        
        _disposed = true;
    }
}