#include <stdint.h>
#include <string.h>

#include "constants.h"
#include "parameters.h"
#include "tweakey.h"


#define LANE_BITS  64
#define LANE_BYTES (LANE_BITS/8)
#define LANES_NB (TWEAKEY_BYTES/LANE_BYTES)


void tweakey_state_init(
    tweakey_state *TK,
    const uint8_t key[KEY_BYTES],
    const uint8_t tweak[TWEAK_BYTES]
)
{
    memcpy(TK->TK,             tweak, TWEAK_BYTES);
    memcpy(TK->TK+TWEAK_BYTES, key,   KEY_BYTES);
}


void tweakey_state_extract(
    const tweakey_state *TK,
    uint8_t round_tweakey[ROUND_TWEAKEY_BYTES], /* output */
    uint8_t i                                   /* round constant */
)
{
    memset(round_tweakey, 0, ROUND_TWEAKEY_BYTES);

    for (const uint8_t *lane=TK->TK; lane<TK->TK+TWEAKEY_BYTES; lane+=LANE_BYTES)
    {
        for (size_t j=0; j<LANE_BYTES; j++)
        {
            round_tweakey[j] ^= lane[j];
        }
    }

    round_tweakey[0] ^= i;
}


static void _permute_state(tweakey_state *TK)
{
    uint8_t TK_old[TWEAKEY_BYTES];
    memcpy(TK_old, TK->TK, sizeof(TK_old));

    /* TODO: homogenize indices; here j=lane; k=byte */

    for (size_t j=0; j<TWEAKEY_BYTES; j+=LANE_BYTES)
    {
        for (size_t k=0; k<LANE_BYTES; k++)
        {
            TK->TK[j+h[k]] = TK_old[j+k];
        }
    }
}

static void _multiply_state(tweakey_state *TK)
{
    /* Lane 0 is multiplied by Id; lane 1 by P_0, lane 2 by P_1... */

    for (size_t lane=1; lane<LANES_NB; lane++)
    {
        const uint8_t* P_lane = P[lane-1];

        /* TODO: homogenize indices; here b=byte */

        for (size_t b=0; b<LANE_BYTES; b++)
        {
            size_t offset = lane*LANE_BYTES + b;
            TK->TK[offset] = P_lane[TK->TK[offset]];
        }
    }
}

void tweakey_state_update(tweakey_state *TK)
{
    _permute_state(TK);
    _multiply_state(TK);
}