namespace Just.Core;

public enum Base32EncodeOptions
{
    None = 0x00,
    LowerCase = 0x01,
    NoPadding = 0x02,
    LowerCaseNoPadding = 0x03,
}

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

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

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

        var size = Encode(input, output, options);

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

    [Pure]
    public static int Encode(ReadOnlySpan<byte> input, Span<char> output, Base32EncodeOptions options = Base32EncodeOptions.None)
    {
        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];

        ReadOnlySpan<char> alphabet = (options & Base32EncodeOptions.LowerCase) == Base32EncodeOptions.LowerCase
            ? AlphabetLower
            : Alphabet;
        Span<byte> alphabetKeys = stackalloc byte[8];

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

            output[i++] = alphabet[alphabetKeys[0]];
            output[i++] = alphabet[alphabetKeys[1]];

            if (numCharsToOutput < 3)
            {
                i = FillWithPadding(output, i, numCharsToOutput, options);
                break;
            }
            output[i++] = alphabet[alphabetKeys[2]];
            output[i++] = alphabet[alphabetKeys[3]];

            if (numCharsToOutput < 5)
            {
                i = FillWithPadding(output, i, numCharsToOutput, options);
                break;
            }
            output[i++] = alphabet[alphabetKeys[4]];

            if (numCharsToOutput < 6)
            {
                i = FillWithPadding(output, i, numCharsToOutput, options);
                break;
            }
            output[i++] = alphabet[alphabetKeys[5]];
            output[i++] = alphabet[alphabetKeys[6]];

            if (numCharsToOutput < 8)
            {
                i = FillWithPadding(output, i, numCharsToOutput, options);
                break;
            }
            output[i++] = alphabet[alphabetKeys[7]];
        }

        return i;

        static int FillWithPadding(Span<char> output, int i, int numCharsToOutput, Base32EncodeOptions options)
        {
            if ((options & Base32EncodeOptions.NoPadding) == Base32EncodeOptions.NoPadding)
            {
                return i;
            }
            var pad = 8 - numCharsToOutput;
            output[i..(i+pad)].Fill(Padding);
            return i + pad;
        }
    }

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

        var outputLength = 5 * input.Length / 8;
        Span<byte> output = outputLength > MaxBytesStack
            ? new byte[outputLength]
            : stackalloc 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 rem = input.Length % 8;
        if (rem == 1 || rem == 3 || rem == 6)
        {
            throw new FormatException("Invalid Base32 string length.");
        }

        var outputLength = 5 * input.Length / 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
            ? new char[input.Length]
            : stackalloc 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)
            {
                if ((byteIndex & ~(0x1f << (bitPos - bits))) > 0)
                {
                    throw new FormatException("Invalid Base32 string. Inconsistent tail bits.");
                }
                break;
            }
        }

        return outputIndex;
    }

    // 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 = input[offset++];
        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;
    }
}