namespace Just.Core;

public static class Base32
{
    public const string Alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
    public const char Padding = '=';
    public const int MaxBytesStack = 250;

    [Pure]
    public static string Encode(ReadOnlySpan<byte> input)
    {
        if (input.IsEmpty) return string.Empty;

        int outLength = 8 * ((input.Length + 4) / 5);
        Span<char> output = input.Length <= MaxBytesStack
            ? stackalloc char[outLength]
            : new char[outLength];

        var size = Encode(input, output);

        return new string(output[..size]);
    }

    [Pure]
    public static int Encode(ReadOnlySpan<byte> input, Span<char> output)
    {
        if (input.IsEmpty) return 0;

        int outputLength = 8 * ((input.Length + 4) / 5);
        if (output.Length < outputLength)
        {
            throw new ArgumentException("Encoded input can not fit in output span.", nameof(output));
        }

        output = output[..outputLength];

        int i = 0;
        ReadOnlySpan<char> alphabet = Alphabet;
        Span<byte> alphabetKeys = stackalloc byte[8];

        for (int offset = 0; offset < input.Length;)
        {
            alphabetKeys.Clear();
            int numCharsToOutput = GetNextGroup(input, ref offset, alphabetKeys);

            output[i++] = (numCharsToOutput > 0) ? alphabet[alphabetKeys[0]] : Padding;
            output[i++] = (numCharsToOutput > 1) ? alphabet[alphabetKeys[1]] : Padding;
            output[i++] = (numCharsToOutput > 2) ? alphabet[alphabetKeys[2]] : Padding;
            output[i++] = (numCharsToOutput > 3) ? alphabet[alphabetKeys[3]] : Padding;
            output[i++] = (numCharsToOutput > 4) ? alphabet[alphabetKeys[4]] : Padding;
            output[i++] = (numCharsToOutput > 5) ? alphabet[alphabetKeys[5]] : Padding;
            output[i++] = (numCharsToOutput > 6) ? alphabet[alphabetKeys[6]] : Padding;
            output[i++] = (numCharsToOutput > 7) ? alphabet[alphabetKeys[7]] : Padding;
        }
        return i;
    }

    [Pure]
    public static byte[] Decode(ReadOnlySpan<char> input)
    {
        input = input.TrimEnd(Padding);
        if (input.IsEmpty) return [];

        var outputLength = 5 * ((input.Length + 7) / 8);
        Span<byte> output = outputLength <= MaxBytesStack
            ? stackalloc byte[outputLength]
            : new byte[outputLength];

        var size = Decode(input, output);

        return output[..size].ToArray();
    }

    [Pure]
    public static int Decode(ReadOnlySpan<char> input, Span<byte> output)
    {
        input = input.TrimEnd(Padding);

        var outputLength = 5 * ((input.Length + 7) / 8);
        if (output.Length < outputLength)
        {
            throw new ArgumentException("Decoded input can not fit in output span.", nameof(output));
        }

        output = output[..outputLength];
        output.Clear();

        Span<char> inputspan = outputLength <= MaxBytesStack
            ? stackalloc char[input.Length]
            : new char[input.Length];
        input.ToUpperInvariant(inputspan);

        int bitIndex = 0;
        int inputIndex = 0;
        int outputBits = 0;
        int outputIndex = 0;
        int bitPos;
        int outBitPos;
        int bits;

        ReadOnlySpan<char> alphabet = Alphabet;
        while (inputIndex < input.Length)
        {
            var byteIndex = alphabet.IndexOf(inputspan[inputIndex]);
            if (byteIndex < 0)
            {
                throw new FormatException("Provided string contains invalid characters.");
            }

            bitPos = 5 - bitIndex;
            outBitPos = 8 - outputBits;
            bits = bitPos < outBitPos ? bitPos : outBitPos;

            output[outputIndex] <<= bits;
            output[outputIndex] |= (byte)(byteIndex >> (bitPos - bits));

            outputBits += bits;
            if (outputBits >= 8)
            {
                outputIndex++;
                outputBits = 0;
            }

            bitIndex += bits;
            if (bitIndex >= 5)
            {
                inputIndex++;
                bitIndex = 0;
            }
            else if (inputIndex == input.Length -1) break;
        }

        return outputIndex + (outputBits + 7) / 8;
    }

    // returns the number of bytes that were output
    [Pure]
    private static int GetNextGroup(ReadOnlySpan<byte> input, ref int offset, Span<byte> alphabetKeys)
    {
        var retVal = (input.Length - offset) switch
        {
            1 => 2,
            2 => 4,
            3 => 5,
            4 => 7,
            _ => 8,
        };
        uint b1 = (offset < input.Length) ? input[offset++] : 0U;
        uint b2 = (offset < input.Length) ? input[offset++] : 0U;
        uint b3 = (offset < input.Length) ? input[offset++] : 0U;
        uint b4 = (offset < input.Length) ? input[offset++] : 0U;
        uint b5 = (offset < input.Length) ? input[offset++] : 0U;

        alphabetKeys[0] = (byte)(b1 >> 3);
        alphabetKeys[1] = (byte)(((b1 & 0x07) << 2) | (b2 >> 6));
        alphabetKeys[2] = (byte)((b2 >> 1) & 0x1f);
        alphabetKeys[3] = (byte)(((b2 & 0x01) << 4) | (b3 >> 4));
        alphabetKeys[4] = (byte)(((b3 & 0x0f) << 1) | (b4 >> 7));
        alphabetKeys[5] = (byte)((b4 >> 2) & 0x1f);
        alphabetKeys[6] = (byte)(((b4 & 0x3) << 3) | (b5 >> 5));
        alphabetKeys[7] = (byte)(b5 & 0x1f);

        return retVal;
    }
}