using System.Buffers;
using System.Collections.Concurrent;
using System.IO.Compression;

namespace ZB.MOM.WW.CBDD.Core.Compression;

/// <summary>
///     Compression codec registry and utility service.
/// </summary>
public sealed class CompressionService
{
    private readonly ConcurrentDictionary<CompressionCodec, ICompressionCodec> _codecs = new();

    /// <summary>
    ///     Initializes a new instance of the <see cref="CompressionService" /> class.
    /// </summary>
    /// <param name="additionalCodecs">Optional additional codecs to register.</param>
    public CompressionService(IEnumerable<ICompressionCodec>? additionalCodecs = null)
    {
        RegisterCodec(new NoneCompressionCodec());
        RegisterCodec(new BrotliCompressionCodec());
        RegisterCodec(new DeflateCompressionCodec());

        if (additionalCodecs == null)
            return;

        foreach (var codec in additionalCodecs) RegisterCodec(codec);
    }

    /// <summary>
    ///     Registers or replaces a compression codec implementation.
    /// </summary>
    /// <param name="codec">The codec implementation to register.</param>
    public void RegisterCodec(ICompressionCodec codec)
    {
        ArgumentNullException.ThrowIfNull(codec);
        _codecs[codec.Codec] = codec;
    }

    /// <summary>
    ///     Attempts to resolve a registered codec implementation.
    /// </summary>
    /// <param name="codec">The codec identifier to resolve.</param>
    /// <param name="compressionCodec">When this method returns, contains the resolved codec when found.</param>
    /// <returns>
    ///     <see langword="true" /> when a codec is registered for <paramref name="codec" />; otherwise,
    ///     <see langword="false" />.
    /// </returns>
    public bool TryGetCodec(CompressionCodec codec, out ICompressionCodec compressionCodec)
    {
        return _codecs.TryGetValue(codec, out compressionCodec!);
    }

    /// <summary>
    ///     Gets a registered codec implementation.
    /// </summary>
    /// <param name="codec">The codec identifier to resolve.</param>
    /// <returns>The registered codec implementation.</returns>
    public ICompressionCodec GetCodec(CompressionCodec codec)
    {
        if (_codecs.TryGetValue(codec, out var compressionCodec))
            return compressionCodec;

        throw new InvalidOperationException($"Compression codec '{codec}' is not registered.");
    }

    /// <summary>
    ///     Compresses payload bytes using the selected codec and level.
    /// </summary>
    /// <param name="input">The payload bytes to compress.</param>
    /// <param name="codec">The codec to use.</param>
    /// <param name="level">The compression level.</param>
    /// <returns>The compressed payload bytes.</returns>
    public byte[] Compress(ReadOnlySpan<byte> input, CompressionCodec codec, CompressionLevel level)
    {
        return GetCodec(codec).Compress(input, level);
    }

    /// <summary>
    ///     Decompresses payload bytes using the selected codec.
    /// </summary>
    /// <param name="input">The compressed payload bytes.</param>
    /// <param name="codec">The codec to use.</param>
    /// <param name="expectedLength">
    ///     The expected decompressed byte length, or a negative value to skip exact-length
    ///     validation.
    /// </param>
    /// <param name="maxDecompressedSizeBytes">The maximum allowed decompressed byte length.</param>
    /// <returns>The decompressed payload bytes.</returns>
    public byte[] Decompress(ReadOnlySpan<byte> input, CompressionCodec codec, int expectedLength,
        int maxDecompressedSizeBytes)
    {
        return GetCodec(codec).Decompress(input, expectedLength, maxDecompressedSizeBytes);
    }

    /// <summary>
    ///     Compresses and then decompresses payload bytes using the selected codec.
    /// </summary>
    /// <param name="input">The payload bytes to roundtrip.</param>
    /// <param name="codec">The codec to use.</param>
    /// <param name="level">The compression level.</param>
    /// <param name="maxDecompressedSizeBytes">The maximum allowed decompressed byte length.</param>
    /// <returns>The decompressed payload bytes after roundtrip.</returns>
    public byte[] Roundtrip(ReadOnlySpan<byte> input, CompressionCodec codec, CompressionLevel level,
        int maxDecompressedSizeBytes)
    {
        byte[] compressed = Compress(input, codec, level);
        return Decompress(compressed, codec, input.Length, maxDecompressedSizeBytes);
    }

    private static byte[] CompressWithCodecStream(ReadOnlySpan<byte> input, Func<Stream, Stream> streamFactory)
    {
        using var output = new MemoryStream(input.Length);
        using (var codecStream = streamFactory(output))
        {
            codecStream.Write(input);
            codecStream.Flush();
        }

        return output.ToArray();
    }

    private static byte[] DecompressWithCodecStream(
        ReadOnlySpan<byte> input,
        Func<Stream, Stream> streamFactory,
        int expectedLength,
        int maxDecompressedSizeBytes)
    {
        if (maxDecompressedSizeBytes <= 0)
            throw new ArgumentOutOfRangeException(nameof(maxDecompressedSizeBytes));

        using var compressed = new MemoryStream(input.ToArray(), false);
        using var codecStream = streamFactory(compressed);
        using var output = expectedLength > 0
            ? new MemoryStream(expectedLength)
            : new MemoryStream();

        byte[] buffer = ArrayPool<byte>.Shared.Rent(8192);
        try
        {
            var totalWritten = 0;
            while (true)
            {
                int bytesRead = codecStream.Read(buffer, 0, buffer.Length);
                if (bytesRead <= 0)
                    break;

                totalWritten += bytesRead;
                if (totalWritten > maxDecompressedSizeBytes)
                    throw new InvalidDataException(
                        $"Decompressed payload exceeds max allowed size ({maxDecompressedSizeBytes} bytes).");

                output.Write(buffer, 0, bytesRead);
            }

            if (expectedLength >= 0 && totalWritten != expectedLength)
                throw new InvalidDataException(
                    $"Expected decompressed length {expectedLength}, actual {totalWritten}.");

            return output.ToArray();
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    private sealed class NoneCompressionCodec : ICompressionCodec
    {
        /// <summary>
        ///     Gets the codec identifier.
        /// </summary>
        public CompressionCodec Codec => CompressionCodec.None;

        /// <summary>
        ///     Returns a copy of the input payload without compression.
        /// </summary>
        /// <param name="input">The payload bytes to copy.</param>
        /// <param name="level">The requested compression level.</param>
        /// <returns>The copied payload bytes.</returns>
        public byte[] Compress(ReadOnlySpan<byte> input, CompressionLevel level)
        {
            return input.ToArray();
        }

        /// <summary>
        ///     Validates and returns an uncompressed payload copy.
        /// </summary>
        /// <param name="input">The payload bytes to validate and copy.</param>
        /// <param name="expectedLength">The expected payload length, or a negative value to skip exact-length validation.</param>
        /// <param name="maxDecompressedSizeBytes">The maximum allowed payload size in bytes.</param>
        /// <returns>The copied payload bytes.</returns>
        public byte[] Decompress(ReadOnlySpan<byte> input, int expectedLength, int maxDecompressedSizeBytes)
        {
            if (input.Length > maxDecompressedSizeBytes)
                throw new InvalidDataException(
                    $"Decompressed payload exceeds max allowed size ({maxDecompressedSizeBytes} bytes).");

            if (expectedLength >= 0 && expectedLength != input.Length)
                throw new InvalidDataException(
                    $"Expected decompressed length {expectedLength}, actual {input.Length}.");

            return input.ToArray();
        }
    }

    private sealed class BrotliCompressionCodec : ICompressionCodec
    {
        /// <summary>
        ///     Gets the codec identifier.
        /// </summary>
        public CompressionCodec Codec => CompressionCodec.Brotli;

        /// <summary>
        ///     Compresses payload bytes using Brotli.
        /// </summary>
        /// <param name="input">The payload bytes to compress.</param>
        /// <param name="level">The compression level.</param>
        /// <returns>The compressed payload bytes.</returns>
        public byte[] Compress(ReadOnlySpan<byte> input, CompressionLevel level)
        {
            return CompressWithCodecStream(input, stream => new BrotliStream(stream, level, true));
        }

        /// <summary>
        ///     Decompresses Brotli-compressed payload bytes.
        /// </summary>
        /// <param name="input">The compressed payload bytes.</param>
        /// <param name="expectedLength">
        ///     The expected decompressed byte length, or a negative value to skip exact-length
        ///     validation.
        /// </param>
        /// <param name="maxDecompressedSizeBytes">The maximum allowed decompressed byte length.</param>
        /// <returns>The decompressed payload bytes.</returns>
        public byte[] Decompress(ReadOnlySpan<byte> input, int expectedLength, int maxDecompressedSizeBytes)
        {
            return DecompressWithCodecStream(input,
                stream => new BrotliStream(stream, CompressionMode.Decompress, true), expectedLength,
                maxDecompressedSizeBytes);
        }
    }

    private sealed class DeflateCompressionCodec : ICompressionCodec
    {
        /// <summary>
        ///     Gets the codec identifier.
        /// </summary>
        public CompressionCodec Codec => CompressionCodec.Deflate;

        /// <summary>
        ///     Compresses payload bytes using Deflate.
        /// </summary>
        /// <param name="input">The payload bytes to compress.</param>
        /// <param name="level">The compression level.</param>
        /// <returns>The compressed payload bytes.</returns>
        public byte[] Compress(ReadOnlySpan<byte> input, CompressionLevel level)
        {
            return CompressWithCodecStream(input, stream => new DeflateStream(stream, level, true));
        }

        /// <summary>
        ///     Decompresses Deflate-compressed payload bytes.
        /// </summary>
        /// <param name="input">The compressed payload bytes.</param>
        /// <param name="expectedLength">
        ///     The expected decompressed byte length, or a negative value to skip exact-length
        ///     validation.
        /// </param>
        /// <param name="maxDecompressedSizeBytes">The maximum allowed decompressed byte length.</param>
        /// <returns>The decompressed payload bytes.</returns>
        public byte[] Decompress(ReadOnlySpan<byte> input, int expectedLength, int maxDecompressedSizeBytes)
        {
            return DecompressWithCodecStream(input,
                stream => new DeflateStream(stream, CompressionMode.Decompress, true), expectedLength,
                maxDecompressedSizeBytes);
        }
    }
}