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

namespace GeoSus.Client
{
    // Klientská strana šifrování - generuje session key, šifruje RSA, AES-CBC session
    // Používá AES-CBC místo AES-GCM pro kompatibilitu s Unity
    public class ClientEncryption : IDisposable
    {
        private byte[] _sessionKey;
        private byte[] _sessionIv;
        private long _nonceCounter;
        private readonly object _lock = new object();
        
        // Kontrola, zda je session key nastaven
        public bool HasSessionKey => _sessionKey != null && _sessionIv != null;
        
        // Generuje nový session key a IV
        public void GenerateSessionKey()
        {
            _sessionKey = new byte[32]; // AES-256
            _sessionIv = new byte[16];  // CBC IV (16 bytes)
            
            using (var rng = RandomNumberGenerator.Create())
            {
                rng.GetBytes(_sessionKey);
                rng.GetBytes(_sessionIv);
            }
        }
        
        public byte[] SessionKey => _sessionKey ?? throw new InvalidOperationException("Session key not generated");
        public byte[] SessionIV => _sessionIv ?? throw new InvalidOperationException("Session IV not generated");
        
        // Zašifruje session key pomocí RSA public key serveru
        public (string EncryptedKey, string EncryptedIV) EncryptSessionKeyForServer(string rsaPublicKeyPem)
        {
            if (_sessionKey == null || _sessionIv == null)
                throw new InvalidOperationException("Session key not generated");
                
            using (var rsa = RSA.Create())
            {
                // Parse PEM - extrahuj Base64 obsah
                var pemLines = rsaPublicKeyPem.Split('\n');
                var base64 = new StringBuilder();
                foreach (var line in pemLines)
                {
                    var trimmed = line.Trim();
                    if (!trimmed.StartsWith("-----") && !string.IsNullOrEmpty(trimmed))
                    {
                        base64.Append(trimmed);
                    }
                }
                
                var keyBytes = Convert.FromBase64String(base64.ToString());
                
                // Unity kompatibilní import - parsujeme SubjectPublicKeyInfo ručně
                ImportSubjectPublicKeyInfoManual(rsa, keyBytes);
                
                // Používáme OaepSHA1 pro Unity kompatibilitu (OaepSHA256 není podporován)
                var encryptedKey = rsa.Encrypt(_sessionKey, RSAEncryptionPadding.OaepSHA1);
                var encryptedIv = rsa.Encrypt(_sessionIv, RSAEncryptionPadding.OaepSHA1);
                
                return (Convert.ToBase64String(encryptedKey), Convert.ToBase64String(encryptedIv));
            }
        }
        
        // Ručně parsuje SubjectPublicKeyInfo (DER) a importuje RSA klíč - Unity kompatibilní
        private static void ImportSubjectPublicKeyInfoManual(RSA rsa, byte[] subjectPublicKeyInfo)
        {
            // SubjectPublicKeyInfo ::= SEQUENCE {
            //   algorithm AlgorithmIdentifier,
            //   subjectPublicKey BIT STRING }
            // RSAPublicKey ::= SEQUENCE { modulus INTEGER, publicExponent INTEGER }
            
            int index = 0;
            
            // Outer SEQUENCE
            if (subjectPublicKeyInfo[index++] != 0x30)
                throw new InvalidOperationException("Invalid SubjectPublicKeyInfo");
            ReadLength(subjectPublicKeyInfo, ref index);
            
            // AlgorithmIdentifier SEQUENCE - skip it
            if (subjectPublicKeyInfo[index++] != 0x30)
                throw new InvalidOperationException("Invalid AlgorithmIdentifier");
            int algLen = ReadLength(subjectPublicKeyInfo, ref index);
            index += algLen;
            
            // BIT STRING containing RSAPublicKey
            if (subjectPublicKeyInfo[index++] != 0x03)
                throw new InvalidOperationException("Invalid BIT STRING");
            ReadLength(subjectPublicKeyInfo, ref index);
            index++; // Skip unused bits byte (should be 0)
            
            // RSAPublicKey SEQUENCE
            if (subjectPublicKeyInfo[index++] != 0x30)
                throw new InvalidOperationException("Invalid RSAPublicKey");
            ReadLength(subjectPublicKeyInfo, ref index);
            
            // Modulus INTEGER
            byte[] modulus = ReadInteger(subjectPublicKeyInfo, ref index);
            
            // Exponent INTEGER
            byte[] exponent = ReadInteger(subjectPublicKeyInfo, ref index);
            
            var parameters = new RSAParameters
            {
                Modulus = modulus,
                Exponent = exponent
            };
            rsa.ImportParameters(parameters);
        }
        
        private static int ReadLength(byte[] data, ref int index)
        {
            int length = data[index++];
            if ((length & 0x80) != 0)
            {
                int numBytes = length & 0x7F;
                length = 0;
                for (int i = 0; i < numBytes; i++)
                {
                    length = (length << 8) | data[index++];
                }
            }
            return length;
        }
        
        private static byte[] ReadInteger(byte[] data, ref int index)
        {
            if (data[index++] != 0x02)
                throw new InvalidOperationException("Expected INTEGER");
            int length = ReadLength(data, ref index);
            
            // Skip leading zero if present (used for positive sign in DER)
            int originalLength = length;
            int start = index;
            if (length > 1 && data[start] == 0x00)
            {
                start++;
                length--;
            }
            
            byte[] result = new byte[length];
            Buffer.BlockCopy(data, start, result, 0, length);
            index += originalLength;
                
            return result;
        }
        
        // Šifruje zprávu pomocí AES-256-CBC s HMAC
        public byte[] Encrypt(byte[] plaintext)
        {
            if (_sessionKey == null || _sessionIv == null)
                throw new InvalidOperationException("Session key not set");
                
            lock (_lock)
            {
                // Generuj unikátní IV pro tuto zprávu
                var iv = GetNextIV();
                
                using (var aes = Aes.Create())
                {
                    aes.Key = _sessionKey;
                    aes.IV = iv;
                    aes.Mode = CipherMode.CBC;
                    aes.Padding = PaddingMode.PKCS7;
                    
                    byte[] ciphertext;
                    using (var encryptor = aes.CreateEncryptor())
                    using (var ms = new MemoryStream())
                    {
                        using (var cs = new CryptoStream(ms, encryptor, CryptoStreamMode.Write))
                        {
                            cs.Write(plaintext, 0, plaintext.Length);
                        }
                        ciphertext = ms.ToArray();
                    }
                    
                    // Compute HMAC pro integritu
                    byte[] hmac;
                    using (var hmacSha = new HMACSHA256(_sessionKey))
                    {
                        var toSign = new byte[iv.Length + ciphertext.Length];
                        Buffer.BlockCopy(iv, 0, toSign, 0, iv.Length);
                        Buffer.BlockCopy(ciphertext, 0, toSign, iv.Length, ciphertext.Length);
                        hmac = hmacSha.ComputeHash(toSign);
                    }
                    
                    // Výstup: [16 bytes IV][32 bytes HMAC][ciphertext]
                    var result = new byte[16 + 32 + ciphertext.Length];
                    Buffer.BlockCopy(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 pomocí AES-256-CBC s HMAC ověřením
        public byte[] Decrypt(byte[] encrypted)
        {
            if (_sessionKey == null)
                throw new InvalidOperationException("Session key not set");
                
            if (encrypted.Length < 48) return null;
            
            try
            {
                var iv = new byte[16];
                var hmac = new byte[32];
                var ciphertext = new byte[encrypted.Length - 48];
                
                Buffer.BlockCopy(encrypted, 0, iv, 0, 16);
                Buffer.BlockCopy(encrypted, 16, hmac, 0, 32);
                Buffer.BlockCopy(encrypted, 48, ciphertext, 0, ciphertext.Length);
                
                // Ověř HMAC
                byte[] expectedHmac;
                using (var hmacSha = new HMACSHA256(_sessionKey))
                {
                    var toVerify = new byte[iv.Length + ciphertext.Length];
                    Buffer.BlockCopy(iv, 0, toVerify, 0, iv.Length);
                    Buffer.BlockCopy(ciphertext, 0, toVerify, iv.Length, ciphertext.Length);
                    expectedHmac = hmacSha.ComputeHash(toVerify);
                }
                
                // Constant-time compare
                var diff = 0;
                for (int i = 0; i < 32; i++)
                {
                    diff |= hmac[i] ^ expectedHmac[i];
                }
                if (diff != 0) return null; // HMAC mismatch
                
                using (var aes = Aes.Create())
                {
                    aes.Key = _sessionKey;
                    aes.IV = iv;
                    aes.Mode = CipherMode.CBC;
                    aes.Padding = PaddingMode.PKCS7;
                    
                    using (var decryptor = aes.CreateDecryptor())
                    using (var ms = new MemoryStream(ciphertext))
                    using (var cs = new CryptoStream(ms, decryptor, CryptoStreamMode.Read))
                    using (var output = new MemoryStream())
                    {
                        cs.CopyTo(output);
                        return output.ToArray();
                    }
                }
            }
            catch (CryptographicException)
            {
                return null;
            }
        }
        
        private byte[] GetNextIV()
        {
            if (_sessionIv == null)
                throw new InvalidOperationException("Session IV not set");
                
            var iv = new byte[16];
            Buffer.BlockCopy(_sessionIv, 0, iv, 0, 8);
            
            var counter = System.Threading.Interlocked.Increment(ref _nonceCounter);
            var counterBytes = BitConverter.GetBytes(counter);
            Buffer.BlockCopy(counterBytes, 0, iv, 8, 8);
            
            return iv;
        }
        
        public void Dispose()
        {
            if (_sessionKey != null)
            {
                Array.Clear(_sessionKey, 0, _sessionKey.Length);
                _sessionKey = null;
            }
        }
    }
}
