using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Text;

namespace SuperCryptEX
{
    public static class Class_SHA3
    {
        #region Parameters

        // Parameters
        private static readonly ushort[] block_size_array = { 224, 256, 384, 512 };         // Block Size                       (224, 256, 384, 512)-Bit
        private static readonly ushort[] block_size_array_in_bytes = { 28, 32, 48, 64 };
        private static readonly ushort[] capacity_array = { 448, 512, 768, 1024 };          // Capacity (c)                     (c = Block Size * 2)
        private static readonly ushort[] capacity_array_in_bytes = { 56, 64, 96, 128 };
        private static readonly ushort[] rate_array = { 1152, 1088, 832, 576 };             // Rate (r)                         (r = b - c)
        private static readonly ushort[] rate_array_in_bytes = { 144, 136, 104, 72 };

        private const byte matrix_size = 5 * 5;
        private const byte state_length = 6;                                                // Panjang State (l)                (1, 2, 3, 4, 5, 6)  -> Default : 6
        private const byte state_width = 64;                                                // Ukuran State (w)                 (w = 2 ^ l)         -> Default : 64-Bit
        private const byte state_width_in_bytes = 8;
        private const byte rounds_number = 24;                                              // Jumlah Round untuk Keccak-f (Nr) (Nr = 12 + 2 * l)    -> Default : 24 Rounds
        private const ushort permutation_width = 1600;                                      // Ukuran Permutasi (b)             (b = 25 * w)        -> Default: 1600-Bit

        /* Untuk SHA3-512 :
         * Length            = 6 (Dari 1, 2, 3, 4, 5, 6)
         * State Width (w)   = 2 ^ l         = 2 ^ 6         = 64    = [ 8 Bytes ]
         * Block (b)         = 5 x 5 x w     = 5 x 5 x 6     = 1600  = [ 200 Bytes ]
         * Rate (r)          = 576           = [ 72 Bytes ]
         * Capacity (c)      = Block - Rate  = 1600 - 576    = 1024  = [ 128 Bytes ]
         * Round Number (Nr) = 12 + 2 * l    = 12 + 2 * 6    = 24 Rounds
         */

        // Variabel Konstanta - Rotation Offset
        private static readonly byte[,] RotationOffset = {  // ((t + 1) * (t + 2) / 2) % (state_width = 64)
        {0, 36, 3, 41, 18},
        {1, 44, 10, 45, 2},
        {62, 6, 43, 15, 61},
        {28, 55, 25, 21, 56},
        {27, 20, 39, 8, 14} };

        // Variabe; Konstanta - Rotation Constants untuk Panjang State (l = 6), Jumlah Round (Nr = 24)
        private static readonly ulong[] RC = {
        0x0000000000000001,
        0x0000000000008082,
        0x800000000000808A,
        0x8000000080008000,
        0x000000000000808B,
        0x0000000080000001,
        0x8000000080008081,
        0x8000000000008009,
        0x000000000000008A,
        0x0000000000000088,
        0x0000000080008009,
        0x000000008000000A,
        0x000000008000808B,
        0x800000000000008B,
        0x8000000000008089,
        0x8000000000008003,
        0x8000000000008002,
        0x8000000000000080,
        0x000000000000800A,
        0x800000008000000A,
        0x8000000080008081,
        0x8000000000008080,
        0x0000000080000001,
        0x8000000080008008 };

        // Variables
        private static readonly ulong[,] B = new ulong[5, 5];
        private static readonly ulong[] C = new ulong[5];
        private static readonly ulong[] D = new ulong[5];
        #endregion

        #region Constructor

        // SHA3 Constructor
        public static HEX SHA3(FileStream MessageFS, BlockSize blockSize)
        {
            return Keccak(GetBlockSizeIndex(blockSize), MessageFS).ToHex();
        }
        #endregion

        #region Functions

        // Core Function - Keccak
        public static byte[] Keccak(int blockSizeIndex /* Index : 0 = (224), 1 = (256), 2 = (384), 3 = (512)*/, FileStream MessageFS)
        {
            ushort blockSize = block_size_array[blockSizeIndex];
            ushort rate_in_bytes = rate_array_in_bytes[blockSizeIndex];

            // Inisialisasi Sebuah Byte State Ukuran 5 x 5 dengan NullByte (Byte 0)
            ulong[,] State = new ulong[5, 5];
            for (int i = 0; i < 5; ++i)
            {
                for (int j = 0; j < 5; ++j)
                {
                    State[i, j] = 0;
                }
            }

            // Inisialisasi Sebuah Byte untuk Menampung File Sebanyak Rate
            byte[] fileToRead = new byte[rate_in_bytes];
            int hasRead;

            // [#] Absorbing Phase :
            // Baca File Sebanyak Kelipatan Rate, Lakukan State Change pada Setiap Blocknya
            while ((hasRead = MessageFS.Read(fileToRead, 0, rate_in_bytes)) == rate_in_bytes)
            {
                State_Change(State, fileToRead, rate_in_bytes);
            }

            // Saat End of File, Jika Tidak Sampai Rate, Lakukan Padding
            Pad_Last_Block(State, fileToRead.SubArray(0, hasRead), rate_in_bytes);

            // [#] Squeezing Phase :
            // Squeezing Sebanyak Rate, Kemudian Ambil 512-bit / 8 = 64 Bytes Message Digest 
            return Squeezing(State, rate_in_bytes).Take(blockSize / 8).ToArray();
        }

        // Core Function - State Change (Dilakukan Sebanyak Rate)
        private static ulong[,] State_Change(ulong[,] State, byte[] Message, ushort Rate_in_bytes)
        {
            int position;
            for (byte iBlock = 0; iBlock < 5; ++iBlock)
                for (byte jBlock = 0; jBlock < 5; ++jBlock)
                {
                    // Position = 0, 5, 10, 15, 20 | 1, 6, 11, 16, 21 | 2, 7, 12, 17, 22 | 3, 8, 13, 18, 23 | 4, 9, 14, 19, 24
                    position = iBlock + (jBlock * 5);

                    // Secara Default (pos < 9), Panjang Setiap Sub-String Adalah 72 Bytes (9 Words of 4 Bytes = 9 x 8 = 72)
                    // Rate = 72 Bytes, State Width = 8 Bytes
                    if (position < (Rate_in_bytes / state_width_in_bytes)) // If (pos < 9)
                    {
                        position *= state_width_in_bytes;
                        for (byte i = 0; i < state_width_in_bytes; ++i)
                        {
                            /* Simbol : 
                             * "<<" = Shift Left
                             * "^"  = XOR
                             */
                            State[iBlock, jBlock] ^= (ulong)Message[position + i] << (i * 8);
                        }
                    }
                    else break;
                }

            // Fungsi Permutasi Keccak_f
            for (byte i = 0; i < rounds_number; i++) // i < 24
            {
                State = Keccak_f(State, RC[i]);
            }

            return State;
        }

        // Core Function - Fungsi Permutasi Keccak-f
        private static ulong[,] Keccak_f(ulong[,] A, ulong RC_i)
        {
            Iota(Chi(Rho_Pi(Theta(A))), RC_i);
            return A;
        }

        // Core Function - Theta Permutation (Fungsi Inti pada Keccak-f)
        private static ulong[,] Theta(ulong[,] A)
        {
            // Rumus : C[i] = A[i, 0] XOR A[i, 1] XOR A[i, 2] XOR A[i, 3] XOR A[i, 4]
            for (int i = 0; i < 5; i++)
            {
                C[i] = A[i, 0] ^ A[i, 1] ^ A[i, 2] ^ A[i, 3] ^ A[i, 4];
            }

            // Rumus : D[i] = C[(i - 1) mod 5, 0] XOR ROT(C[(i + 1) mod 5, 1])
            for (int i = 0; i < 5; i++)
            {
                D[i] = C[(i + 4) % 5] ^ Rotate(C[(i + 1) % 5], 1, state_width);
            }

            // Rumus : A'[i, j] = A[i, j] XOR D[i]
            for (int i = 0; i < 5; i++)
            {
                for (int j = 0; j < 5; j++)
                {
                    A[i, j] = A[i, j] ^ D[i];
                }
            }

            return A;
        }

        // Core Function - Rho & Pi Permutation (Fungsi Inti pada Keccak-f)
        private static ulong[,] Rho_Pi(ulong[,] A)
        {
            // Rumus : B[j, (2i + 3j) mod 5] = ROT(A[i, j], Offset[i, j])
            for (int i  = 0; i < 5; i++)
            {
                for (int j = 0; j < 5; j++)
                {
                    // (Pi)                     (Rho)
                    B[j, (2 * i + 3 * j) % 5] = Rotate(A[i, j], RotationOffset[i, j], state_width);
                }
            }

            return A;
        }

        // Core Function - Chi Permutation (Fungsi Inti pada Keccak-f)
        private static ulong[,] Chi(ulong[,] A)
        {
            // Rumus : A[i, j] = B[i, j] XOR ((inverse_B[(i + 1) mod 5, j]) AND B[(i + 2) mod 5, j])
            for (int i = 0; i < 5; i++)
            {
                for (int j = 0; j < 5; j++)
                {
                    A[i, j] = B[i, j] ^ ((~B[(i + 1) % 5, j]) & B[(i + 2) % 5, j]);
                }
            }

            return A;
        }

        // Core Function - Iota Permutation (Fungsi Inti pada Keccak-f)
        private static ulong[,] Iota(ulong[,] A, ulong RC_i)
        {
            // Rumus : A[0, 0] = A[0, 0] XOR RC[i]
            A[0, 0] = A[0, 0] ^ RC_i;

            return A;
        }

        // Core Function - Fungsi untuk Melakukan Fase Squeezing
        private static byte[] Squeezing(ulong[,] State, ushort Rate_in_bytes)
        {
            byte[] hash = new byte[Rate_in_bytes];
            byte[] temp = new byte[8];
            int position, current = 0;

            for (int jBlock = 0; jBlock < 5; ++jBlock)
            {
                for (int iBlock = 0; iBlock < 5; ++iBlock)
                {
                    position = iBlock + jBlock * 5;
                    if (position < (Rate_in_bytes / state_width_in_bytes))
                    {
                        temp = BitConverter.GetBytes(State[iBlock, jBlock]);
                        for (int j = 0; j < 8; ++j)
                        {
                            hash[current + j] = temp[j];
                        }
                        current += 8;
                    }
                }
            }

            return hash;
        }

        // Function - Padding Block Input Terakhir Sebanyak Rate
        private static ulong[,] Pad_Last_Block(ulong[,] State, byte[] MessageB, ushort Rate_in_bytes)
        {
            // Delta = Selisih dari (Buffer Terakhir hingga End of File yang Terbaca) dengan Rate
            ushort delta = (ushort)(Rate_in_bytes - MessageB.Length);
            byte[] Pad_Message = new byte[Rate_in_bytes];
            int pos = 0;

            for (int i = 0; i < MessageB.Length; ++i, ++pos)
            {
                Pad_Message[pos] = MessageB[i];
            }

            if (delta == 1)
            {
                Pad_Message[pos] = 0x86; //0x61; //0110 0001
            }
            else
            {
                Pad_Message[pos] = 0x06; //0x60; //0110 0000
                delta -= 2;

                while (delta > 0)
                {
                    Pad_Message[pos + delta] = 0x00;
                    --delta;
                }

                Pad_Message[Pad_Message.Length - 1] = 0x80; //0x01;
            }

            State_Change(State, Pad_Message, Rate_in_bytes);

            return State;
        }

        // Function - Rotasi Byte pada Block State
        private static ulong Rotate(ulong x, byte n, byte w)
        {
            return ((x << (n % w)) | (x >> (w - (n % w))));
        }

        // Function - Get Index untuk Block Size, Capacity, dan Rate
        private static int GetBlockSizeIndex(BlockSize blockSize)
        {
            int index = 0;

            if (blockSize == BlockSize.Bit_224)
            {
                index = 0;
            }
            else if (blockSize == BlockSize.Bit_256)
            {
                index = 1;
            }
            else if (blockSize == BlockSize.Bit_384)
            {
                index = 2;
            }
            else if (blockSize == BlockSize.Bit_512)
            {
                index = 3;
            }

            /*
            int constant_index = Array.FindIndex(security_array, s => s == security);
            if (constant_index == -1) throw new ArgumentException("Security May be One of " + security_array);
            */
            return index;
        }

        #endregion
    }

    // Class Tambahan untuk Mengubah Output Hash Menjadi HEX
    public class HEX
    {
        private readonly string data;

        public static implicit operator string(HEX value)
        {
            return value.data;
        }

        public HEX()
        {
            data = "";
        }

        public HEX(byte[] b)
        {
            data = "";
            for (var i = 0; i < b.Length; ++i)
                data += String.Format("{0:x2}", b[i]);
        }

        public HEX(string str)
        {
            str = str.Trim();
            str = str.ToLower();
            if (str.Length % 2 == 1) str = str + "0";
            ValidateHEX(str);
            data = str;
        }

        private static void ValidateHEX(string value)
        {
            foreach (var el in value)
            {
                if (!IsHex(el))
                {
                    throw new ArgumentException("Not a HEX string.");
                }
            }
        }

        private static bool IsHex(char c)
        {
            if ((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'))
                return true;
            return false;
        }

        public byte[] ToByteArray()
        {
            var bytes = new byte[data.Length / 2];
            for (var i = 0; i < bytes.Length; ++i)
            {
                bytes[i] = byte.Parse(data.Substring(i * 2, 2), NumberStyles.HexNumber);
            }

            return bytes;
        }

        public int Length()
        {
            return data.Length;
        }
    }

    // Class Tambahan untuk Hash Utility
    internal static class Utils
    {
        public static byte[] SubArray(this IReadOnlyList<byte> byteArray, int position, int length)
        {
            var subArray = new byte[length];
            for (var i = 0; i < length; ++i)
            {
                subArray[i] = byteArray[i + position];
            }
            return subArray;
        }

        public static HEX ToHex(this byte[] b)
        {
            return new HEX(b);
        }

        public static byte[] ToByteArray(this string str)
        {
            string temp = str.Replace("\r", "");
            var r = Encoding.UTF8.GetBytes(temp);
            return r;
        }
    }

}
