// stdlib/crypto/shake_stream.sy — streaming SHAKE-128 / SHAKE-256 wrapper
//
// Kyber needs to absorb a seed ONCE and then squeeze output in chunks as
// long as it needs (rejection sampling for matrix A doesn't know in advance
// how many bytes it'll consume). Today's VM only exposes one-shot XOF.
//
// This module fakes streaming on top of one-shot by maintaining (seed,
// cumulative_bytes_requested) and re-absorbing every time we need more
// output. Correct but O(n^2) in total bytes consumed — fine for correctness
// tests, unacceptable for production. Phase 3 §1.2 replaces the internals
// with real streaming builtins, and the public API here stays identical.
//
// Aislinn McCubbins, Apr 11 2026

// ── Streaming SHAKE-128 state ────────────────────────────
// A Shake128State is a triple:
//   seed:     the absorbed input (Tensor of bytes)
//   buf:      currently cached output (Tensor of bytes)
//   cursor:   number of bytes already consumed from buf
//
// Invariant: buf always contains shake_128(seed, len(buf)), so re-absorbing
// with a longer length gives a strict prefix-extension of the current buf.
// SHAKE's one-shot and streaming variants agree on the output stream, so
// this is correctness-preserving.

struct Shake128State {
    seed: Tensor,
    buf: Tensor,
    cursor: Int,
}

// Construct a new streaming hasher seeded with `seed_bytes`.
fn shake128_new(seed_bytes: Tensor) -> Shake128State {
    let empty = t_zeros(0)
    return Shake128State {
        seed: seed_bytes,
        buf: empty,
        cursor: 0,
    }
}

// Absorb additional bytes. For streaming usage this must be called BEFORE
// the first squeeze. Today's fallback concatenates onto the seed and resets
// the cache; once real streaming lands, this will be a true absorb call.
fn shake128_absorb(s: Shake128State, more: Tensor) -> Shake128State {
    let new_seed = t_cat(s.seed, more)
    let empty = t_zeros(0)
    return Shake128State {
        seed: new_seed,
        buf: empty,
        cursor: 0,
    }
}

// Squeeze `n` more bytes from the stream, returning [new_state, bytes].
// The returned bytes are a Tensor of length `n` (f32-encoded 0..255).
//
// Strategy: if cursor + n exceeds the current buf length, re-hash with a
// bigger buffer (double-and-cover). Then slice [cursor .. cursor+n] and
// advance cursor.
fn shake128_squeeze(s: Shake128State, n: Int) -> Array {
    let needed = s.cursor + n
    let have = tensor_len(s.buf)

    let mut new_buf = s.buf
    if have < needed {
        // Re-absorb with enough room. Double the needed size so that
        // repeated small squeezes don't pay O(n) re-hash cost each time.
        let alloc = needed * 2
        new_buf = shake_128(s.seed, alloc)
    }

    let out = tensor_slice(new_buf, s.cursor, s.cursor + n)
    let new_state = Shake128State {
        seed: s.seed,
        buf: new_buf,
        cursor: s.cursor + n,
    }

    let mut result = array_new()
    result = array_push(result, new_state)
    result = array_push(result, out)
    return result
}

// ── Streaming SHAKE-256 state (same shape) ───────────────

struct Shake256State {
    seed: Tensor,
    buf: Tensor,
    cursor: Int,
}

fn shake256_new(seed_bytes: Tensor) -> Shake256State {
    let empty = t_zeros(0)
    return Shake256State {
        seed: seed_bytes,
        buf: empty,
        cursor: 0,
    }
}

fn shake256_absorb(s: Shake256State, more: Tensor) -> Shake256State {
    let new_seed = t_cat(s.seed, more)
    let empty = t_zeros(0)
    return Shake256State {
        seed: new_seed,
        buf: empty,
        cursor: 0,
    }
}

fn shake256_squeeze(s: Shake256State, n: Int) -> Array {
    let needed = s.cursor + n
    let have = tensor_len(s.buf)

    let mut new_buf = s.buf
    if have < needed {
        let alloc = needed * 2
        new_buf = shake_256(s.seed, alloc)
    }

    let out = tensor_slice(new_buf, s.cursor, s.cursor + n)
    let new_state = Shake256State {
        seed: s.seed,
        buf: new_buf,
        cursor: s.cursor + n,
    }

    let mut result = array_new()
    result = array_push(result, new_state)
    result = array_push(result, out)
    return result
}

// ── Convenience: domain-separated PRF ────────────────────
// Kyber uses PRF(seed, nonce) = SHAKE-256(seed || nonce_byte, out_len).
// This is always one-shot in Kyber (the output length is known: 2*eta*n/8),
// so we expose a direct helper that bypasses the streaming state machine.
fn prf_256(seed: Tensor, nonce: Int, out_len: Int) -> Tensor {
    let nonce_t = tensor_from_ints([nonce])
    let input = t_cat(seed, nonce_t)
    return shake_256(input, out_len)
}

// Kyber's XOF for matrix expansion is SHAKE-128 with a 34-byte input:
//   XOF(rho, i, j) = SHAKE-128(rho || i_byte || j_byte) [streaming]
// Used by the matrix-A sampler with rejection.
fn xof_128_init(rho: Tensor, i: Int, j: Int) -> Shake128State {
    let idx = tensor_from_ints([i, j])
    let seed = t_cat(rho, idx)
    return shake128_new(seed)
}
