using PCC.App.BinaryFormatter;
using System.Security;
using System.Security.Cryptography;
using System.Text;

namespace PCC.App.Security;

public static class RsaUtility
{
    public static (string pub, string pri) GeneratePem(int bits)
    {
        var rsa = RSA.Create(bits);
        var pub = rsa.ExportRSAPublicKeyPem();
        var pri = rsa.ExportRSAPrivateKeyPem();
        return (pub, pri);
    }

    public static (byte[] pub, byte[] pri) GeneratePKCS1(int bits)
    {
        var rsa = RSA.Create(bits);
        var pub = rsa.ExportRSAPublicKey();
        var pri = rsa.ExportRSAPrivateKey();
        return (pub, pri);
    }

    public static RSA FromPKCS1PrivateKey(byte[] privateKey)
    {
        var rsa = RSA.Create();
        rsa.ImportRSAPrivateKey(privateKey, out _);
        return rsa;
    }

    public static RSA FromPKCS1PublicKey(byte[] publicKey)
    {
        var rsa = RSA.Create();
        rsa.ImportRSAPublicKey(publicKey, out _);
        return rsa;
    }

    public static RSA FromPem(string pub)
    {
        var rsa = RSA.Create();
        rsa.ImportFromPem(pub);
        return rsa;
    }

    public static byte[] Encrypt(RSA encPub, byte[] payload, int skSize = 256, int ivSize = 128) => EncryptAndSignature(encPub, payload, null, skSize, ivSize);

    public static byte[] EncryptAndSignature(RSA encPub, ReadOnlySpan<byte> payload, RSA? sigPri, int skSize = 256, int ivSize = 128)
    {
        // EK= Rsa(encPub, [SK] + [IV] )
        // ED= AES(SK,IV,[rawData] + [Sign(sigPri,rawData)] )
        // ret= [EK] + [ED]

        // generate Symmetric key: SK,IV
        var sk = RandomNumberGenerator.GetBytes(skSize / 8);
        var iv = RandomNumberGenerator.GetBytes(ivSize / 8);

        //Debug.Print($"SK: {Convert.ToHexString(sk)}");
        //Debug.Print($"IV: {Convert.ToHexString(iv)}");

        using var aes = Aes.Create();
        aes.Mode = CipherMode.CBC;
        aes.Padding = PaddingMode.PKCS7;
        aes.KeySize = skSize;
        aes.BlockSize = ivSize;

        aes.Key = sk;
        aes.IV = iv;

        // EK= Rsa(encPub, [SK] + [IV])
        byte[] ek;
        {
            var msSkIv = new MemoryStream();
            var bw = new BinaryWriter(msSkIv, Encoding.UTF8, true);
            bw.WriteBlock(sk);
            bw.WriteBlock(iv);
            bw.Close();
            var skIv = msSkIv.ToArray();
            ek = encPub.Encrypt(skIv, RSAEncryptionPadding.Pkcs1);
        }

        // ED= AES(SK,IV,[rawData] + [Sign(sigPri,rawData)] )
        byte[] ed;
        {
            var msEd = new MemoryStream();
            using var cryptoTransform = aes.CreateEncryptor();
            var es = new CryptoStream(msEd, cryptoTransform, CryptoStreamMode.Write);
            var eBw = new BinaryWriter(es, Encoding.UTF8, true);
            eBw.WriteBlock(payload);

            if (sigPri != null)
            {
                var sign = sigPri.SignData(payload, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
                eBw.WriteBlock(sign);
            }

            eBw.Close();

            es.Close();
            ed = msEd.ToArray();
        }

        // return [EK] + [ED]
        byte[] ret;
        {
            var msRet = new MemoryStream();
            var bw = new BinaryWriter(msRet, Encoding.UTF8, true);
            bw.WriteBlock(ek);
            bw.WriteBlock(ed);
            bw.Close();
            ret = msRet.ToArray();
        }

        return ret;
    }

    public static byte[] Decrypt(RSA decPri, byte[] dataToDecrypt) => DecryptAndVerifySignature(decPri, dataToDecrypt, null);

    public static byte[] DecryptAndVerifySignature(RSA decPri, byte[] dataToDecrypt, RSA? verPub)
    {
        // extract [EK] [ED]
        byte[] ek, ed;
        {
            var ms = new MemoryStream(dataToDecrypt);
            var br = new BinaryReader(ms);

            ek = br.ReadBlock();
            ed = br.ReadBlock();
        }

        byte[] sk, iv;
        {
            // decrypt EK by pri
            var skIv = decPri.Decrypt(ek, RSAEncryptionPadding.Pkcs1);
            var msSkIv = new MemoryStream(skIv);
            var br = new BinaryReader(msSkIv);
            sk = br.ReadBlock();
            iv = br.ReadBlock();
        }

        //Debug.Print($"SK: {Convert.ToHexString(sk)}");
        //Debug.Print($"IV: {Convert.ToHexString(iv)}");

        // decData= [payload] + [Signature]
        byte[] decData;
        {
            // decrypt by SK,IV
            var aes = Aes.Create();
            aes.Mode = CipherMode.CBC;
            aes.Padding = PaddingMode.PKCS7;
            aes.KeySize = 128;
            aes.BlockSize = 128;

            aes.Key = sk;
            aes.IV = iv;

            var msEd = new MemoryStream(ed);
            var msDec = new MemoryStream();
            var ds = new CryptoStream(msEd, aes.CreateDecryptor(), CryptoStreamMode.Read);
            ds.CopyTo(msDec);
            ds.Close();
            decData = msDec.ToArray();
        }

        // extract [payload] [Signature]
        byte[] payload;
        byte[]? signature = null;
        {
            var ms = new MemoryStream(decData);
            var br = new BinaryReader(ms);
            payload = br.ReadBlock();
            if (verPub != null)
            {
                if (ms.Position != ms.Length) signature = br.ReadBlock();
                else throw new SecurityException("Missing signature");
            }
        }

        if (signature == null || verPub == null) return payload;

        // Validate signature
        var r = verPub.VerifyData(payload, signature, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
        if (r == false) throw new SecurityException("Signature verify fail");

        // return payload
        return payload;
    }
}