using System;
using System.Linq;
using System.Numerics;
using System.Text;

namespace FrostFS.SDK.Cryptography;

public static class Base58
{
    public const string Alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

    public static byte[] Base58CheckDecode(this string input)
    {
        if (input is null)
            throw new ArgumentNullException(nameof(input));
        
        byte[] buffer = Decode(input);

        if (buffer.Length < 4) 
            throw new FormatException();
        
        byte[] checksum = buffer[0..(buffer.Length - 4)].Sha256().Sha256();

        if (!buffer.AsSpan(buffer.Length - 4).SequenceEqual(checksum[..4].AsSpan()))
            throw new FormatException();
       
        var ret = buffer[..^4];
        Array.Clear(buffer, 0, buffer.Length);
        return ret;
    }

    public static string Base58CheckEncode(this ReadOnlySpan<byte> data)
    {
        byte[] checksum = data.ToArray().Sha256().Sha256();
        Span<byte> buffer = stackalloc byte[data.Length + 4];
        data.CopyTo(buffer);
        checksum[..4].AsSpan().CopyTo(buffer[data.Length..]);
        var ret = Encode(buffer);
        buffer.Clear();
        return ret;
    }

    public static byte[] Decode(string input)
    {
        // Decode Base58 string to BigInteger 
        var bi = BigInteger.Zero;
        for (int i = 0; i < input.Length; i++)
        {
            int digit = Alphabet.IndexOf(input[i]);
            if (digit < 0)
                throw new FormatException($"Invalid Base58 character '{input[i]}' at position {i}");
          
            bi = bi * Alphabet.Length + digit;
        }

        int leadingZeroCount = input.TakeWhile(c => c == Alphabet[0]).Count();
        var leadingZeros = new byte[leadingZeroCount];
                
        if (bi.IsZero)
            return leadingZeros;

        var bytesBigEndian = bi.ToByteArray().Reverse();

        var firstNonZeroIndex = 0;

        while(bytesBigEndian.ElementAt(firstNonZeroIndex) == 0x0)
            firstNonZeroIndex++;
        
        var bytesWithoutLeadingZeros = bytesBigEndian.Skip(firstNonZeroIndex).ToArray();

        return ArrayHelper.Concat(leadingZeros, bytesWithoutLeadingZeros);
    }

    public static string Encode(ReadOnlySpan<byte> input)
    {
        var data = input.ToArray().Reverse().Concat(new byte[] { 0 }).ToArray();
        BigInteger value = new(data);

        // Encode BigInteger to Base58 string
        var sb = new StringBuilder();

        while (value > 0)
        {
            value = BigInteger.DivRem(value, Alphabet.Length, out var remainder);
            sb.Insert(0, Alphabet[(int)remainder]);
        }

        // Append `1` for each leading 0 byte
        for (int i = 0; i < input.Length && input[i] == 0; i++)
        {
            sb.Insert(0, Alphabet[0]);
        }
        return sb.ToString();
    }
}