with Interfaces; use Interfaces;
with SMG_OAEP; use SMG_OAEP;
with SMG_Keccak; use SMG_Keccak;
with Ada.Exceptions; use Ada.Exceptions;
with Ada.Text_IO; use Ada.Text_IO;
with Ada.Strings.Fixed; use Ada.Strings.Fixed;
with System;

procedure SMG_Keccak.Test is
  package Octet_IO is new Ada.Text_IO.Modular_IO(Interfaces.Unsigned_8); 

  --types
  type Keccak_Perms is (None, Theta, Rho, Pi, Chi, Iota);
  type Test_Vector is array(Keccak_Perms) of State;
  type Test_Round is array(Round_Index) of Test_Vector;

  --helper methods
  procedure print_state(S: in State; Title: in String) is
    Hex: array(0..15) of Character := ("0123456789ABCDEF");
    Len: constant Natural := Z_Length / 4;
    HexString: String(1..Len);
    W: ZWord;
  begin
    Put_Line("---------" & Title & "---------");
    for Y in XYCoord loop
      for X in XYCoord loop
        W := S(X,Y);
        for Z in 0..Len-1 loop
          HexString(Natural(Len-Z)) := Hex(Natural(W mod 16));
          W := W / 16;
        end loop;
        Put(HexString & " ");
      end loop;
      Put_Line("");
    end loop;
  end;

  procedure print_bytestream(B: in Bytestream; Title: in String) is
  begin
    Put_Line("---" & Title & "---");
    for Pos in B'First..B'Last loop
      Octet_IO.Put(Item => B(Pos), Width => 2, Base => 16);
    end loop;
    New_Line;
  end print_bytestream;

  function read_state(File: in FILE_TYPE; Oct: Positive :=8) return State is
    S: State;
    Line1: String := "0000000000000000 " &
                     "0000000000000000 " &
                     "0000000000000000 " & 
                     "0000000000000000 " &
                     "0000000000000000";
    StartPos, EndPos: Positive;
    Len: Positive := Oct*2;
  begin
    for Y in XYCoord loop
      Line1 := Get_Line(File);
      StartPos := Line1'First;
      EndPos := StartPos + Len-1;

      for X in XYCoord loop
        S(X,Y) := ZWord'value("16#" & Line1(StartPos..EndPos) & "#");
        StartPos := EndPos + 2;	--one space to skip
        EndPos := StartPos + Len - 1;
      end loop;
    end loop;
    return S;
  end read_state;

  --reads a full test round from specified file (pre-defined format)
  function read_from_file (filename : in String;
                           T        : out Test_Round)
                           return Boolean is 
    file: FILE_TYPE;
    InputMarker: String := "lanes as 64-bit words:";
    octets: Positive := 8;
    RoundNo: Round_Index;
  begin
    -- try to open the input file
    begin
      open(file, In_File, filename);
    exception
      when others =>
        Put_Line(Standard_Error, 
                 "Can not open the file '" & filename & "'. Does it exist?");
        return False;
    end;

  -- find & read input state first
    RoundNo := -1;
    loop
      declare
        Line: String := Get_Line(file);
      begin
        --check if this is test data of any known kind
        if index(Line, InputMarker, 1) > 0 then
          T(0)(None) := read_state(file, octets);
          print_state(T(0)(None), "Read Input State");
        elsif index(Line, "Round ", 1) > 0 then
          RoundNo := RoundNo +1;
        elsif index(Line, "theta", 1) > 0 then
          T(RoundNo)(Theta) := read_state(file, octets);
          if (RoundNo > 0) then
            T(RoundNo)(None) := T(RoundNo-1)(Iota);  -- previous state as input
          end if;
        elsif index(Line, "rho", 1) > 0 then
          T(RoundNo)(Rho) := read_state(file, octets);
        elsif index(Line, "pi", 1) > 0 then
          T(RoundNo)(Pi) := read_state(file, octets);
        elsif index(Line, "chi", 1) > 0 then
          T(RoundNo)(Chi) := read_state(file, octets);
        elsif index(Line, "iota", 1) > 0 then
          T(RoundNo)(Iota) := read_state(file, octets);
        end if;	
        exit when End_Of_File(file);
      end;
    end loop;
    Close(file);
    return True;
  end read_from_file;

  -- performs one single round of Keccak, step by step
  -- each permutation is tested separately
  -- test fails with exception raised at first output not matching expected
  procedure test_one_round(T: Test_Vector; Round: Round_Index) is
    Input: State;
    Expected: State;
    Output: State;
    Test_One_Round_Fail: Exception;
  begin
    Input := T(None);
    for I in Keccak_Perms range Theta..Iota loop
      Expected := T(I);
      case I is
        when Theta => Output := SMG_Keccak.Theta(Input);
        when Rho   => Output := SMG_Keccak.Rho(Input);
        when Pi    => Output := SMG_Keccak.Pi(Input);
        when Chi => Output := SMG_Keccak.Chi(Input);
        when Iota => Output := SMG_Keccak.Iota(RC(Round), Input);
        when others => null;
      end case;

      if (Output /= Expected) then
        print_state(Output, "----------real output-------");
        print_state(Expected, "----------expected output--------");
        raise Test_One_Round_Fail;
      else
        Put_Line("PASSED: " & Keccak_Perms'Image(I));
      end if;
      -- get ready for next permutation
      Input := Expected;
    end loop;
  end test_one_round;

  procedure test_bytes_to_word_conversion is
    bytes: Byteword := (others => 0);
    obtained_bytes: Byteword := (others => 0);
    expected: ZWord;
    msbexpected: ZWord;
    obtained: ZWord;
  begin
    expected := 16#8FA4F19E0287BBE7#;
    msbexpected := 16#F1258F7940E1DDE7#;
    bytes := (231, 187, 135, 2, 158, 241, 164, 143); --LSByte order!

    obtained := BytesToWordLE(bytes);
    obtained_bytes := WordToBytesBE(msbexpected);

    if obtained /= expected then
      Put_Line("FAIL: bytes to word");
      Put_Line("Expected: " & ZWord'Image(expected));
      Put_Line("Obtained: " & ZWord'Image(obtained));
    else
      Put_Line("PASSED: bytes to word");
    end if;

    if obtained_bytes /= bytes then
      Put_Line("FAIL: word to bytes");
      Put("Expected: ");
      for I in Byteword'Range loop
        Put(Interfaces.Unsigned_8'Image(bytes(I)));
      end loop;
      Put_Line("");
      Put_Line("Obtained: ");
      for I in Byteword'Range loop
        Put(Interfaces.Unsigned_8'Image(obtained_bytes(I)));
      end loop;
      Put_Line("");
    else
      Put_Line("PASSED: word to bytes");
    end if;
  end test_bytes_to_word_conversion;

  procedure test_flip is
    B      : Byteword := (231, 187, 135, 2, 158, 241, 164, 143);
    Expected: Byteword := (143, 164, 241, 158, 2, 135, 187, 231);
    Output : Byteword;
  begin
    Output := FlipOctets( B );
    if Output /= Expected then
      Put_Line( "FAILED: flip octets" );
      Put_Line( "Expected: " );
      for I of Expected loop
        Put(Interfaces.Unsigned_8'Image(I));
      end loop;
      new_line(1);
      Put_Line( "Output: " );
      for I of Output loop
        Put(Interfaces.Unsigned_8'Image(I));
      end loop;
      new_line(1);
    else
      Put_Line( "PASSED: flip octets" );
    end if;
  end test_flip;

  procedure test_sponge is
    Bitrate   : constant Keccak_Rate := 1344 / 8;
    Input     : Bytestream(1..1);
    Hex       : array(0..15) of Character := ("0123456789ABCDEF");
    ExpHex    : constant String :=
--for 129 aka 1000 0001
--              "4BAB86C1E1C067A913F62EFD0D65AF58A9268A80A1D8A606F8"&
--              "72F535CDE3F1D66704093C78E6F1A8AFA7D9C6D1C016BA88ED"&
--              "3577450B2B339745ACF467E6DE357A14EC1B1E005E75156FF2"&
--              "B131F224CB2851F9D70D88ACEE2A7C05B9EF0724C1C8653E01"&
--              "A55B0A51145C73F23D6C31E0087D2C414B08939F1E67480867"&
--              "923B3FC7706675678A50260A62DBBFC38A144D75179DFBCA7B"&
--              "DA3CD87C7B8FFD0F8F12149EE6EAF0322BD48A079B3039A62C"&
--              "E1A8A5E9BCB03CE38C61ACF3AA9B9FBB159D7EB212054F7DBA"&
--              "D7ACE8EF0F70B0863E9E6E018F6AE23D74B01707ED59452B73"&
--              "862579B07C3B20DB9E5AA2479E98C0DBE6FD290FCE7F55BDA7"&
--              "8C1ED4A6C6A61E373F0A9D154AFF2DC86673658AF2494AAD09"&
--              "B8409C9B48D217FF354797BFED51807272B3C3AFE52F80FE4A"&
--              "4180ABE14296FF09A5024AED31F310F870288E91CB58569EE1"&
--              "FAA2D558E404D2ADFDD96480AE51CBF7FE6D19E8F0FEBC21DA"&
--              "88D597CA0C2998B95DDA72EF4C749F965688A1E133698E9E71"&
--              "15F3AEC61ADB7CCB279504FA716E7A059564A8DC5E20535A33"&
--              "E781C116B5B72B803E204BB25D91192756A575C1D9513282E1"&
--              "D8204AABE5DAA0A7774A296F2E4B9A87A87F72CD16A00639E7"&
--              "9EF280227975E05CED112346D3AE1DCE1E12EBA9A5673B29C7"&
--              "12AE3546D9CD3E09BCA2F2AAD46FFFCC2418E08604BFE30650"&
--             "549CF10FE6339D769D628828";

--for 19 aka 00010011
              "18F5F5E68207EA4A2F5A992AA4415DDC5DA7EBE12D74CF9E"&
              "0F33205D4E998E09883471A929959B146733C0D1FAD57919"&
              "C6F5B5AA64130266AB5D36666AA896FC69D14E821873138E"&
              "1C53B92FF6106428E5E407AC9E7084DEC1E8819CFB09DD54"&
              "0DE980C631F67BF0CE55F1D6274E5E4274C1D966551B4A6F"&
              "DFD780710872F4A8D4DF380F530EFF443355AAC86CA3AC50"&
              "F65C0339B107B479411B8EFB8860066A2016513E8B3B9E23"&
              "9A5EEF664EA5B50143F9152E45890078D174A75C837633A8"&
              "76E9842316CBDACAFACB87D48E57C6A6948CE71F4A934AB0"&
              "E73E8107F28B3E4CE2CA9A5828F3097126D7141FCF8E6DA4"&
              "FDA017805782BD761663B43AB607914E6FC853443A240E93"&
              "F13C228394C036D0E51CD926D033825951279D4A658F6F81"&
              "5406A8695685C255237CDC6CDD16EC743A42D32BBDDF50A3"&
              "6033E3121B19049AA127138C58411C7B9E7DF4A130C5A3A1"&
              "D7253EC804379A75BF443F54575E96DF71757FEFB8EF4634"&
              "861B2816BAC62B54A88D520D373E7AAAB6F564DEDECC7140"&
              "0B7113BFF9B88AA62EFEC914401FD54B394C54E4B14E7423"&
              "99847A7D1E3EEAC173202B275F7A79E4E6721AA29A9EB974"&
              "D30F72D666A9F7BD7C3FD0EA39239289EB8BA56E3EB4DF94"&
              "D516ADBEE177CE1462379F8A22C82A315D7A4C7BDAE1D0A5"&
              "4A4878C81F56011CDD53D49D2366CCF2D9BBCA0494BD72B4"&
              "17E0C6987118B8F6";
    Output    : Bytestream( 1 .. ExpHex'Length / 2 );
  begin
--    Input(Input'First) := 129; --1000 0001--
    Input(Input'First) := 19; --00010011 MSB--
--    Input(Input'First) := 200; --11001000 MSB--
    Put_Line("---sponge test---");
    Sponge(Input, Output, Bitrate);
    Put_Line("Input is:");
    for I of Input loop
      Put(Interfaces.Unsigned_8'Image(I));
    end loop;
    new_line(1);

    Put_Line("Output is:");
    for I of Output loop
      Put(Interfaces.Unsigned_8'Image(I));
    end loop;
    new_line(1);

    -- check output
    declare
      PO, PE: Natural;
      Pass: Boolean := True;
    begin
      
      for I in 0..Output'Length-1 loop
        PO := Output'First + I;
        PE := ExpHex'First + I*2;
        if Hex(Natural(Output(PO) mod 16)) /= ExpHex(PE+1) or
           Hex(Natural(Output(PO) / 16)) /= ExpHex(PE) then
          Put_Line("FAIL: test_sponge");
          Octet_IO.Put(Output(Output'First+I), Width=>2, Base=>16);
          Put_Line(" vs exp: " &  ExpHex(ExpHex'First+I*2..ExpHex'First+I*2+1));
          Pass := False;
          exit;
        end if;
      end loop;
      if Pass then
        Put_Line("PASS: test_sponge");
      end if;
    end;
  end test_sponge;

  procedure test_keccak_function(T: in Test_Round) is
    S: State;
  begin
    Put_Line("---Full Keccak Function test---");
    S := Keccak_Function(T(Round_Index'First)(None));
    if S /= T(Round_Index'Last)(Iota) then
      Put_Line("FAILED: full keccak function test");
    else
      Put_Line("PASSED: full keccak function test");
    end if;
  end test_keccak_function;

  procedure test_bytestream_conversion is
    S: String := "Aa*/";
    E: Bytestream( 0..3 ) := (65, 97, 42, 47);
    B: Bytestream( 0..3 );
    SS: String := "  t ";
  begin
    Put_Line("---Testing string to bytestream conversion---");
    ToBytestream( S, B );
    if E /= B then
      Put_Line("FAILED: string to bytestream conversion.");
      for I in 0..3 loop
        Put_Line("Got " & Interfaces.Unsigned_8'Image(B(I)) & " vs " &
                 Interfaces.Unsigned_8'Image(E(I)));
      end loop;
    else
      Put_Line("PASSED: string to bytestream conversion.");
    end if;

    Put_Line("---Testing bytestream to string conversion---");
    ToString( B, SS );
    if SS /= S then
      Put_Line("FAILED: bytestream to string conversion");
      Put_Line("EXPECTED: " & S);
      Put_Line("OUTPUT: " & SS);
    else
      Put_Line("PASSED: bytestream to string conversion");
    end if;
  end test_bytestream_conversion;

  procedure test_hash_keccak is
    S: String := "X";
    O: String := "abc";
    B: Bytestream( 0 .. 2 );
    BB: Bytestream( 1.. 3);
    Exp: Bytestream( 1 .. 3 ) := (225, 98, 227);
  begin
    BB(BB'First) := 26;
    Put_Line("----Testing hash keccak on string " & S & "----");
    HashKeccak(S, O);
    ToBytestream( O, B );
    if B /= Exp then
      Put_Line("FAILED: testing hash keccak on string");
      Put_Line("Output:");
      for I of B loop
        Put( Interfaces.Unsigned_8'Image( I ) );
      end loop;
      new_line(1);
      Put_Line("Expected:");
      for I of Exp loop
        Put( Interfaces.Unsigned_8'Image( I ) );
      end loop;
    else
      Put_Line("PASSED: testing hash keccak on string");
    end if;
    new_line(1);
  end test_hash_keccak;

  procedure test_xor_strings is
    S1     : String := "ABC";
    S2     : String := "CBA";
    Exp    : String := "...";
    Result : String := "...";
  begin
    Exp( Exp'First     ) := Character'Val( 2 );
    Exp( Exp'First + 1 ) := Character'Val( 0 );
    Exp( Exp'First + 2 ) := Character'Val( 2 );

    Put_Line("----Testing xor on strings---");
    XOR_Strings( S1, S2, Result);
    Put_Line("S1 is " & S1);
    Put_Line("S2 is " & S2);
    Put_Line("Result is: ");
    for C of Result loop
      Put( Natural'Image( Character'Pos( C ) ) );
    end loop;
    new_line(1);

    if Result /= Exp then
      Put_Line("FAILED: xor on strings");
    else
      Put_Line("PASSED: xor on strings");
    end if;
  end test_xor_strings;

  procedure test_oaep is
    Msg     : String := "abcdefghij jihgfedcba123456789";
    LongMsg : String( 1..1000 ) := ( others => 'T' );
    Encr    : OAEP_Block := ( others => ' ' );
    Decr    : OAEP_HALF  := ( others => ' ' );
    Entropy : OAEP_Block := ( others => 'e' );
    Len     : Natural;
    Flag    : Boolean;
    C       : Character;
    MaxLen  : constant := 245;
  begin
    Put_Line("----Testing OAEP Encrypt----");
    OAEP_Encrypt( Msg, Entropy, Encr );

    Put_Line("----Testing OAEP Decrypt----");
    OAEP_Decrypt( Encr, Len, Decr, Flag );

    if Flag = False or 
       Len /= Msg'Length * 8 or 
       Decr( Decr'First .. Decr'First + Msg'Length - 1 ) /= Msg
       then
      Put_Line("FAILED: oaep test");
      Put_Line("Msg is: "  & Msg);
      Put_Line("Decr is: " & Decr);
      Put_Line("Flag is: " & Boolean'Image( Flag ) );
      Put_Line("Len is: "  & Natural'Image( Len ) );
    else
      Put_Line("PASSED: oaep test");
    end if;

    -- test decrypt on invalid (non-OAEP) string
    Flag := True;
    C := Encr( Encr'First );
    Encr( Encr'First ) := Character'Val( Character'Pos( C ) / 2 );
    Decr := ( others => ' ' );
    OAEP_Decrypt( Encr, Len, Decr, Flag );

    if Flag = True then
      Put_Line("FAILED: oaep test with invalid package");
    else
      Put_Line("PASSED: oaep test with invalid package");
    end if; 

    -- test encrypt on message longer than maximum payload (1096 bits)
    Flag := False;
    Len := 0;
    LongMsg( 1..Msg'Length ) := Msg;
    Encr := ( others => '.' );
    OAEP_Encrypt( LongMsg, Entropy, Encr);
    OAEP_Decrypt( Encr, Len, Decr, Flag);

    if Flag = False or 
       Len /= MaxLen * 8 or
       Decr( Decr'First .. Decr'First + Len / 8 - 1 ) /= 
             LongMsg( LongMsg'First..LongMsg'First + MaxLen - 1 ) 
       then
      Put_Line("FAILED: oaep test with too long message");
      Put_Line("Msg is: "  & LongMsg);
      Put_Line("Decr is: " & Decr);
      Put_Line("Flag is: " & Boolean'Image( Flag ) );
      Put_Line("Len is: "  & Natural'Image( Len ) );
    else
      Put_Line("PASSED: oaep test with too long message");
    end if;

  end test_oaep;

  procedure test_all_byterates is
    Input : constant String := "hello, world";
    Bin   : Bytestream( 0 .. Input'Length - 1 ) := ( others => 0 );
    Bout  : Bytestream( 0 .. 100 ) := ( others => 0 );
  begin
    ToBytestream( Input, Bin );
    Put_Line("Testing all bitrates:");
      for Byterate in Keccak_Rate'Range loop
        Sponge(Bin, Bout, Byterate);
        Put_Line("PASSED: keccak with byterate " & Integer'Image(Byterate));
      end loop;
  end test_all_byterates;

  procedure test_bitreverse is
    A: Interfaces.Unsigned_8;
    B: Interfaces.Unsigned_8;
    Pass: Boolean := True;
  begin
    for I in 0..255 loop
      -- reverse the bits
      A := Interfaces.Unsigned_8(I);
      B := 0;
      for J in 0..7 loop
        B := B*2 + A mod 2;
        A := A / 2;
      end loop;
      -- compare with tabled value
      if B /= Reverse_Table(I) then
        Pass := False;
        Put_Line("FAIL: reverse table value is wrong Table(" &
                 Integer'Image(I) & ")= " & 
                 Interfaces.Unsigned_8'Image(Reverse_Table(I)) &
                 " but should be " &
                 Interfaces.Unsigned_8'Image(B));
      end if;
    end loop;
    if Pass then 
      Put_Line("PASS: bitreverse table is correct.");
    end if;
  end test_bitreverse;

  -- end of helper methods 

	--variables
  T: Test_Round;
begin
  Put_Line("-----Testing with zero state as input------");
  if (not read_from_file("testvectorszero.txt", T)) then
    return;
  end if;

  for I in Round_Index loop
    Put_Line("---round " & Round_Index'Image(I) & "---");
    test_one_round(T(I), I);		
  end loop;

  -- test also Keccak_Function as a whole --
  test_keccak_function(T);

  Put_Line("-----Testing with non-zero state as input------");
  if (not read_from_file("testvectorsnonzero.txt", T)) then 
    return;
  end if;

  for I in Round_Index loop
    Put_Line("---round " & Round_Index'Image(I) & "---");
    test_one_round(T(I), I);		
  end loop;

  -- test also Keccak_Function as a whole --
  test_keccak_function(T);

  -- test BytesToWord and WordToBytes
  test_bytes_to_word_conversion;

  -- test Sponge construction
  test_sponge;

  -- test flipping octets
  test_flip;

  -- test bitstream conversion
  test_bytestream_conversion;

  -- test hash keccak (strings version)
  Put_Line("Sytem's BIT order is " & 
           System.Bit_Order'Image(System.Default_Bit_Order));
  test_hash_keccak;

  -- test oaep encrypt + decrypt
  test_oaep;

  -- test xor on strings
  test_xor_strings;

  -- test ALL byterates of the Keccak sponge
  test_all_byterates;

  -- test the values in the lookup table for bit reversing
  test_bitreverse;

end SMG_Keccak.Test;