namespace DbaMessageInserter;

/// <summary>
/// Builds optimal binary length-limited Huffman code lengths with the
/// package-merge algorithm. DEFLATE limits literal/length and distance codes
/// to 15 bits, and code-length alphabet codes to 7 bits.
/// </summary>
internal static class LengthLimitedHuffmanBuilder
{
    public static byte[] Build(
        IReadOnlyList<long> frequencies,
        int maximumCodeLength)
    {
        if (frequencies is null)
            throw new ArgumentNullException(nameof(frequencies));
        if (frequencies.Count < 2)
            throw new ArgumentException("At least two symbol slots are required.", nameof(frequencies));
        if (maximumCodeLength <= 0 || maximumCodeLength > 15)
            throw new ArgumentOutOfRangeException(nameof(maximumCodeLength));

        var adjustedFrequencies = new long[frequencies.Count];
        var activeSymbols = new List<int>();

        for (int symbol = 0; symbol < frequencies.Count; symbol++)
        {
            long frequency = frequencies[symbol];
            if (frequency < 0)
                throw new ArgumentOutOfRangeException(nameof(frequencies), "Frequencies cannot be negative.");

            adjustedFrequencies[symbol] = frequency;
            if (frequency > 0)
                activeSymbols.Add(symbol);
        }

        // A canonical Huffman alphabet needs at least two leaves for the most
        // broadly compatible DEFLATE representation. Dummy leaves are never
        // emitted as data symbols.
        if (activeSymbols.Count == 0)
        {
            adjustedFrequencies[0] = 1;
            adjustedFrequencies[1] = 1;
            activeSymbols.Add(0);
            activeSymbols.Add(1);
        }
        else if (activeSymbols.Count == 1)
        {
            int dummySymbol = activeSymbols[0] == 0 ? 1 : 0;
            adjustedFrequencies[dummySymbol] = 1;
            activeSymbols.Add(dummySymbol);
        }

        int activeCount = activeSymbols.Count;
        if (activeCount > (1 << maximumCodeLength))
        {
            throw new ArgumentException(
                $"{activeCount} active symbols cannot fit in {maximumCodeLength}-bit Huffman codes.",
                nameof(frequencies));
        }

        int sequence = 0;
        var leaves = new List<PackageNode>(activeCount);
        foreach (int symbol in activeSymbols)
        {
            leaves.Add(new PackageNode(
                adjustedFrequencies[symbol],
                symbol,
                sequence++,
                symbol,
                null,
                null));
        }
        leaves.Sort(PackageNodeComparer.Instance);

        int selectedNodeCount = checked(2 * activeCount - 2);
        List<PackageNode> current = leaves;

        // The final list represents the maximumCodeLength-th package-merge
        // level. Keeping only 2n-2 cheapest entries at every level is safe and
        // bounds both memory and runtime.
        for (int level = 1; level < maximumCodeLength; level++)
        {
            var packages = new List<PackageNode>(current.Count / 2);

            for (int index = 0; index + 1 < current.Count; index += 2)
            {
                PackageNode left = current[index];
                PackageNode right = current[index + 1];
                packages.Add(new PackageNode(
                    checked(left.Weight + right.Weight),
                    Math.Min(left.MinimumSymbol, right.MinimumSymbol),
                    sequence++,
                    -1,
                    left,
                    right));
            }

            packages.Sort(PackageNodeComparer.Instance);
            current = MergeAndLimit(leaves, packages, selectedNodeCount);
        }

        if (current.Count < selectedNodeCount)
            throw new InvalidOperationException("Package-merge produced too few nodes.");

        var lengths = new byte[frequencies.Count];
        var stack = new Stack<PackageNode>();

        for (int index = 0; index < selectedNodeCount; index++)
        {
            stack.Push(current[index]);

            while (stack.Count > 0)
            {
                PackageNode node = stack.Pop();
                if (node.Symbol >= 0)
                {
                    if (lengths[node.Symbol] == byte.MaxValue)
                        throw new InvalidOperationException("Huffman code length overflow.");
                    lengths[node.Symbol]++;
                }
                else
                {
                    stack.Push(node.Right!);
                    stack.Push(node.Left!);
                }
            }
        }

        ValidateCompleteTree(lengths, maximumCodeLength, activeCount);
        return lengths;
    }

    private static List<PackageNode> MergeAndLimit(
        IReadOnlyList<PackageNode> leaves,
        IReadOnlyList<PackageNode> packages,
        int maximumCount)
    {
        var result = new List<PackageNode>(maximumCount);
        int leafIndex = 0;
        int packageIndex = 0;

        while (result.Count < maximumCount &&
               (leafIndex < leaves.Count || packageIndex < packages.Count))
        {
            bool takeLeaf = packageIndex >= packages.Count ||
                (leafIndex < leaves.Count &&
                 PackageNodeComparer.Instance.Compare(
                     leaves[leafIndex],
                     packages[packageIndex]) <= 0);

            result.Add(takeLeaf
                ? leaves[leafIndex++]
                : packages[packageIndex++]);
        }

        return result;
    }

    private static void ValidateCompleteTree(
        IReadOnlyList<byte> lengths,
        int maximumCodeLength,
        int expectedActiveCount)
    {
        var counts = new int[maximumCodeLength + 1];
        int actualActiveCount = 0;

        foreach (byte length in lengths)
        {
            if (length == 0) continue;
            if (length > maximumCodeLength)
                throw new InvalidOperationException("Package-merge exceeded the code-length limit.");

            counts[length]++;
            actualActiveCount++;
        }

        if (actualActiveCount != expectedActiveCount)
            throw new InvalidOperationException("Package-merge lost or duplicated a Huffman leaf.");

        int unusedCodeSpace = 1;
        for (int bits = 1; bits <= maximumCodeLength; bits++)
        {
            unusedCodeSpace = checked((unusedCodeSpace << 1) - counts[bits]);
            if (unusedCodeSpace < 0)
                throw new InvalidOperationException("Package-merge produced an oversubscribed tree.");
        }

        if (unusedCodeSpace != 0)
            throw new InvalidOperationException("Package-merge produced an incomplete tree.");
    }

    private sealed record PackageNode(
        long Weight,
        int MinimumSymbol,
        int Sequence,
        int Symbol,
        PackageNode? Left,
        PackageNode? Right);

    private sealed class PackageNodeComparer : IComparer<PackageNode>
    {
        public static PackageNodeComparer Instance { get; } = new();

        public int Compare(PackageNode? x, PackageNode? y)
        {
            if (ReferenceEquals(x, y)) return 0;
            if (x is null) return -1;
            if (y is null) return 1;

            int result = x.Weight.CompareTo(y.Weight);
            if (result != 0) return result;

            result = x.MinimumSymbol.CompareTo(y.MinimumSymbol);
            if (result != 0) return result;

            return x.Sequence.CompareTo(y.Sequence);
        }
    }
}
