using System.Numerics;
using System.Text;
using System.Text.Json;

namespace Encryption;

public class RsaKeyGenerator
{
    private static readonly Random _random = new(Environment.TickCount);
    private BigInteger _d;
    private BigInteger _e;
    private BigInteger _n;
    private BigInteger _p;
    private BigInteger _q;
    private BigInteger _r;

    public RsaKeyGenerator(int bitLength)
    {
        //The "e" value for low compute time RSA encryption.
        //Only has two bits of value 1.
        const int e = 0x10001;

        //Generating primes, checking if the GCD of (n-1)(p-1) and e is 1.
        do
        {
            _q = FindPrime(bitLength / 2);
        } while (_q % e == 1);

        do
        {
            _p = FindPrime(bitLength / 2);
        } while (_p % e == 1);

        //Setting n as QP, phi (represented here as r) to tortiary.
        _n = _q * _p;
        _r = (_p - 1) * (_q - 1);

        //Computing D such that ed = 1%x.
        _d = ModularInverse(e, _r);
        
        _e = e;
    }

    /// <summary>
    ///     Finds a prime of the given bit length, to be used as n and p in RSA key calculations.
    /// </summary>
    /// <param name="bitlength"></param>
    /// <returns></returns>
    private static BigInteger FindPrime(int bitlength)
    {
        //Generating a random number of bit length.
        if (bitlength % 8 != 0) throw new Exception("Invalid bit length for key given, cannot generate primes.");

        //Filling bytes with pseudorandom.
        var randomBytes = new byte[bitlength / 8 + 1];
        _random.NextBytes(randomBytes);
        //Making the extra byte 0x0 so the BigInts are unsigned (little endian).
        randomBytes[randomBytes.Length - 1] = 0x0;

        //Setting the bottom bit and top two bits of the number.
        //This ensures the number is odd, and ensures the high bit of N is set when generating keys.
        SetBitInByte(0, ref randomBytes[0]);
        SetBitInByte(7, ref randomBytes[randomBytes.Length - 2]);
        SetBitInByte(6, ref randomBytes[randomBytes.Length - 2]);

        while (true)
        {
            //Performing a Rabin-Miller primality test.
            var isPrime = RabinMillerTest(randomBytes, 40);
            if (isPrime)
            {
                break;
            }

            IncrementByteArrayLE(ref randomBytes, 2);
            var upper_limit = new byte[randomBytes.Length];

            //Clearing upper bit for unsigned, creating upper and lower bounds.
            upper_limit[randomBytes.Length - 1] = 0x0;
            var upper_limit_bi = new BigInteger(upper_limit);
            var lower_limit = upper_limit_bi - 20;
            var current = new BigInteger(randomBytes);

            if (lower_limit < current && current < upper_limit_bi)
                //Failed to find a prime, returning -1.
                //Reached limit with no solutions.
                return new BigInteger(-1);
        }

        //Returning working BigInt.
        return new BigInteger(randomBytes);
    }

    /// <summary>
    ///     A Rabin Miller primality test which returns true or false.
    /// </summary>
    /// <param name="num">The number to check for being likely prime.</param>
    /// <returns></returns>
    private static bool RabinMillerTest(BigInteger source, int certainty)
    {
        //Filter out basic primes.
        if (source == 2 || source == 3) return true;
        //Below 2, and % 0? Not prime.
        if (source < 2 || source % 2 == 0) return false;

        //Finding even integer below number.
        var d = source - 1;
        var s = 0;

        while (d % 2 == 0)
        {
            d /= 2;
            s += 1;
        }

        //Getting a random BigInt using bytes.
        var rng = new Random(Environment.TickCount);
        var bytes = new byte[source.ToByteArray().LongLength];
        BigInteger a;

        //Looping to check random factors.
        for (var i = 0; i < certainty; i++)
        {
            do
            {
                //Generating new random bytes to check as a factor.
                rng.NextBytes(bytes);
                a = new BigInteger(bytes);
            } while (a < 2 || a >= source - 2);

            //Checking for x=1 or x=s-1.
            var x = BigInteger.ModPow(a, d, source);
            if (x == 1 || x == source - 1) continue;

            //Iterating to check for prime.
            for (var r = 1; r < s; r++)
            {
                x = BigInteger.ModPow(x, 2, source);
                if (x == 1)
                    return false;
                if (x == source - 1) break;
            }

            if (x != source - 1) return false;
        }

        //All tests have failed to prove composite, so return prime.
        return true;
    }

    /// <summary>
    ///     An overload wrapper for the RabinMillerTest which accepts a byte array.
    /// </summary>
    /// <param name="bytes"></param>
    /// <param name="acc_amt"></param>
    /// <returns></returns>
    private static bool RabinMillerTest(byte[] bytes, int acc_amt)
    {
        var b = new BigInteger(bytes);
        return RabinMillerTest(b, acc_amt);
    }

    /// <summary>
    ///     Performs a modular inverse on u and v,
    ///     such that d = gcd(u,v);
    /// </summary>
    /// <returns>D, such that D = gcd(u,v).</returns>
    private static BigInteger ModularInverse(BigInteger u, BigInteger v)
    {
        //Declaring new variables on the heap.
        BigInteger inverse, u1, u3, v1, v3, t1, t3, q = new();
        //Staying on the stack, quite small, so no need for extra memory time.
        BigInteger iteration;

        //Stating initial variables.
        u1 = 1;
        u3 = u;
        v1 = 0;
        v3 = v;

        //Beginning iteration.
        iteration = 1;
        while (v3 != 0)
        {
            //Divide and sub q, t3 and t1.
            q = u3 / v3;
            t3 = u3 % v3;
            t1 = u1 + q * v1;

            //Swap variables for next pass.
            u1 = v1;
            v1 = t1;
            u3 = v3;
            v3 = t3;
            iteration = -iteration;
        }

        if (u3 != 1)
            //No inverse, return 0.
            return 0;
        if (iteration < 0)
            inverse = v - u1;
        else
            inverse = u1;

        //Return.
        return inverse;
    }


    /// <summary>
    ///     Returns the greatest common denominator of both BigIntegers given.
    /// </summary>
    /// <returns>The GCD of A and B.</returns>
    private static BigInteger GCD(BigInteger a, BigInteger b)
    {
        //Looping until the numbers are zero values.
        while (a != 0 && b != 0)
            if (a > b)
                a %= b;
            else
                b %= a;

        //Returning check.
        return a == 0 ? b : a;
    }

    /// <summary>
    ///     Sets a bit in a given ref byte, using an index from 0-7 from the right.
    /// </summary>
    /// <param name="bitNumFromRight">The index of the bit number from the lesser side of the byte.</param>
    /// <param name="toSet">The referenced byte to set.</param>
    private static void SetBitInByte(int bitNumFromRight, ref byte toSet)
    {
        var mask = (byte)(1 << bitNumFromRight);
        toSet |= mask;
    }

    /// <summary>
    ///     Increments the byte array as a whole, by a given amount. Assumes little endian.
    ///     Assumes unsigned randomBytes.
    /// </summary>
    private static void IncrementByteArrayLE(ref byte[] randomBytes, int amt)
    {
        var n = new BigInteger(randomBytes);
        n += amt;
        randomBytes = n.ToByteArray();
    }

    /// <summary>
    ///     Decrements the byte array as a whole, by a given amount. Assumes little endian.
    ///     Assumes unsigned randomBytes.
    /// </summary>
    private static void DecrementByteArrayLE(ref byte[] randomBytes, int amt)
    {
        var n = new BigInteger(randomBytes);
        n -= amt;
        randomBytes = n.ToByteArray();
    }

    public RsaKeyPair GetKeys()
    {
        return new RsaKeyPair(new RsaPrivateKey(_d, _n), new RsaPublicKey(_e, _n));
    }
}

public record RsaPublicKey(BigInteger E, BigInteger N)
{
    public override string ToString()
    {
        return Convert.ToBase64String(Encoding.UTF8.GetBytes($"e:{E};n:{N}"));
    }

    public string ToJsonString()
    {
        var e = E.ToString();
        var n = N.ToString();
        
        var obj = new Dictionary<string, string>
        {
            ["E"] = e,
            ["N"] = n,
        };
        
        return JsonSerializer.Serialize(obj);
    }
}

public record RsaPrivateKey(BigInteger D, BigInteger N)
{
    public override string ToString()
    {
        return Convert.ToBase64String(Encoding.UTF8.GetBytes($"d:{D};n:{N}"));
    }
}

public record RsaKeyPair(RsaPrivateKey PrivateKey, RsaPublicKey PublicKey);