// absolutely not production-ready

func crypto.rotl32[x: i64, r: i64] : i64
    return ((x << r) | (x >> (32 - r))) & 0xffffffff

func crypto.rotr64[x: i64, y: i64] : i64
    y = y & 63
    if y == 0
        return x
    let lhs_mask: i64 = ((1 << (64 - y)) - 1)
    let r: i64 = (x >> y) & lhs_mask
    let l: i64 = (x << (64 - y))
    return (r | l)

func crypto.blake2b._G[v: ptr, a: i64, b: i64, c: i64, d: i64, x: i64, y: i64] : void
    mem.write64(v + a * 8, mem.read64(v + a * 8) + mem.read64(v + b * 8) + x)
    mem.write64(v + d * 8, crypto.rotr64(mem.read64(v + d * 8) ^ mem.read64(v + a * 8), 32))
    mem.write64(v + c * 8, mem.read64(v + c * 8) + mem.read64(v + d * 8))
    mem.write64(v + b * 8, crypto.rotr64(mem.read64(v + b * 8) ^ mem.read64(v + c * 8), 24))
    mem.write64(v + a * 8, mem.read64(v + a * 8) + mem.read64(v + b * 8) + y)
    mem.write64(v + d * 8, crypto.rotr64(mem.read64(v + d * 8) ^ mem.read64(v + a * 8), 16))
    mem.write64(v + c * 8, mem.read64(v + c * 8) + mem.read64(v + d * 8))
    mem.write64(v + b * 8, crypto.rotr64(mem.read64(v + b * 8) ^ mem.read64(v + c * 8), 63))

func crypto.blake2b._compress[h: ptr, block: ptr, t0: i64, t1: i64, last: i64, iv: Array, sigma: Array] : void
    let v: ptr = mem.alloc(16 * 8)
    let m: ptr = mem.alloc(16 * 8)

    for j in 0..8
        mem.write64(v + j * 8, mem.read64(h + j * 8))
        mem.write64(v + (j + 8) * 8, array.nth(iv, j))

    mem.write64(v + 12 * 8, mem.read64(v + 12 * 8) ^ t0)
    mem.write64(v + 13 * 8, mem.read64(v + 13 * 8) ^ t1)
    if last
        mem.write64(v + 14 * 8, mem.read64(v + 14 * 8) ^ 0xffffffffffffffff)

    for j in 0..16
        let w = 0
        for k in 0..8
            w = w | ((block[j * 8 + k] & 0xff) << (8 * k))
        mem.write64(m + j * 8, w)

    for r in 0..12
        let base: i64 = r * 16
        crypto.blake2b._G(v, 0, 4, 8, 12, mem.read64(m + array.nth(sigma, base + 0) * 8), mem.read64(m + array.nth(sigma, base + 1) * 8))
        crypto.blake2b._G(v, 1, 5, 9, 13, mem.read64(m + array.nth(sigma, base + 2) * 8), mem.read64(m + array.nth(sigma, base + 3) * 8))
        crypto.blake2b._G(v, 2, 6, 10, 14, mem.read64(m + array.nth(sigma, base + 4) * 8), mem.read64(m + array.nth(sigma, base + 5) * 8))
        crypto.blake2b._G(v, 3, 7, 11, 15, mem.read64(m + array.nth(sigma, base + 6) * 8), mem.read64(m + array.nth(sigma, base + 7) * 8))
        crypto.blake2b._G(v, 0, 5, 10, 15, mem.read64(m + array.nth(sigma, base + 8) * 8), mem.read64(m + array.nth(sigma, base + 9) * 8))
        crypto.blake2b._G(v, 1, 6, 11, 12, mem.read64(m + array.nth(sigma, base + 10) * 8), mem.read64(m + array.nth(sigma, base + 11) * 8))
        crypto.blake2b._G(v, 2, 7, 8, 13, mem.read64(m + array.nth(sigma, base + 12) * 8), mem.read64(m + array.nth(sigma, base + 13) * 8))
        crypto.blake2b._G(v, 3, 4, 9, 14, mem.read64(m + array.nth(sigma, base + 14) * 8), mem.read64(m + array.nth(sigma, base + 15) * 8))

    for j in 0..8
        mem.write64(h + j * 8, mem.read64(h + j * 8) ^ (mem.read64(v + j * 8) ^ mem.read64(v + (j + 8) * 8)))
    
    mem.zero_and_free(v, 16 * 8)
    mem.zero_and_free(m, 16 * 8)

func crypto.blake2b.hash[outlen: i64, key: ptr, keylen: i64, input: ptr, inputlen: i64] : ptr
    if outlen == 0 || outlen > 64 || keylen > 64
        panic("invalid length passed to crypto.blake2b.hash")
    let out: ptr = mem.alloc(outlen)
    
    // IV[i] = floor(2**w * frac(sqrt(prime(i+1)))), where prime(i)
    // is the i:th prime number ( 2, 3, 5, 7, 11, 13, 17, 19 )
    // and sqrt(x) is the square root of x.
    // https://www.ietf.org/rfc/rfc7693#section-2.6
    let iv: Array = [0x6a09e667f3bcc908, 0xbb67ae8584caa73b, 0x3c6ef372fe94f82b, 0xa54ff53a5f1d36f1, 0x510e527fade682d1, 0x9b05688c2b3e6c1f, 0x1f83d9abfb41bd6b, 0x5be0cd19137e2179]

    // https://crypto.stackexchange.com/questions/108987/rationale-for-blake2-message-schedule
    let sigma: Array = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,14,10,4,8,9,15,13,6,1,12,0,2,11,7,5,3,11,8,12,0,5,2,15,13,10,14,3,6,7,1,9,4,7,9,3,1,13,12,11,14,2,6,5,10,4,0,15,8,9,0,5,7,2,4,10,15,14,1,11,12,6,8,3,13,2,12,6,10,0,11,8,3,4,13,7,5,15,14,1,9,12,5,1,15,14,13,4,10,0,7,6,3,9,2,8,11,13,11,7,14,12,1,3,9,5,0,15,4,8,6,2,10,6,15,14,9,11,3,0,8,12,2,13,7,1,4,10,5,10,2,8,4,7,6,1,5,15,11,9,14,3,12,13,0,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,14,10,4,8,9,15,13,6,1,12,0,2,11,7,5,3]
    
    let h: ptr = mem.alloc(8 * 8)
    let t0 = 0
    let t1 = 0
    let block: ptr = mem.alloc(128)
    let c = 0

    for i in 0..8
        mem.write64(h + i * 8, array.nth(iv, i))

    mem.write64(h, mem.read64(h) ^ (0x01010000 ^ ((keylen << 8) ^ outlen)))

    if keylen > 0
        mem.zero(block, 128)
        mem.copy(key, block, keylen)
        c = 128
    else
        c = 0
    
    for i in 0..inputlen
        if c == 128
            t0 = t0 + c
            if t0 < c
                t1 = t1 + 1
            crypto.blake2b._compress(h, block, t0, t1, 0, iv, sigma)
            c = 0
        block[c] = input[i]
        c = c + 1
    
    t0 = t0 + c
    if t0 < c
        t1 = t1 + 1
    
    if c < 128
        for i in (c)..128
            block[i] = 0
    crypto.blake2b._compress(h, block, t0, t1, 1, iv, sigma)

    for i in 0..outlen
        out[i] = (mem.read64(h + (i >> 3) * 8) >> (8 * (i & 7))) & 0xff

    mem.zero_and_free(h, 8 * 8)
    mem.zero_and_free(block, 128)
    array.free(iv)
    array.free(sigma)
    return out

func crypto.chacha20._quarter_round[state: ptr, a: i64, b: i64, c: i64, d: i64] : void
    let va: i64 = mem.read32(state + a * 4)
    let vb: i64 = mem.read32(state + b * 4)
    let vc: i64 = mem.read32(state + c * 4)
    let vd: i64 = mem.read32(state + d * 4)
    va = (va + vb) & 0xffffffff
    vd = vd ^ va
    vd = crypto.rotl32(vd, 16)
    vc = (vc + vd) & 0xffffffff
    vb = vb ^ vc
    vb = crypto.rotl32(vb, 12)
    va = (va + vb) & 0xffffffff
    vd = vd ^ va
    vd = crypto.rotl32(vd, 8)
    vc = (vc + vd) & 0xffffffff
    vb = vb ^ vc
    vb = crypto.rotl32(vb, 7)
    mem.write32(state + a * 4, va)
    mem.write32(state + b * 4, vb)
    mem.write32(state + c * 4, vc)
    mem.write32(state + d * 4, vd)

func crypto.xchacha20._permute[state: ptr] : void
    for i in 0..10
        crypto.chacha20._quarter_round(state, 0, 4, 8, 12)
        crypto.chacha20._quarter_round(state, 1, 5, 9, 13)
        crypto.chacha20._quarter_round(state, 2, 6, 10, 14)
        crypto.chacha20._quarter_round(state, 3, 7, 11, 15)
        crypto.chacha20._quarter_round(state, 0, 5, 10, 15)
        crypto.chacha20._quarter_round(state, 1, 6, 11, 12)
        crypto.chacha20._quarter_round(state, 2, 7, 8, 13)
        crypto.chacha20._quarter_round(state, 3, 4, 9, 14)

func crypto.xchacha20._block[key: ptr, nonce: ptr, blocknum: i64, out: ptr] : void
    let sigma: str = "expand 32-byte k"
    let state: ptr = mem.alloc(16 * 4)

    mem.write32(state + 0, mem.read32(sigma + 0))
    mem.write32(state + 4, mem.read32(sigma + 4))
    mem.write32(state + 8, mem.read32(sigma + 8))
    mem.write32(state + 12, mem.read32(sigma + 12))

    for i in 0..8
        mem.write32(state + (4 + i) * 4, mem.read32(key + i * 4))

    mem.write32(state + 12 * 4, blocknum)
    mem.write32(state + 13 * 4, mem.read32(nonce + 0))
    mem.write32(state + 14 * 4, mem.read32(nonce + 4))
    mem.write32(state + 15 * 4, mem.read32(nonce + 8))

    let working: ptr = mem.alloc(16 * 4)
    for i in 0..16
        mem.write32(working + i * 4, mem.read32(state + i * 4))

    crypto.xchacha20._permute(working)

    for i in 0..16
        let v: i64 = (mem.read32(working + i * 4) + mem.read32(state + i * 4)) & 0xffffffff
        mem.write32(out + i * 4, v)
    mem.zero_and_free(working, 16 * 4)
    mem.zero_and_free(state, 16 * 4)

func crypto.xchacha20._hchacha20[key: ptr, input: ptr, out32: ptr] : void
    let sigma: str = "expand 32-byte k"
    let state: ptr = mem.alloc(16 * 4)

    mem.write32(state + 0, mem.read32(sigma + 0))
    mem.write32(state + 4, mem.read32(sigma + 4))
    mem.write32(state + 8, mem.read32(sigma + 8))
    mem.write32(state + 12, mem.read32(sigma + 12))

    for i in 0..8
        mem.write32(state + (4 + i) * 4, mem.read32(key + i * 4))

    for i in 0..4
        mem.write32(state + (12 + i) * 4, mem.read32(input + i * 4))

    crypto.xchacha20._permute(state)

    for i in 0..4
        mem.write32(out32 + i * 4, mem.read32(state + i * 4))
    for i in 0..4
        mem.write32(out32 + 16 + i * 4, mem.read32(state + (12 + i) * 4))
    mem.zero_and_free(state, 16 * 4)

func crypto.xchacha20._stream[key: ptr, nonce: ptr, out: ptr, len: i64] : void
    let subkey: ptr = mem.alloc(32)
    crypto.xchacha20._hchacha20(key, nonce, subkey)

    let nonce12: ptr = mem.alloc(12)
    mem.zero(nonce12, 12)
    mem.copy(nonce + 16, nonce12 + 4, 8)

    let blocknum = 0
    let remaining: i64 = len
    let block: ptr = mem.alloc(64)

    while remaining > 0
        crypto.xchacha20._block(subkey, nonce12, blocknum, block)
        let take = 64
        if remaining < 64
            take = remaining
        mem.copy(block, out + len - remaining, take)
        remaining = remaining - take
        blocknum = blocknum + 1
    mem.zero_and_free(block, 64)
    mem.zero_and_free(nonce12, 12)
    mem.zero_and_free(subkey, 32)

// expects 32-byte key and 24-byte nonce
// no way to verify that until we get actual strings
func crypto.xchacha20.xor_no_auth[key: ptr, nonce: ptr, input: ptr, len: i64] : ptr
    if len <= 0
        let out: ptr = mem.alloc(1)
        out[0] = 0
        return out
    let out: ptr = mem.alloc(len)
    let ks: ptr = mem.alloc(len)
    crypto.xchacha20._stream(key, nonce, ks, len)
    for i in 0..len
        out[i] = input[i] ^ ks[i]
    mem.zero_and_free(ks, len)
    return out

func crypto.x25519.carry[elem: ptr] : void
    for i in 0..16
        let carry: i64 = mem.read64(elem + i * 8) >> 16
        mem.write64(elem + i * 8, mem.read64(elem + i * 8) - (carry << 16))
        if i < 15
            mem.write64(elem + (i + 1) * 8, mem.read64(elem + (i + 1) * 8) + carry)
        else
            mem.write64(elem, mem.read64(elem) + 38 * carry)

func crypto.x25519.fadd[out: ptr, a: ptr, b: ptr] : void
    for i in 0..16
        mem.write64(out + i * 8, mem.read64(a + i * 8) + mem.read64(b + i * 8))

func crypto.x25519.fsub[out: ptr, a: ptr, b: ptr] : void
    for i in 0..16
        mem.write64(out + i * 8, mem.read64(a + i * 8) - mem.read64(b + i * 8))

func crypto.x25519.fmul[out: ptr, a: ptr, b: ptr] : void
    let product: ptr = mem.alloc(31 * 8)
    for i in 0..31
        mem.write64(product + i * 8, 0)
    for i in 0..16
        for j in 0..16
            mem.write64(product + (i + j) * 8, mem.read64(product + (i + j) * 8) + (mem.read64(a + i * 8) * mem.read64(b + j * 8)))
    for i in 0..15
        mem.write64(product + i * 8, mem.read64(product + i * 8) + 38 * mem.read64(product + (i + 16) * 8))
    for i in 0..16
        mem.write64(out + i * 8, mem.read64(product + i * 8))

    crypto.x25519.carry(out)
    crypto.x25519.carry(out)
    mem.zero_and_free(product, 31 * 8)

func crypto.x25519.finverse[out: ptr, input: ptr] : void
    let c: ptr = mem.alloc(16 * 8)
    for i in 0..16
        mem.write64(c + i * 8, mem.read64(input + i * 8))
    
    let i = 253
    while i >= 0
        crypto.x25519.fmul(c, c, c)
        if i != 2 && i != 4
            crypto.x25519.fmul(c, c, input)
        i = i - 1

    for i in 0..16
        mem.write64(out + i * 8, mem.read64(c + i * 8))
    mem.zero_and_free(c, 16 * 8)

func crypto.x25519.swap[p: ptr, q: ptr, bit: i64] : void
    for i in 0..16
        let t: i64 = (-bit) & (mem.read64(p + i * 8) ^ mem.read64(q + i * 8))
        mem.write64(p + i * 8, mem.read64(p + i * 8) ^ t)
        mem.write64(q + i * 8, mem.read64(q + i * 8) ^ t)

func crypto.x25519.unpack[out: ptr, input: ptr] : void
    for i in 0..16
        mem.write64(out + i * 8, input[i * 2] + (input[i * 2 + 1] << 8))
    mem.write64(out + 8 * 15, mem.read64(out + 8 * 15) & 0x7fff)

func crypto.x25519.pack[out: ptr, input: ptr] : void
    let t: ptr = mem.alloc(16 * 8)
    for i in 0..16
        mem.write64(t + i * 8, mem.read64(input + i * 8))
    let m: ptr = mem.alloc(16 * 8)

    crypto.x25519.carry(t)
    crypto.x25519.carry(t)
    crypto.x25519.carry(t)
    for j in 0..2
        mem.write64(m, mem.read64(t) - 0xffed)
        for i in 1..15
            mem.write64(m + i * 8, mem.read64(t + i * 8) - 0xffff - ((mem.read64(m + (i - 1) * 8) >> 16) & 1))
            mem.write64(m + (i - 1) * 8, mem.read64(m + (i - 1) * 8) & 0xffff)
        mem.write64(m + 15 * 8, mem.read64(t + 15 * 8) - 0x7fff - ((mem.read64(m + 14 * 8) >> 16) & 1))
        let carry: i64 = (mem.read64(m + 15 * 8) >> 16) & 1
        mem.write64(m + 14 * 8, mem.read64(m + 14 * 8) & 0xffff)
        crypto.x25519.swap(t, m, 1 - carry)
    
    for i in 0..16
        let v: i64 = mem.read64(t + i * 8)
        out[i * 2] = v & 0xff
        out[i * 2 + 1] = (v >> 8) & 0xff

    mem.zero_and_free(t, 16 * 8)
    mem.zero_and_free(m, 16 * 8)

func crypto.x25519.scalarmult[scalar: ptr, point: ptr] : ptr
    let clamped: ptr = mem.alloc(32)
    let a: ptr = mem.alloc(16 * 8)
    let b: ptr = mem.alloc(16 * 8)
    let c: ptr = mem.alloc(16 * 8)
    let d: ptr = mem.alloc(16 * 8)
    let e: ptr = mem.alloc(16 * 8)
    let f: ptr = mem.alloc(16 * 8)
    let x: ptr = mem.alloc(16 * 8)

    // "A" parameter in Curve25519
    // 121665 = (486662 - 2)/4
    let magic: ptr = mem.alloc(16 * 8)
    mem.zero(magic, 16 * 8)
    mem.write64(magic, 0xdb41)
    mem.write64(magic + 8, 1)

    // copy and clamp scalar
    mem.copy(scalar, clamped, 32)
    clamped[0] = clamped[0] & 0xf8
    clamped[31] = (clamped[31] & 0x7f) | 0x40

    // load point
    crypto.x25519.unpack(x, point)

    // initialize ladder state
    for i in 0..16
        mem.write64(a + i * 8, 0)
        mem.write64(b + i * 8, mem.read64(x + i * 8))
        mem.write64(c + i * 8, 0)
        mem.write64(d + i * 8, 0)
    mem.write64(a, 1)
    mem.write64(d, 1)

    let i = 254
    while i >= 0
        let bit: i64 = (clamped[i >> 3] >> (i & 7)) & 1
        crypto.x25519.swap(a, b, bit)
        crypto.x25519.swap(c, d, bit)
        crypto.x25519.fadd(e, a, c)
        crypto.x25519.fsub(a, a, c)
        crypto.x25519.fadd(c, b, d)
        crypto.x25519.fsub(b, b, d)
        crypto.x25519.fmul(d, e, e)
        crypto.x25519.fmul(f, a, a)
        crypto.x25519.fmul(a, c, a)
        crypto.x25519.fmul(c, b, e)
        crypto.x25519.fadd(e, a, c)
        crypto.x25519.fsub(a, a, c)
        crypto.x25519.fmul(b, a, a)
        crypto.x25519.fsub(c, d, f)
        crypto.x25519.fmul(a, c, magic)
        crypto.x25519.fadd(a, a, d)
        crypto.x25519.fmul(c, c, a)
        crypto.x25519.fmul(a, d, f)
        crypto.x25519.fmul(d, b, x)
        crypto.x25519.fmul(b, e, e)
        crypto.x25519.swap(a, b, bit)
        crypto.x25519.swap(c, d, bit)
        i = i - 1

    crypto.x25519.finverse(c, c)
    crypto.x25519.fmul(a, a, c)
    let out: ptr = mem.alloc(32)
    crypto.x25519.pack(out, a)

    mem.zero_and_free(clamped, 32)
    mem.zero_and_free(a, 16 * 8)
    mem.zero_and_free(b, 16 * 8)
    mem.zero_and_free(c, 16 * 8)
    mem.zero_and_free(d, 16 * 8)
    mem.zero_and_free(e, 16 * 8)
    mem.zero_and_free(f, 16 * 8)
    mem.zero_and_free(x, 16 * 8)
    mem.free(magic)
    return out