using System.Security.Cryptography;

namespace KyberDotnet
{
	public class Kyber
	{
		public const int q = 3329;
		public const int n = 256;

		public int k = 2;
		public int eta1 = 3, eta2 = 2;
		public int DU = 10, DV = 4;


		protected Kyber(int k, int eta1, int eta2, int DU, int DV)
		{
			this.k = k;
			this.eta1 = eta1;
			this.eta2 = eta2;
			this.DU = DU;
			this.DV = DV;
		}

		public static Kyber Create512() => new(k: 2, eta1: 3, eta2: 2, DU: 10, DV: 4);
		public static Kyber Create768() => new(k: 3, eta1: 2, eta2: 2, DU: 10, DV: 4);
		public static Kyber Create1024() => new(k: 4, eta1: 2, eta2: 2, DU: 11, DV: 5);

		static byte[] GetRandom32()
		{
			return RandomNumberGenerator.GetBytes(32);
		}

		//	--------------------------------------------------------------------- CPAPKE

		/// <summary>
		/// Random public and private key generation.
		/// </summary>
		/// <returns>Public and secret key tuple.</returns>
		(byte[] pk, byte[] sk) CPAPKE_KeyGen()
		{
			var d = GetRandom32();

			var rs = Hashes.G(d);
			var rho = rs.Take(32).ToArray();
			var sigma = rs.TakeLast(32).ToArray();

			var Ahat = new Poly[k, k];
			for (byte i = 0; i < k; i++)
			{
				for (byte j = 0; j < k; j++)
				{
					Ahat[i, j] = Utils.Parse(Hashes.XOF(rho, j, i));	//j,i !
				}
			}

			var s = new Poly[k];
			var e = new Poly[k];

			byte N = 0;

			for (int i = 0; i < k; i++)
			{
				s[i] = Utils.CBD(Hashes.PRF(sigma, N++), eta1);
			}
			for (int i = 0; i < k; i++)
			{
				e[i] = Utils.CBD(Hashes.PRF(sigma, N++), eta1);
			}

			//s -> sHat
			foreach (var poly in s) NTT.Generate(poly);
			//e -> eHat
			foreach (var poly in e) NTT.Generate(poly);

			//tHat = Ahat * sHat + eHat
			var tHat = NTT.MatrixMultiply(Ahat, s);
			Poly.Add(tHat, e);
			Poly.Reduce(tHat);

			//pk = Enc...||rho
			var pk = new List<byte>(Utils.Encode(tHat, 12));
			pk.AddRange(rho);

			var sk = Utils.Encode(s, 12);

			return (pk.ToArray(), sk);
		}

		/// <summary>
		/// Message encryption using a public key.
		/// </summary>
		/// <param name="pk">Public key byte array.</param>
		/// <param name="m">Message to encode of length 32 bytes.</param>
		/// <param name="rand">Random 32 byte array.</param>
		/// <returns>Encrypted message byte array.</returns>
		byte[] CPAPKE_Encrypt(byte[] pk, byte[] m, byte[] rand)
		{
			var tHat = Utils.DecodeVector(pk, 12, k);

			//rho from end of pk
			var rho = pk.TakeLast(32).ToArray();

			var AhatT = new Poly[k, k];
			for (byte i = 0; i < k; i++)
			{
				for (byte j = 0; j < k; j++)
				{
					AhatT[i, j] = Utils.Parse(Hashes.XOF(rho, i, j));	//i,j !
				}
			}

			var r = new Poly[k];
			var e1 = new Poly[k];

			byte N = 0;

			for (int i = 0; i < k; i++)
			{
				r[i] = Utils.CBD(Hashes.PRF(rand, N++), eta1);
			}
			for (int i = 0; i < k; i++)
			{
				e1[i] = Utils.CBD(Hashes.PRF(rand, N++), eta2);
			}
			var e2 = Utils.CBD(Hashes.PRF(rand, N), eta2);

			//r -> rHat
			foreach (var poly in r) NTT.Generate(poly);

			//u = INTT(AhatT * rHat) + e1
			var u = NTT.MatrixMultiply(AhatT, r);
			foreach (var poly in u) NTT.Invert(poly);
			Poly.Add(u, e1);
			Poly.Reduce(u);

			//v = INTT(tHatT * rHat) + e2 + Decompress(Decode(m))
			var v = NTT.DotProduct(tHat, r);
			NTT.Invert(v);
			v.Add(e2);

			var vAdd = Utils.Decode(m, 1);
			Utils.Decompress(vAdd, 1);

			v.Add(vAdd);
			v.Reduce();

			//u -> Compress(u)
			Utils.Compress(u, DU);
			//c1 = Encode(Compress(u))
			var c1 = new List<byte>(Utils.Encode(u, DU));

			//v -> Compress(v)
			Utils.Compress(v, DV);
			//c = c1 || Encode(Compress(v))
			c1.AddRange(Utils.Encode(v, DV));

			return c1.ToArray();
		}

		/// <summary>
		/// Message decryption.
		/// </summary>
		/// <param name="sk">Secret key byte array.</param>
		/// <param name="c">Encrypted message.</param>
		/// <returns>Decrypted 32 byte message.</returns>
		byte[] CPAPKE_Decrypt(byte[] sk, byte[] c)
		{
			//c1 -> u
			var u = Utils.DecodeVector(c, DU, k);
			Utils.Decompress(u, DU);

			//c2 -> v
			var v = Utils.Decode(c.Skip(DU * k * (n / 8)).ToArray(), DV);
			Utils.Decompress(v, DV);

			//sk -> sHat
			var sHat = Utils.DecodeVector(sk, 12, k);

			//u -> NTT(u)
			foreach (var poly in u) NTT.Generate(poly);

			//INTT(sHatT * NTT(U))
			var m = NTT.DotProduct(sHat, u);
			NTT.Invert(m);

			//v = v - INTT(...)
			v.Sub(m);
			v.Reduce();

			//m = Encode(Compress(v - ...))
			Utils.Compress(v, 1);
			return Utils.Encode(v, 1);
		}

		//	--------------------------------------------------------------------- CCAKEM

		/// <summary>
		/// KEM Key generation.
		/// </summary>
		/// <returns>Public and secret key tuple.</returns>
		public (byte[] pk, byte[] sk) CCAKEM_KeyGen()
		{
			var z = GetRandom32();
			var keys = CPAPKE_KeyGen();

			//sk'||pk||H(pk)||z
			var sk = new List<byte>();
			sk.AddRange(keys.sk);
			sk.AddRange(keys.pk);
			sk.AddRange(Hashes.H(keys.pk));
			sk.AddRange(z);
			return (keys.pk, sk.ToArray());
		}

		/// <summary>
		/// KEM Encapsulation.
		/// </summary>
		/// <param name="pk">Public key byte array.</param>
		/// <returns>Ciphertext and shared secret tuple.</returns>
		public (byte[] c, byte[] K) CCAKEM_Enc(byte[] pk)
		{
			//m = H(B32)
			var m = Hashes.H(GetRandom32());

			//(K', r) = G(m||H(pk))
			var g = new List<byte>(m);
			g.AddRange(Hashes.H(pk));
			var kr = Hashes.G(g.ToArray());

			var K = kr.Take(32).ToList();
			var r = kr.TakeLast(32).ToArray();

			//c = Enc(pk,m,r)
			var c = CPAPKE_Encrypt(pk, m, r);

			//K = KDF(K'||H(c))
			K.AddRange(Hashes.H(c));
			return (c, Hashes.KDF(K.ToArray()));
		}

		/// <summary>
		/// KEM Decapsulation.
		/// </summary>
		/// <param name="c">Ciphertext.</param>
		/// <param name="sk">Secret key.</param>
		/// <returns>Shared secret.</returns>
		public byte[] CCAKEM_Dec(byte[] c, byte[] sk)
		{
			var pk = sk.Skip(12 * k * n / 8).Take(12 * k * n / 8 + 32).ToArray();
			var h = sk.Skip(24 * k * n / 8 + 32).Take(32).ToArray();
			var z = sk.Skip(24 * k * n / 8 + 64).Take(32).ToArray();

			//m' = Dec(pk,m',r')
			var mD = CPAPKE_Decrypt(sk, c);

			//m'||h
			var mh = new List<byte>(mD);
			mh.AddRange(h);

			//(K',r') = G(m'||h)
			var kr = Hashes.G(mh.ToArray());
			var KD = kr.Take(32).ToArray();
			var rD = kr.TakeLast(32).ToArray();

			//c' = Enc(pk,m',r')
			var cD = CPAPKE_Encrypt(pk, mD, rD);

			//c == c'
			if (Enumerable.SequenceEqual(c, cD))
			{
				//KDF(K'||H(c))
				var kc = new List<byte>(KD);
				kc.AddRange(Hashes.H(c));
				return Hashes.KDF(kc.ToArray());
			}
			else
			{
				//return KDF(z||H(c))
				var zc = new List<byte>(z);
				zc.AddRange(Hashes.H(c));
				return Hashes.KDF(zc.ToArray());
			}
		}
	}
}
