struct SHA256_CTX {
  data : array<u32, 64>, 
  bitlen: array<u32, 2>, 
  state: array<u32, 8>, 
  // info: array<u32, 8>,
}

const SHA256_BLOCK_SIZE = 32; 
const INPUT_BLOCK_SIZE  = 64; 
const INPUT_BLOCK_MAX   = 8; 
const INPUT_PACKED_SIZE = (INPUT_BLOCK_SIZE * INPUT_BLOCK_MAX / 4); 

struct InputBlocks { 
  data: array<u32, 128>,
}
struct OutputHash {
  data: array<u32, 8>, 
} 

@group(0) @binding(0) var<storage, read> inputs: array<InputBlocks>; 
@group(0) @binding(1) var<storage, read> sizes: array<u32>; 
@group(0) @binding(2) var<storage, read_write> results: array<OutputHash>; 


const k = array<u32, 64> (
  0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
  0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
  0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
  0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
  0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
  0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
  0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
  0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2
);

const masks = array<u32, 4> (
  0x000000FFu, 
  0x0000FF00u,
  0x00FF0000u, 
  0xFF000000u, 
);

const full_masks = array<u32, 5> (
  0x00000000u, 
  0x000000FFu, 
  0x0000FFFFu,
  0x00FFFFFFu, 
  0xFFFFFFFFu, 
); 

fn ROTLEFT(a : u32, b : u32) -> u32{return (((a) << (b)) | ((a) >> (32-(b))));}
fn ROTRIGHT(a : u32, b : u32) -> u32{return (((a) >> (b)) | ((a) << (32-(b))));}

fn CH(e : u32, f : u32, g : u32) -> u32{return (((e) & (f)) ^ (~(e) & (g)));}
fn MAJ(a : u32, b : u32, c: u32) -> u32{return (((a) & (b)) ^ ((a) & (c)) ^ ((b) & (c)));}
fn EP0(x : u32) -> u32{return (ROTRIGHT(x,2) ^ ROTRIGHT(x,13) ^ ROTRIGHT(x,22));}
fn EP1(x : u32) -> u32{return (ROTRIGHT(x,6) ^ ROTRIGHT(x,11) ^ ROTRIGHT(x,25));}
fn SIG0(x : u32) -> u32{return (ROTRIGHT(x,7) ^ ROTRIGHT(x,18) ^ ((x) >> 3));}
fn SIG1(x : u32) -> u32{return (ROTRIGHT(x,17) ^ ROTRIGHT(x,19) ^ ((x) >> 10));}

fn flip_endian(x: u32) -> u32 {
  var y: u32 = 0;
  // 11 22 33 44 => 44 33 22 11
  // 11 >> 3*8 << 0*8;  
  // 22 >> 2*8 << 1*8;  
  // 33 >> 1*8 << 2*8;  
  // 44 >> 0*8 << 3*8;  
  for (var i: u32; i < 4; i++) {
    y |= ((x >> (i*8)) & 0xFF) << ((3-i) * 8); 
  } 
  return y; 
}

fn sha256_init(ctx: ptr<function, SHA256_CTX>) {
  
  // CTX INIT
  (*ctx).bitlen[0] = 0;
  (*ctx).bitlen[1] = 0;
  (*ctx).state[0] = 0x6a09e667;
  (*ctx).state[1] = 0xbb67ae85;
  (*ctx).state[2] = 0x3c6ef372;
  (*ctx).state[3] = 0xa54ff53a;
  (*ctx).state[4] = 0x510e527f;
  (*ctx).state[5] = 0x9b05688c;
  (*ctx).state[6] = 0x1f83d9ab;
  (*ctx).state[7] = 0x5be0cd19;
} 

fn sha256_transform(ctx: ptr<function, SHA256_CTX>) {
  var a : u32 = (*ctx).state[0];
  var b : u32 = (*ctx).state[1];
  var c : u32 = (*ctx).state[2];
  var d : u32 = (*ctx).state[3];
  var e : u32 = (*ctx).state[4];
  var f : u32 = (*ctx).state[5];
  var g : u32 = (*ctx).state[6];
  var h : u32 = (*ctx).state[7];
  var t1 : u32;
  var t2 : u32;
  

  for (var x: u32 = 16; x < 64; x++) {
    // W_16 = w_0 + sig0(w_1) + w_9 + sig1(w_14)
    (*ctx).data[x]= (*ctx).data[x-16] + SIG0((*ctx).data[x-15]) + (*ctx).data[x-7] + SIG1((*ctx).data[x-2]);
    // if (x >= 16 && x < 24) {
    //   var z = x - 16;
    //   (*ctx).info[z] = (*ctx).data[x];
    // }
  } 

  
  for (var x: u32 = 0; x < 64; x++) {
    t1 = h + EP1(e) + CH(e, f, g) + k[x] + (*ctx).data[x];  
    t2 = EP0(a) + MAJ(a,b,c); 
    h = g; 
    g = f;
    f = e;
    e = d + t1; 
    d = c; 
    c = b; 
    b = a; 
    a = t1 + t2; 
  }

  (*ctx).state[0] += a;
  (*ctx).state[1] += b;
  (*ctx).state[2] += c;
  (*ctx).state[3] += d;
  (*ctx).state[4] += e;
  (*ctx).state[5] += f;
  (*ctx).state[6] += g;
  (*ctx).state[7] += h;

} 

fn sha256_append_len(ctx: ptr<function, SHA256_CTX>, len: u32) {
  if ((*ctx).bitlen[0] > 0xffffffff - (len)) {
    (*ctx).bitlen[1]++; 
  }
  (*ctx).bitlen[0] += len; 
} 


fn sha256_update(ctx: ptr<function, SHA256_CTX>, index: u32) {
  var len = sizes[index];
  var block_index = 0u; 
  while len > 64 {
    for (var i : u32 = 0; i < 16; i++) {
      (*ctx).data[i] = flip_endian(inputs[index].data[16 * block_index + i]);
    }
    sha256_transform(ctx);
    sha256_append_len(ctx, 512);
    block_index += 1; 
    len -= 64; 
  } 
  var imax = len/4;  
  // if  (len % 4 > 0) {imax += 1;}  
  for (var i : u32 = 0; i < (imax); i++) {
      (*ctx).data[i] = flip_endian(inputs[index].data[16 * block_index + i]);
  }    
  for(var i = imax; i < 16; i ++) {
        (*ctx).data[i] = 0; 
  }

  var offset   : u32     = (len) % 4; 
  var last_word: u32     = (len - offset) / 4; 
  // if (offset == 0) {last_word += 1;}
  var padding_word: u32  = (0x80u << (8 * offset));
  (*ctx).data[last_word] = flip_endian((inputs[index].data[16*block_index + last_word] & full_masks[offset]) | padding_word); 

  // for (var i: u32 = 0; i < 8; i++) {
  //    (*ctx).info[i] = (*ctx).data[i];  
  // }
  // for (var i: u32 = 14; i < 16; i++) {
  //    (*ctx).info[i] = (*ctx).data[i];  
  // }

  if (len > 56) {
    sha256_transform(ctx);
    for(var i = 0; i < 14; i ++) {
       (*ctx).data[i] = 0; 
    }
  }  

  var final_len: u32 = (len*8);
  sha256_append_len(ctx, final_len);
  var upper = (*ctx).bitlen[0];  
  var lower = (*ctx).bitlen[1];

  (*ctx).data[15] = upper; 
  (*ctx).data[14] = lower; 

  sha256_transform(ctx);
}

@compute @workgroup_size(64,1,1) 
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  var ctx: SHA256_CTX; 
  var index: u32 = global_id.x; 
  sha256_init(&ctx); 
  sha256_update(&ctx, index); 

  for(var i: u32 = 0; i < 8; i++) {
      results[index].data[i] = flip_endian(ctx.state[i]);
  }
  // for(var i: u32 = 0; i < 8; i++) {
  //      results[index].data[i] = inputs[index].data[i];
  // }

  // results[index].data[0] = ctx.info;
}