namespace GeoSus.Server;

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

// Šifrování komunikace - RSA handshake + AES-256-CBC session (kompatibilní s Unity)
public class ServerEncryption : IDisposable
{
    private readonly RSA _rsa;
    private readonly string _publicKeyPem;
    
    public ServerEncryption(int keySizeBits = 2048)
    {
        _rsa = RSA.Create(keySizeBits);
        _publicKeyPem = ExportPublicKeyPem();
    }
    
    public string PublicKeyPem => _publicKeyPem;
    
    // Dešifruje session key od klienta
    public (byte[] Key, byte[] IV) DecryptSessionKey(string encryptedKeyBase64, string encryptedIvBase64)
    {
        var encryptedKey = Convert.FromBase64String(encryptedKeyBase64);
        var encryptedIv = Convert.FromBase64String(encryptedIvBase64);
        
        // Používáme OaepSHA1 pro Unity kompatibilitu
        var key = _rsa.Decrypt(encryptedKey, RSAEncryptionPadding.OaepSHA1);
        var iv = _rsa.Decrypt(encryptedIv, RSAEncryptionPadding.OaepSHA1);
        
        return (key, iv);
    }
    
    private string ExportPublicKeyPem()
    {
        var publicKey = _rsa.ExportSubjectPublicKeyInfo();
        var base64 = Convert.ToBase64String(publicKey);
        var sb = new StringBuilder();
        sb.AppendLine("-----BEGIN PUBLIC KEY-----");
        for (int i = 0; i < base64.Length; i += 64)
        {
            sb.AppendLine(base64.Substring(i, Math.Min(64, base64.Length - i)));
        }
        sb.AppendLine("-----END PUBLIC KEY-----");
        return sb.ToString();
    }
    
    public void Dispose()
    {
        _rsa.Dispose();
    }
}

// Session šifrování pro konkrétní klientské spojení - AES-256-CBC + HMAC-SHA256
public class SessionCrypto : IDisposable
{
    private readonly byte[] _key;
    private readonly object _lock = new();
    
    public SessionCrypto(byte[] key, byte[] iv)
    {
        if (key.Length != 32) 
            throw new ArgumentException("Key musí být 32 bajtů pro AES-256");
        if (iv.Length != 16) 
            throw new ArgumentException("IV musí být 16 bajtů pro AES-CBC");
            
        _key = key;
        // IV se nepoužívá přímo - každá zpráva má svůj unikátní IV
    }
    
    // Šifruje zprávu s AES-CBC a HMAC
    public byte[] Encrypt(byte[] plaintext)
    {
        lock (_lock)
        {
            using var aes = Aes.Create();
            aes.Key = _key;
            aes.Mode = CipherMode.CBC;
            aes.Padding = PaddingMode.PKCS7;
            aes.GenerateIV(); // Generujeme nový IV pro každou zprávu
            
            byte[] ciphertext;
            using (var encryptor = aes.CreateEncryptor())
            {
                ciphertext = encryptor.TransformFinalBlock(plaintext, 0, plaintext.Length);
            }
            
            // Počítáme HMAC přes IV + ciphertext (používáme AES klíč pro HMAC)
            byte[] hmac;
            using (var hmacSha = new HMACSHA256(_key))
            {
                var dataToSign = new byte[aes.IV.Length + ciphertext.Length];
                Buffer.BlockCopy(aes.IV, 0, dataToSign, 0, aes.IV.Length);
                Buffer.BlockCopy(ciphertext, 0, dataToSign, aes.IV.Length, ciphertext.Length);
                hmac = hmacSha.ComputeHash(dataToSign);
            }
            
            // Výstup: [16 bytes IV][32 bytes HMAC][ciphertext]
            var result = new byte[16 + 32 + ciphertext.Length];
            Buffer.BlockCopy(aes.IV, 0, result, 0, 16);
            Buffer.BlockCopy(hmac, 0, result, 16, 32);
            Buffer.BlockCopy(ciphertext, 0, result, 48, ciphertext.Length);
            
            return result;
        }
    }
    
    // Dešifruje zprávu a validuje HMAC
    public byte[]? Decrypt(byte[] encrypted)
    {
        if (encrypted.Length < 48) return null; // 16 IV + 32 HMAC + min data
        
        try
        {
            var iv = new byte[16];
            var receivedHmac = new byte[32];
            var ciphertext = new byte[encrypted.Length - 48];
            
            Buffer.BlockCopy(encrypted, 0, iv, 0, 16);
            Buffer.BlockCopy(encrypted, 16, receivedHmac, 0, 32);
            Buffer.BlockCopy(encrypted, 48, ciphertext, 0, ciphertext.Length);
            
            // Ověříme HMAC (používáme AES klíč pro HMAC)
            byte[] computedHmac;
            using (var hmacSha = new HMACSHA256(_key))
            {
                var dataToVerify = new byte[16 + ciphertext.Length];
                Buffer.BlockCopy(iv, 0, dataToVerify, 0, 16);
                Buffer.BlockCopy(ciphertext, 0, dataToVerify, 16, ciphertext.Length);
                computedHmac = hmacSha.ComputeHash(dataToVerify);
            }
            
            // Constant-time porovnání
            if (!CryptographicOperations.FixedTimeEquals(receivedHmac, computedHmac))
            {
                return null; // HMAC mismatch - zpráva byla změněna
            }
            
            // Dešifrujeme
            using var aes = Aes.Create();
            aes.Key = _key;
            aes.IV = iv;
            aes.Mode = CipherMode.CBC;
            aes.Padding = PaddingMode.PKCS7;
            
            using var decryptor = aes.CreateDecryptor();
            return decryptor.TransformFinalBlock(ciphertext, 0, ciphertext.Length);
        }
        catch (CryptographicException)
        {
            return null;
        }
    }
    
    public void Dispose()
    {
        Array.Clear(_key, 0, _key.Length);
    }
}

// Klientská strana - generuje session key a šifruje ho RSA public key
public static class ClientEncryptionHelper
{
    public static (byte[] Key, byte[] IV) GenerateSessionKey()
    {
        var key = RandomNumberGenerator.GetBytes(32); // AES-256
        var iv = RandomNumberGenerator.GetBytes(16);  // AES-CBC IV (16 bytes)
        return (key, iv);
    }
    
    public static (string EncryptedKey, string EncryptedIV) EncryptSessionKey(
        string rsaPublicKeyPem, byte[] key, byte[] iv)
    {
        using var rsa = RSA.Create();
        rsa.ImportFromPem(rsaPublicKeyPem);
        
        var encryptedKey = rsa.Encrypt(key, RSAEncryptionPadding.OaepSHA256);
        var encryptedIv = rsa.Encrypt(iv, RSAEncryptionPadding.OaepSHA256);
        
        return (Convert.ToBase64String(encryptedKey), Convert.ToBase64String(encryptedIv));
    }
}