using System;

using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Math;

namespace Org.BouncyCastle.Crypto.Engines
{
    /**
    * support class for constructing intergrated encryption ciphers
    * for doing basic message exchanges on top of key agreement ciphers
    */
    public class IesEngine
    {
        private readonly IBasicAgreement     agree;
        private readonly IDerivationFunction kdf;
        private readonly IMac                mac;
        private readonly BufferedBlockCipher cipher;
		private readonly byte[]              macBuf;

		private bool				forEncryption;
        private ICipherParameters	privParam, pubParam;
        private IesParameters		param;

        /**
        * set up for use with stream mode, where the key derivation function
        * is used to provide a stream of bytes to xor with the message.
        *
        * @param agree the key agreement used as the basis for the encryption
        * @param kdf the key derivation function used for byte generation
        * @param mac the message authentication code generator for the message
        */
        public IesEngine(
            IBasicAgreement     agree,
            IDerivationFunction kdf,
            IMac                mac)
        {
            this.agree = agree;
            this.kdf = kdf;
            this.mac = mac;
            this.macBuf = new byte[mac.GetMacSize()];
//            this.cipher = null;
        }

        /**
        * set up for use in conjunction with a block cipher to handle the
        * message.
        *
        * @param agree the key agreement used as the basis for the encryption
        * @param kdf the key derivation function used for byte generation
        * @param mac the message authentication code generator for the message
        * @param cipher the cipher to used for encrypting the message
        */
        public IesEngine(
            IBasicAgreement     agree,
            IDerivationFunction kdf,
            IMac                mac,
            BufferedBlockCipher cipher)
        {
            this.agree = agree;
            this.kdf = kdf;
            this.mac = mac;
            this.macBuf = new byte[mac.GetMacSize()];
            this.cipher = cipher;
        }

        /**
        * Initialise the encryptor.
        *
        * @param forEncryption whether or not this is encryption/decryption.
        * @param privParam our private key parameters
        * @param pubParam the recipient's/sender's public key parameters
        * @param param encoding and derivation parameters.
        */
        public void Init(
            bool                     forEncryption,
            ICipherParameters            privParameters,
            ICipherParameters            pubParameters,
            ICipherParameters            iesParameters)
        {
            this.forEncryption = forEncryption;
            this.privParam = privParameters;
            this.pubParam = pubParameters;
            this.param = (IesParameters)iesParameters;
        }

        private byte[] DecryptBlock(
            byte[]  in_enc,
            int     inOff,
            int     inLen,
            byte[]  z)
        {
            byte[]          M = null;
            KeyParameter    macKey = null;
            KdfParameters   kParam = new KdfParameters(z, param.GetDerivationV());
            int             macKeySize = param.MacKeySize;

            kdf.Init(kParam);

            inLen -= mac.GetMacSize();

            if (cipher == null)     // stream mode
            {
				byte[] Buffer = GenerateKdfBytes(kParam, inLen + (macKeySize / 8));

                M = new byte[inLen];

                for (int i = 0; i != inLen; i++)
                {
                    M[i] = (byte)(in_enc[inOff + i] ^ Buffer[i]);
                }

                macKey = new KeyParameter(Buffer, inLen, (macKeySize / 8));
            }
            else
            {
                int cipherKeySize = ((IesWithCipherParameters)param).CipherKeySize;
				byte[] Buffer = GenerateKdfBytes(kParam, (cipherKeySize / 8) + (macKeySize / 8));

                cipher.Init(false, new KeyParameter(Buffer, 0, (cipherKeySize / 8)));

				M = cipher.DoFinal(in_enc, inOff, inLen);

				macKey = new KeyParameter(Buffer, (cipherKeySize / 8), (macKeySize / 8));
            }

            byte[] macIV = param.GetEncodingV();

            mac.Init(macKey);
            mac.BlockUpdate(in_enc, inOff, inLen);
            mac.BlockUpdate(macIV, 0, macIV.Length);
            mac.DoFinal(macBuf, 0);

			inOff += inLen;

			for (int t = 0; t < macBuf.Length; t++)
            {
                if (macBuf[t] != in_enc[inOff + t])
                {
                    throw (new InvalidCipherTextException("IMac codes failed to equal."));
                }
            }

            return M;
        }

        private byte[] EncryptBlock(
            byte[]  input,
            int     inOff,
            int     inLen,
            byte[]  z)
        {
            byte[]          C = null;
            KeyParameter    macKey = null;
            KdfParameters   kParam = new KdfParameters(z, param.GetDerivationV());
            int             c_text_length = 0;
            int             macKeySize = param.MacKeySize;

            if (cipher == null)     // stream mode
            {
				byte[] Buffer = GenerateKdfBytes(kParam, inLen + (macKeySize / 8));

                C = new byte[inLen + mac.GetMacSize()];
                c_text_length = inLen;

				for (int i = 0; i != inLen; i++)
                {
                    C[i] = (byte)(input[inOff + i] ^ Buffer[i]);
                }

                macKey = new KeyParameter(Buffer, inLen, (macKeySize / 8));
            }
            else
            {
                int cipherKeySize = ((IesWithCipherParameters)param).CipherKeySize;
				byte[] Buffer = GenerateKdfBytes(kParam, (cipherKeySize / 8) + (macKeySize / 8));

                cipher.Init(true, new KeyParameter(Buffer, 0, (cipherKeySize / 8)));

                c_text_length = cipher.GetOutputSize(inLen);
				byte[] tmp = new byte[c_text_length];

				int len = cipher.ProcessBytes(input, inOff, inLen, tmp, 0);
				len += cipher.DoFinal(tmp, len);

				C = new byte[len + mac.GetMacSize()];
				c_text_length = len;

				Array.Copy(tmp, 0, C, 0, len);

				macKey = new KeyParameter(Buffer, (cipherKeySize / 8), (macKeySize / 8));
            }

            byte[] macIV = param.GetEncodingV();

            mac.Init(macKey);
            mac.BlockUpdate(C, 0, c_text_length);
            mac.BlockUpdate(macIV, 0, macIV.Length);
            //
            // return the message and it's MAC
            //
            mac.DoFinal(C, c_text_length);
            return C;
        }

		private byte[] GenerateKdfBytes(
			KdfParameters	kParam,
			int				length)
		{
			byte[] buf = new byte[length];

			kdf.Init(kParam);

			kdf.GenerateBytes(buf, 0, buf.Length);

			return buf;
		}

		public byte[] ProcessBlock(
            byte[]  input,
            int     inOff,
            int     inLen)
        {
            agree.Init(privParam);

			BigInteger z = agree.CalculateAgreement(pubParam);

			// TODO Is a fixed length result expected?
			byte[] zBytes = z.ToByteArrayUnsigned();

            return forEncryption
				?	EncryptBlock(input, inOff, inLen, zBytes)
                :	DecryptBlock(input, inOff, inLen, zBytes);
        }
    }

}