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 Interfaces; use Interfaces;

procedure SMG_Keccak.Test is
  --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_bitstream(B: in Bitstream; Title: in String) is
    Hex       : array(0..15) of Character := ("0123456789ABCDEF");
    HexString : String(1..B'Length/4);
    C         : Natural;
    Pos       : Natural;
  begin
    for I in 1..B'Length/4 loop
      Pos := (I-1)*4 + B'First;
      C := Natural( B(Pos) ) * 8 +
           Natural( B(Pos + 1) ) * 4 +
           Natural( B(Pos + 2) ) * 2 +
           Natural( B(Pos + 3) );
			HexString(I) := Hex(C);
    end loop;
    Put_Line("---" & Title & "---");
    Put_Line(HexString);
  end print_bitstream;

  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_bits_to_word_conversion is
    bits: Bitword := (others => 0);
    obtained_bits: Bitword := (others => 0);
    expected: ZWord;
    obtained: ZWord;
  begin
    expected := 16#E7DDE140798F25F1#;
    bits := (1,1,1,0, 0,1,1,1, 1,1,0,1, 1,1,0,1, 1,1,1,0, 0,0,0,1, 0,1,0,0, 
             0,0,0,0, 0,1,1,1, 1,0,0,1, 1,0,0,0, 1,1,1,1, 0,0,1,0, 0,1,0,1, 
             1,1,1,1, 0,0,0,1);
    obtained := BitsToWord(bits);
    obtained_bits := WordToBits(expected);

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

    if obtained_bits /= bits then
      Put_Line("FAIL: word to bits");
      Put("Expected: ");
      for I in Bitword'Range loop
        Put(Bit'Image(bits(I)));
      end loop;
      Put_Line("");
      Put_Line("Obtained: ");
      for I in Bitword'Range loop
        Put(Bit'Image(obtained_bits(I)));
      end loop;
      Put_Line("");
    else
      Put_Line("PASSED: word to bits");
    end if;
  end test_bits_to_word_conversion;

  procedure test_sponge is
    Bitrate   : constant Keccak_Rate := 1344;
    Input     : Bitstream(1..5) := (1, 1, 0, 0, 1);
    Output    : Bitstream(1..Bitrate*2);
    Hex       : array(0..15) of Character := ("0123456789ABCDEF");
    HexString : String(1..Bitrate/2);
    C         : Natural;
    ExpHex    : String(1..Bitrate/2);
    Error     : Natural;
    Pos       : Natural;
  begin
    ExpHex := "B57B7DAED6330F79BA5783C5D45EABFFA1461FAC6CEA09BD"&
              "AAC114F17E23E5B349EECBC907E07FA36ECF8374079811E8"&
              "5E49243D04182C389E68C733BE698468423DB9891D3A7B10"&
              "320E0356AB4AB916F68C0EA20365A1D4DBA48218CA89CBB8"& 
              "6D08A34E04544D4100FFE9CB138EADC2D3FC0E8CC2BC15A7"&
              "5B950776970BFC310F33BF609630D73CAD918CF54657589E"&
              "42CF7CBF20DE677D2AB7E49389F6F6C3B3FE2992905325CE"&
              "60931C1515043595ADC1619CB7E034EF52BDC485D03B7FDD"&
              "7345E849FFB4C4426195C8D88C1E7BF9ADA41B92E006C3DA"& 
              "F1ED0FD63ADD9408A3FC815F727457692727637687C1F79D"& 
              "837DE20798E64C878181C02DF56A533F684459E8A03C8EF6"& 
              "234854531110E6CD9BDEFEA85E35C802B1ACDDF29C9332E2"& 
              "53C0FA72F3ED1ABA274838CFE6EF8BD572E89E1C2135F6A7"& 
              "5BC5D6EA4F85C9A757E68E561A56AC0FC19F1F086C43272F";

    Put_Line("---sponge test---");
    Sponge(Input, Bitrate, Output);
    Put_Line("Input is:");
    for I of Input loop
      Put(Bit'Image(I));
    end loop;
    new_line(1);

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

    Error := 0;
    for I in 1..Output'Length/4 loop
      Pos := Output'First + (I-1)*4;
      C := Natural( Output( Pos ) ) * 8 +
           Natural( Output( Pos + 1 ) ) * 4 +
           Natural( Output( Pos + 2 ) ) * 2 +
           Natural( Output( Pos + 3 ) );
			Hexstring(I) := Hex(C);
      if Hexstring(I) /= ExpHex(I) then
        Error := Error + 1;
      end if;
    end loop;
    Put_Line("Expected: ");
    Put_Line(ExpHex);
    Put_Line("Obtained: ");
    Put_Line(Hexstring);
    Put_Line("Errors found: " & Natural'Image(Error));

  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;

  -- 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 BitsToWord and WordToBits
  test_bits_to_word_conversion;

  -- test Sponge construction
  test_sponge;

end SMG_Keccak.Test;