using System.Text;

namespace HuffmanEditor;

public static class Huffman
{
    public static byte[] EncodeString(string str)
    {
        return GenerateBytes(str);
    }

    public static string DecodeString(byte[] bytes)
    {
        return DecypherBytes(bytes);
    }

    private static HuffmanNode GenerateTree(string str)
    {
        var frequencies = new int[65536];
        
        foreach (var character in str)
            frequencies[character]++;
        
        var priorityQueue = new PriorityQueue<HuffmanNode, int>();
        
        for (var i = 0; i < frequencies.Length; i++)
            if (frequencies[i] != 0)
                priorityQueue.Enqueue(new HuffmanNode((char)i, frequencies[i]), frequencies[i]);
        
        while (priorityQueue.Count > 1)
        {
            var left = priorityQueue.Dequeue();
            var right = priorityQueue.Dequeue();
            priorityQueue.Enqueue(new HuffmanNode(left, right), left.Frequency + right.Frequency);
        }

        return priorityQueue.Dequeue();
    }

    private static SortedDictionary<char, ushort> GenerateLengthDictionary(HuffmanNode node)
    {
        var dict = new SortedDictionary<char, ushort>();

        void TraverseDFS(HuffmanNode node, ushort depth)
        {
            if (node == null) return;
            if (node.Symbol != null)
            {
                dict.Add(node.Symbol.Value, depth);
                return;
            }
            TraverseDFS(node.Left!, (ushort)(depth + 1));
            TraverseDFS(node.Right!, (ushort)(depth + 1));
        }

        TraverseDFS(node.Left!, 1);
        TraverseDFS(node.Right!, 1);

        return dict;
    }

    private static Dictionary<char, HuffmanCode> RecreateDictionary(SortedDictionary<char, ushort> lenDict)
    {
        var sorted = lenDict.OrderBy(x => x.Value).ThenBy(x => x.Key).ToList();
        var codes = new Dictionary<char, HuffmanCode>();

        var currentCode = new HuffmanCode { Code = 0, Length = 0 };

        foreach (var code in sorted)
        {
            if (code.Value > currentCode.Length)
            {
                currentCode.Code <<= (code.Value - currentCode.Length);
                currentCode.Length = code.Value;
            }
            codes.Add(code.Key, currentCode);
            currentCode.Code++;
        }

        return codes;
    }

    private static HuffmanNode RecreateTree(SortedDictionary<char, ushort> lenDict)
    {
        var dict = RecreateDictionary(lenDict);

        var root = new HuffmanNode();

        foreach (var code in dict)
        {
            var currentNode = root;

            for (var i = code.Value.Length - 1; i >= 0; i--)
            {
                if (((code.Value.Code >> i) & 1) == 1)
                {
                    currentNode.Right ??= new HuffmanNode();
                    currentNode = currentNode.Right;
                }
                else
                {
                    currentNode.Left ??= new HuffmanNode();
                    currentNode = currentNode.Left;
                }
            }

            currentNode.Symbol = code.Key;
        }

        return root;
    }

    private static byte[] GenerateBytes(string str)
    {
        var lenDict = GenerateLengthDictionary(GenerateTree(RSFFilter(str)));
        //var lenDict = GenerateLengthDictionary(GenerateTree(str));

        var bytes = new List<byte>
        {
            (byte)(lenDict.Count >> 8),
            (byte)(lenDict.Count & 0xFF),
        };

        foreach (var kvp in lenDict)
            bytes.AddRange(
                [
                    (byte)(kvp.Key >> 8),
                    (byte)(kvp.Key & 0xFF),
                    (byte)(kvp.Value >> 8),
                    (byte)(kvp.Value & 0xFF),
                ]
            );

        bytes.AddRange(
                [
                    (byte)(str.Length >> 24),
                    (byte)(str.Length >> 16 & 0xFF),
                    (byte)(str.Length >> 8 & 0xFF),
                    (byte)(str.Length & 0xFF),
                ]
            );

        var dict = RecreateDictionary(lenDict);

        byte writeBuffer = 0, counter = 0;

        foreach (var code in str)
        {
            var thisCode = dict[code];
            for (var i = 0; i < thisCode.Length; i++)
            {
                writeBuffer = (byte)((writeBuffer << 1) | ((thisCode.Code >> (thisCode.Length - i - 1)) & 1));
                counter++;
                if (counter == 8)
                {
                    counter = 0;
                    bytes.Add(writeBuffer);
                    writeBuffer = 0;
                }
            }
        }
        while (true)
        {
            writeBuffer = (byte)(writeBuffer << 1);
            counter++;
            if (counter == 8)
            {
                bytes.Add(writeBuffer);
                break;
            }
        }

        return bytes.ToArray();
    }

    private static string DecypherBytes(byte[] bytes)
    {
        var currentByte = 0;
        var letters = bytes[currentByte++] << 8 | bytes[currentByte++];
        var result = string.Empty;

        var lenDict = new SortedDictionary<char, ushort>();

        for (var i = 0; i < letters; i++)
            lenDict.Add((char)(bytes[currentByte++] << 8 | bytes[currentByte++]), (ushort)(bytes[currentByte++] << 8 | bytes[currentByte++]));

        var characters = bytes[currentByte++] << 24 | bytes[currentByte++] << 16 | bytes[currentByte++] << 8 | bytes[currentByte++];

        var tree = RecreateTree(lenDict);
        var currentNode = tree;
        
        for (var i = 0; i < characters - 1; )
        {
            var thisByte = bytes[currentByte++];

            for (var j = 0; j < 8; j++)
            {
                if ((thisByte >> (7 - j) & 0x01) == 1)
                    currentNode = currentNode.Right;
                else
                    currentNode = currentNode.Left;
                if (currentNode == null) throw new ArgumentException("IDK man you done bad");
                if (currentNode.Symbol != null)
                {
                    result += currentNode.Symbol.Value;
                    currentNode = tree;
                    i++;
                }
            }
        }

        return RSFUnfilter(result);
        //return result;
    }

    private static string RSFFilter(string input)
    {
        if (string.IsNullOrEmpty(input)) return input;

        // Initialize with input length to reduce resizing overhead
        var sb = new StringBuilder(input.Length);
        var i = 0;

        while (i < input.Length)
        {
            var currentChar = input[i];
            var count = 1;

            // Count consecutive identical characters
            while (i + count < input.Length && input[i + count] == currentChar)
            {
                count++;
            }

            // Apply transformation if repeat count > 3 (4 or more occurrences)
            if (count > 3)
            {
                sb.Append('\uE000');
                sb.Append(count);
                sb.Append(currentChar);
            }
            else
            {
                // Append the characters as-is (1, 2, or 3 times)
                sb.Append(currentChar, count);
            }

            i += count;
        }

        return sb.ToString();
    }

    private static string RSFUnfilter(string input)
    {
        if (string.IsNullOrEmpty(input)) return input;

        var sb = new StringBuilder();
        int i = 0;

        while (i < input.Length)
        {
            if (input[i] == '\uE000')
            {
                i++; // Skip the marker
                int start = i;

                // Parse the numeric repeatCount
                while (i < input.Length && char.IsDigit(input[i]))
                {
                    i++;
                }

                if (int.TryParse(input.AsSpan(start, i - start), out var count) && i < input.Length)
                {
                    var repeatChar = input[i];
                    sb.Append(repeatChar, count);
                    i++; // Move past the character being repeated
                }
            }
            else
            {
                sb.Append(input[i]);
                i++;
            }
        }

        return sb.ToString();
    }

    private struct HuffmanCode
    {
        public int Code;
        public ushort Length;
    }
}

public class HuffmanNode : IComparable<HuffmanNode>
{
    public char? Symbol
    {
        get; set;
    }

    public int Frequency
    {
        get; set;
    }

    public HuffmanNode? Left
    {
        get; set;
    }
    public HuffmanNode? Right
    {
        get; set;
    }

    public HuffmanNode(char symbol, int frequency)
    {
        Symbol = symbol;
        Frequency = frequency;
    }

    public HuffmanNode(HuffmanNode left, HuffmanNode right)
    {
        Symbol = null;
        Left = left;
        Right = right;
        Frequency = left.Frequency + right.Frequency;
    }

    public HuffmanNode()
    {
        Frequency = 0;
        Symbol = null;
        Left = null;
        Right = null;
    }

    public int CompareTo(HuffmanNode? other)
    {
        if (other == null) return 1;
        return Frequency.CompareTo(other.Frequency);
    }
}