#include "tweetnacl.h"
#include "tweetpwhash.h"
#define FOR(i, n) for (i = 0; i < n; ++i)
#define MIN(a,b) (((a)<(b))?(a):(b))
#define sv static void

#define HB crypto_hash_BYTES
#define BB crypto_hashblocks_BLOCKBYTES
#define SB crypto_pwhash_SALTBYTES
#define CB crypto_kdf_CONTEXTBYTES

typedef unsigned char u8;
typedef unsigned long long u64;

sv ts64(u8 *x, u64 u)
{
    int i;
    for (i = 7;i >= 0;--i) { x[i] = u; u >>= 8; }
}

static const u8 iv[64] = {
    0x6a,0x09,0xe6,0x67,0xf3,0xbc,0xc9,0x08,
    0xbb,0x67,0xae,0x85,0x84,0xca,0xa7,0x3b,
    0x3c,0x6e,0xf3,0x72,0xfe,0x94,0xf8,0x2b,
    0xa5,0x4f,0xf5,0x3a,0x5f,0x1d,0x36,0xf1,
    0x51,0x0e,0x52,0x7f,0xad,0xe6,0x82,0xd1,
    0x9b,0x05,0x68,0x8c,0x2b,0x3e,0x6c,0x1f,
    0x1f,0x83,0xd9,0xab,0xfb,0x41,0xbd,0x6b,
    0x5b,0xe0,0xcd,0x19,0x13,0x7e,0x21,0x79
};

typedef struct {
    u8 h[HB];
    u8 m[BB];
    u64 n;
    u64 l;
} ctx;

sv chi(ctx *c)
{
    u64 i;

    FOR(i, HB) c->h[i] = iv[i];
    c->n = c->l = 0;
}

sv chu(ctx *c, const u8 *m, u64 n)
{
    u64 i, t;

    c->l += n;

    while (n) {
        t = MIN(n, BB - c->n);
        FOR(i, t) c->m[c->n + i] = m[i];
        c->n += t;

        crypto_hashblocks(c->h, c->m, c->n);
        c->n &= (BB - 1);
        m += t;
        n -= t;
    }
}

sv chf(u8 *out, ctx *c)
{
    u8 x[256];
    u64 i, n = c->n;

    FOR(i, 256) x[i] = 0;
    FOR(i, n) x[i] = c->m[i];
    x[n] = 128;

    n = 256 - 128 * (n < 112);
    x[n - 9] = c->l >> 61;
    ts64(x + n - 8, c->l << 3);
    crypto_hashblocks(c->h, x, n);

    FOR(i, HB) out[i] = c->h[i];
}

sv prf(u8 *out, const u8 *m, u64 mln, const u8 *k, u64 kln)
{
    ctx c;
    u64 j;
    u8 hk[HB], ik[BB], ok[BB], ih[HB];

    if (kln > BB) {
        chi(&c); chu(&c, k, kln); chf(hk, &c);
        k = hk;
        kln = HB;
    }

    FOR(j, BB) ik[j] = ok[j] = 0;
    FOR(j, kln) ik[j] = ok[j] = k[j];
    FOR(j, BB) ik[j] ^= 0x36;
    FOR(j, BB) ok[j] ^= 0x5c;

    chi(&c); chu(&c, ik, BB); chu(&c, m, mln); chf(ih, &c);
    chi(&c); chu(&c, ok, BB); chu(&c, ih, HB); chf(out, &c);
}

sv hext(u8 *prk, const u8 *ikm, u64 ikmln)
{
    u64 i;
    u8 s[HB];

    FOR(i, HB) s[i] = 0;
    prf(prk, ikm, ikmln, s, HB);
}

sv hexp(u8 *out, u64 outln, const u8* prk, const u64 idx, const u8 *ctx)
{
    u64 l = 0, i = 0, j, bl;
    u8 b[HB + CB + 8 + 1];

    FOR(j, CB) b[HB + j] = ctx[j];
    ts64(b + HB + CB, idx);

    while (l < outln) {
        b[HB + CB + 8] = ++i;

        if (i == 1)
            prf(b, b + HB, CB + 8 + 1, prk, HB);
        else
            prf(b, b, HB + CB + 8 + 1, prk, HB);

        bl = MIN(HB, outln - l);
        FOR(j, bl) out[l + j] = b[j];
        l += bl;
    }
}

int crypto_kdf(u8 *out, u64 outln, const u64 idx, const u8 *ctx, const u8 *ikm, u64 ikmln)
{
    u8 prk[HB];

    if (outln > HB * 255)
        return -1;

    hext(prk, ikm, ikmln);
    hexp(out, outln, prk, idx, ctx);

    return 0;
}

int crypto_pwhash(u8 *out, u64 outln, const u8 *p, u64 pln, const u8 *s, u64 c)
{
    u64 i, j, b, nb, bl, o;
    u8 u[HB], t[HB], bs[SB + 4];

    FOR(i, SB) bs[i] = s[i];
    nb = (outln + HB - 1) / HB;

    for (b = 1, o = 0; b <= nb; b++) {
        bs[SB + 0] = (b >> 24) & 0xff;
        bs[SB + 1] = (b >> 16) & 0xff;
        bs[SB + 2] = (b >> 8) & 0xff;
        bs[SB + 3] = b & 0xff;

        prf(u, bs, SB + 4, p, pln);
        FOR(i, HB) t[i] = u[i];

        FOR(i, c - 1) {
            prf(u, u, HB, p, pln);
            FOR(j, HB) t[j] ^= u[j];
        }

        bl = MIN(HB, outln - o);
        FOR(i, bl) out[o + i] = t[i];
        o += bl;
    }

    return 0;
}