(** SHA-3 (SHA3-256) — pure OCaml, WASM-safe.

    Keccak-f[1600] permutation + SHA-3 multi-rate padding (domain byte
    0x06, NOT the legacy Keccak 0x01). Reference: FIPS 202. No deps. *)

let ( ^: ) = Int64.logxor
let ( &: ) = Int64.logand
let lnot64 = Int64.lognot

let rotl64 x n =
  if n = 0 then x
  else
    Int64.logor (Int64.shift_left x n) (Int64.shift_right_logical x (64 - n))

(* FIPS 202 Table 2 — ρ rotation offsets, indexed lane = x + 5*y. *)
let rho = [|
   0;  1; 62; 28; 27;
  36; 44;  6; 55; 20;
   3; 10; 43; 25; 39;
  41; 45; 15; 21;  8;
  18;  2; 61; 56; 14 |]

(* FIPS 202 §3.2.5 — round constants RC[0..23] for ι. *)
let rc = [|
  0x0000000000000001L; 0x0000000000008082L; 0x800000000000808aL;
  0x8000000080008000L; 0x000000000000808bL; 0x0000000080000001L;
  0x8000000080008081L; 0x8000000000008009L; 0x000000000000008aL;
  0x0000000000000088L; 0x0000000080008009L; 0x000000008000000aL;
  0x000000008000808bL; 0x800000000000008bL; 0x8000000000008089L;
  0x8000000000008003L; 0x8000000000008002L; 0x8000000000000080L;
  0x000000000000800aL; 0x800000008000000aL; 0x8000000080008081L;
  0x8000000000008080L; 0x0000000080000001L; 0x8000000080008008L |]

let keccak_f (a : int64 array) : unit =
  let c = Array.make 5 0L and d = Array.make 5 0L in
  let b = Array.make 25 0L in
  for round = 0 to 23 do
    (* θ *)
    for x = 0 to 4 do
      c.(x) <- a.(x) ^: a.(x + 5) ^: a.(x + 10)
               ^: a.(x + 15) ^: a.(x + 20)
    done;
    for x = 0 to 4 do
      d.(x) <- c.((x + 4) mod 5) ^: (rotl64 c.((x + 1) mod 5) 1)
    done;
    for x = 0 to 4 do
      for y = 0 to 4 do
        a.(x + 5 * y) <- a.(x + 5 * y) ^: d.(x)
      done
    done;
    (* ρ and π: B[y, 2x+3y] = rotl(A[x,y], rho[x,y]) *)
    for x = 0 to 4 do
      for y = 0 to 4 do
        let nx = y and ny = (2 * x + 3 * y) mod 5 in
        b.(nx + 5 * ny) <- rotl64 a.(x + 5 * y) rho.(x + 5 * y)
      done
    done;
    (* χ *)
    for y = 0 to 4 do
      for x = 0 to 4 do
        a.(x + 5 * y) <-
          b.(x + 5 * y)
          ^: ((lnot64 b.((x + 1) mod 5 + 5 * y))
              &: b.((x + 2) mod 5 + 5 * y))
      done
    done;
    (* ι *)
    a.(0) <- a.(0) ^: rc.(round)
  done

let sha3_256_hex (msg : string) : string =
  let rate = 136 (* bytes: (1600 - 2*256) / 8 *) in
  let len = String.length msg in
  (* pad10*1 with SHA-3 domain byte 0x06; last byte ORed with 0x80. *)
  let q = rate - (len mod rate) in
  let padded = Bytes.make (len + q) '\000' in
  Bytes.blit_string msg 0 padded 0 len;
  if q = 1 then
    Bytes.set padded len '\x86'
  else begin
    Bytes.set padded len '\x06';
    Bytes.set padded (len + q - 1) '\x80'
  end;
  let total = Bytes.length padded in
  let a = Array.make 25 0L in
  let nblocks = total / rate in
  for blk = 0 to nblocks - 1 do
    let base = blk * rate in
    (* Absorb: XOR rate bytes into the state, little-endian lanes. *)
    for j = 0 to rate - 1 do
      let lane = j / 8 and sh = (j mod 8) * 8 in
      let byte = Int64.of_int (Char.code (Bytes.get padded (base + j))) in
      a.(lane) <- a.(lane) ^: (Int64.shift_left byte sh)
    done;
    keccak_f a
  done;
  (* Squeeze 32 bytes (fits in the first 4 lanes; rate > 32). *)
  let out = Buffer.create 64 in
  for j = 0 to 31 do
    let lane = j / 8 and sh = (j mod 8) * 8 in
    let byte =
      Int64.to_int
        (Int64.logand (Int64.shift_right_logical a.(lane) sh) 0xFFL)
    in
    Buffer.add_string out (Printf.sprintf "%02x" byte)
  done;
  Buffer.contents out
