using System;
using System.IO;
using System.Security.Cryptography;
using System.Text;

namespace ShrinkDataSaver
{
    public static class SaveEncryptor
    {
        private const int KeySize = 256;
        private const int BlockSize = 128;
        private const int IvBytes = 16;
        private const int SaltBytes = 16;
        private const int Iterations = 10000;

        public static byte[] Encrypt(byte[] data, string password)
        {
            if (data == null) throw new ArgumentNullException(nameof(data));
            if (string.IsNullOrEmpty(password)) throw new ArgumentException("Password must not be empty.");

            var salt = GenerateRandom(SaltBytes);
            var iv = GenerateRandom(IvBytes);
            var key = DeriveKey(password, salt);

            using var aes = CreateAes(key, iv);
            using var encryptor = aes.CreateEncryptor();
            var cipher = encryptor.TransformFinalBlock(data, 0, data.Length);

            var result = new byte[SaltBytes + IvBytes + cipher.Length];
            Buffer.BlockCopy(salt, 0, result, 0, SaltBytes);
            Buffer.BlockCopy(iv, 0, result, SaltBytes, IvBytes);
            Buffer.BlockCopy(cipher, 0, result, SaltBytes + IvBytes, cipher.Length);
            return result;
        }

        public static byte[] Decrypt(byte[] data, string password)
        {
            if (data == null) throw new ArgumentNullException(nameof(data));
            if (string.IsNullOrEmpty(password)) throw new ArgumentException("Password must not be empty.");
            if (data.Length < SaltBytes + IvBytes)
                throw new ArgumentException("Data is too short to be valid encrypted content.");

            var salt = new byte[SaltBytes];
            var iv = new byte[IvBytes];
            var cipher = new byte[data.Length - SaltBytes - IvBytes];

            Buffer.BlockCopy(data, 0, salt, 0, SaltBytes);
            Buffer.BlockCopy(data, SaltBytes, iv, 0, IvBytes);
            Buffer.BlockCopy(data, SaltBytes + IvBytes, cipher, 0, cipher.Length);

            var key = DeriveKey(password, salt);

            using var aes = CreateAes(key, iv);
            using var decryptor = aes.CreateDecryptor();
            return decryptor.TransformFinalBlock(cipher, 0, cipher.Length);
        }


        private static byte[] DeriveKey(string password, byte[] salt)
        {
            using var deriveBytes = new Rfc2898DeriveBytes(
                Encoding.UTF8.GetBytes(password), salt, Iterations, HashAlgorithmName.SHA256);
            return deriveBytes.GetBytes(KeySize / 8);
        }

        private static Aes CreateAes(byte[] key, byte[] iv)
        {
            var aes = Aes.Create();
            aes.KeySize = KeySize;
            aes.BlockSize = BlockSize;
            aes.Mode = CipherMode.CBC;
            aes.Padding = PaddingMode.PKCS7;
            aes.Key = key;
            aes.IV = iv;
            return aes;
        }

        private static byte[] GenerateRandom(int length)
        {
            var bytes = new byte[length];
            using var rng = RandomNumberGenerator.Create();
            rng.GetBytes(bytes);
            return bytes;
        }
    }
}