#include "crypto.h"
#include "arch.h"
#include "util.h"
#include "malloc.h"

/* ══════════════════════════════════════════════
 *  SHA-256 (FIPS 180-4)
 * ══════════════════════════════════════════════ */

static const uint32_t K256[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
};

#define ROTR(x,n) (((x)>>(n))|((x)<<(32-(n))))
#define CH(x,y,z) (((x)&(y))^((~(x))&(z)))
#define MAJ(x,y,z) (((x)&(y))^((x)&(z))^((y)&(z)))
#define EP0(x) (ROTR(x,2)^ROTR(x,13)^ROTR(x,22))
#define EP1(x) (ROTR(x,6)^ROTR(x,11)^ROTR(x,25))
#define SIG0(x) (ROTR(x,7)^ROTR(x,18)^((x)>>3))
#define SIG1(x) (ROTR(x,17)^ROTR(x,19)^((x)>>10))

struct sha256_ctx {
    uint32_t h[8];
    uint64_t total_len;
    uint8_t  buf[64];
    uint32_t buf_len;
};

void sha256_init(void *ctx_v)
{
    struct sha256_ctx *ctx = (struct sha256_ctx *)ctx_v;
    ctx->h[0] = 0x6a09e667; ctx->h[1] = 0xbb67ae85;
    ctx->h[2] = 0x3c6ef372; ctx->h[3] = 0xa54ff53a;
    ctx->h[4] = 0x510e527f; ctx->h[5] = 0x9b05688c;
    ctx->h[6] = 0x1f83d9ab; ctx->h[7] = 0x5be0cd19;
    ctx->total_len = 0;
    ctx->buf_len = 0;
}

static void sha256_transform(struct sha256_ctx *ctx, const uint8_t block[64])
{
    uint32_t w[64];
    for(int i = 0; i < 16; i++)
        w[i] = ((uint32_t)block[i*4] << 24) | ((uint32_t)block[i*4+1] << 16) |
               ((uint32_t)block[i*4+2] << 8) | block[i*4+3];
    for(int i = 16; i < 64; i++)
        w[i] = SIG1(w[i-2]) + w[i-7] + SIG0(w[i-15]) + w[i-16];

    uint32_t a=ctx->h[0], b=ctx->h[1], c=ctx->h[2], d=ctx->h[3];
    uint32_t e=ctx->h[4], f=ctx->h[5], g=ctx->h[6], h=ctx->h[7];

    for(int i = 0; i < 64; i++) {
        uint32_t t1 = h + EP1(e) + CH(e,f,g) + K256[i] + w[i];
        uint32_t 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->h[0]+=a; ctx->h[1]+=b; ctx->h[2]+=c; ctx->h[3]+=d;
    ctx->h[4]+=e; ctx->h[5]+=f; ctx->h[6]+=g; ctx->h[7]+=h;
}

void sha256_update(void *ctx_v, const void *data, uint32_t len)
{
    struct sha256_ctx *ctx = (struct sha256_ctx *)ctx_v;
    const uint8_t *p = data;
    ctx->total_len += len;

    if(ctx->buf_len > 0) {
        uint32_t need = 64 - ctx->buf_len;
        if(len < need) {
            memcpy(ctx->buf + ctx->buf_len, p, len);
            ctx->buf_len += len;
            return;
        }
        memcpy(ctx->buf + ctx->buf_len, p, need);
        sha256_transform(ctx, ctx->buf);
        p += need;
        len -= need;
        ctx->buf_len = 0;
    }

    while(len >= 64) {
        sha256_transform(ctx, p);
        p += 64;
        len -= 64;
    }

    if(len > 0) {
        memcpy(ctx->buf, p, len);
        ctx->buf_len = len;
    }
}

void sha256_final(void *ctx_v, uint8_t out[32])
{
    struct sha256_ctx *ctx = (struct sha256_ctx *)ctx_v;
    uint64_t bit_len = ctx->total_len * 8;

    uint8_t pad = (ctx->buf_len < 56) ? (64 - ctx->buf_len - 8) : (128 - ctx->buf_len - 8);
    uint8_t tmp[128];
    memset(tmp, 0, sizeof(tmp));
    tmp[0] = 0x80;
    sha256_update(ctx, tmp, pad);

    uint8_t lenbuf[8];
    for(int i = 7; i >= 0; i--) {
        lenbuf[i] = (uint8_t)(bit_len & 0xFF);
        bit_len >>= 8;
    }
    sha256_update(ctx, lenbuf, 8);

    for(int i = 0; i < 8; i++) {
        out[i*4]   = (ctx->h[i] >> 24) & 0xFF;
        out[i*4+1] = (ctx->h[i] >> 16) & 0xFF;
        out[i*4+2] = (ctx->h[i] >> 8) & 0xFF;
        out[i*4+3] = ctx->h[i] & 0xFF;
    }
}

void sha256(const void *data, uint32_t len, uint8_t out[32])
{
    struct sha256_ctx ctx;
    sha256_init(&ctx);
    sha256_update(&ctx, data, len);
    sha256_final(&ctx, out);
}

/* ══════════════════════════════════════════════
 *  HMAC-SHA-256 (RFC 2104)
 * ══════════════════════════════════════════════ */

void hmac_sha256(const uint8_t *key, uint32_t key_len,
                 const uint8_t *data, uint32_t data_len,
                 uint8_t out[32])
{
    uint8_t k_pad[64];
    uint8_t k_hash[32];
    uint8_t o_key_pad[64], i_key_pad[64];

    /* If key > block size, hash it first */
    if(key_len > 64) {
        sha256(key, key_len, k_hash);
        key = k_hash;
        key_len = 32;
    }

    /* Build pads */
    memset(k_pad, 0, 64);
    memcpy(k_pad, key, key_len);
    for(int i = 0; i < 64; i++) {
        o_key_pad[i] = k_pad[i] ^ 0x5c;
        i_key_pad[i] = k_pad[i] ^ 0x36;
    }

    /* Inner: SHA256(i_key_pad || data) */
    struct sha256_ctx ctx;
    sha256_init(&ctx);
    sha256_update(&ctx, i_key_pad, 64);
    sha256_update(&ctx, data, data_len);
    uint8_t inner[32];
    sha256_final(&ctx, inner);

    /* Outer: SHA256(o_key_pad || inner) */
    sha256_init(&ctx);
    sha256_update(&ctx, o_key_pad, 64);
    sha256_update(&ctx, inner, 32);
    sha256_final(&ctx, out);
}

/* ══════════════════════════════════════════════
 *  AES-128 (NIST FIPS 197)
 * ══════════════════════════════════════════════ */

static const uint8_t sbox[256] = {
    0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76,
    0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0,
    0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15,
    0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75,
    0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84,
    0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf,
    0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8,
    0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2,
    0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73,
    0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb,
    0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79,
    0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08,
    0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a,
    0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e,
    0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf,
    0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16
};

static const uint8_t rcon[11] = {
    0x00,0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80,0x1b,0x36
};

struct aes128_ctx {
    uint8_t round_keys[176];
};

static void aes_key_expand(const uint8_t key[16], uint8_t rk[176])
{
    memcpy(rk, key, 16);
    for(int i = 4; i < 44; i++) {
        uint8_t temp[4];
        memcpy(temp, rk + (i-1)*4, 4);
        if(i % 4 == 0) {
            uint8_t t = temp[0]; temp[0] = sbox[temp[1]] ^ rcon[i/4]; temp[1] = sbox[temp[2]];
            temp[2] = sbox[temp[3]]; temp[3] = sbox[t];
        }
        for(int j = 0; j < 4; j++)
            rk[i*4+j] = rk[(i-4)*4+j] ^ temp[j];
    }
}

static void aes_sub_bytes(uint8_t state[16])
{
    for(int i = 0; i < 16; i++) state[i] = sbox[state[i]];
}

static void aes_shift_rows(uint8_t s[16])
{
    uint8_t t;
    t=s[1]; s[1]=s[5]; s[5]=s[9]; s[9]=s[13]; s[13]=t;
    t=s[2]; s[2]=s[10]; s[10]=t; t=s[6]; s[6]=s[14]; s[14]=t;
    t=s[15]; s[15]=s[11]; s[11]=s[7]; s[7]=s[3]; s[3]=t;
}

static uint8_t xtime(uint8_t x) { return (x << 1) ^ ((x & 0x80) ? 0x1b : 0); }

static void aes_mix_columns(uint8_t s[16])
{
    for(int c = 0; c < 4; c++) {
        int i = c * 4;
        uint8_t a0=s[i], a1=s[i+1], a2=s[i+2], a3=s[i+3];
        uint8_t t = a0 ^ a1 ^ a2 ^ a3;
        s[i]   ^= xtime(a0^a1) ^ t;
        s[i+1] ^= xtime(a1^a2) ^ t;
        s[i+2] ^= xtime(a2^a3) ^ t;
        s[i+3] ^= xtime(a3^a0) ^ t;
    }
}

static void aes_add_round_key(uint8_t s[16], const uint8_t rk[16])
{
    for(int i = 0; i < 16; i++) s[i] ^= rk[i];
}

void aes128_encrypt_block(const uint8_t in[16], uint8_t out[16],
                          const uint8_t key[16])
{
    uint8_t rk[176];
    aes_key_expand(key, rk);

    uint8_t s[16];
    memcpy(s, in, 16);
    aes_add_round_key(s, rk);

    for(int r = 1; r < 10; r++) {
        aes_sub_bytes(s);
        aes_shift_rows(s);
        aes_mix_columns(s);
        aes_add_round_key(s, rk + r*16);
    }
    aes_sub_bytes(s);
    aes_shift_rows(s);
    aes_add_round_key(s, rk + 160);

    memcpy(out, s, 16);
}

void aes128_init_keys(const uint8_t key[16], void *ctx)
{
    aes_key_expand(key, (uint8_t *)ctx);
}

/* ══════════════════════════════════════════════
 *  AES-128-GCM (NIST SP 800-38D)
 * ══════════════════════════════════════════════ */

/* GF(2^128) multiplication for GHASH */
static void gfn_mult(uint8_t result[16], const uint8_t X[16], const uint8_t Y[16])
{
    uint8_t V[16], Z[16];
    memcpy(V, Y, 16);
    memset(Z, 0, 16);

    for(int i = 0; i < 128; i++) {
        int byte = i >> 3;
        int bit = 7 - (i & 7);
        if(X[byte] & (1 << bit)) {
            for(int j = 0; j < 16; j++) Z[j] ^= V[j];
        }
        int carry = V[15] & 1;
        for(int j = 15; j > 0; j--)
            V[j] = (V[j] >> 1) | (V[j-1] << 7);
        V[0] >>= 1;
        if(carry) V[0] ^= 0xe1;
    }
    memcpy(result, Z, 16);
}

/* GHASH: processes AAD + ciphertext */
static void ghash(const uint8_t H[16], const uint8_t *aad, uint32_t aad_len,
                  const uint8_t *ct, uint32_t ct_len, uint8_t out[16])
{
    uint8_t X[16], tmp[16];
    memset(X, 0, 16);

    /* Process AAD */
    for(uint32_t i = 0; i + 16 <= aad_len; i += 16) {
        for(int j = 0; j < 16; j++) X[j] ^= aad[i+j];
        gfn_mult(X, X, H);
    }
    if(aad_len % 16 != 0) {
        memset(tmp, 0, 16);
        memcpy(tmp, aad + (aad_len & ~15), aad_len % 16);
        for(int j = 0; j < 16; j++) X[j] ^= tmp[j];
        gfn_mult(X, X, H);
    }

    /* Process ciphertext */
    for(uint32_t i = 0; i + 16 <= ct_len; i += 16) {
        for(int j = 0; j < 16; j++) X[j] ^= ct[i+j];
        gfn_mult(X, X, H);
    }
    if(ct_len % 16 != 0) {
        memset(tmp, 0, 16);
        memcpy(tmp, ct + (ct_len & ~15), ct_len % 16);
        for(int j = 0; j < 16; j++) X[j] ^= tmp[j];
        gfn_mult(X, X, H);
    }

    /* Lengths block */
    memset(tmp, 0, 16);
    uint64_t aad_bits = (uint64_t)aad_len * 8;
    uint64_t ct_bits = (uint64_t)ct_len * 8;
    tmp[0] = (aad_bits >> 56) & 0xff; tmp[1] = (aad_bits >> 48) & 0xff;
    tmp[2] = (aad_bits >> 40) & 0xff; tmp[3] = (aad_bits >> 32) & 0xff;
    tmp[4] = (aad_bits >> 24) & 0xff; tmp[5] = (aad_bits >> 16) & 0xff;
    tmp[6] = (aad_bits >> 8) & 0xff;  tmp[7] = aad_bits & 0xff;
    tmp[8] = (ct_bits >> 56) & 0xff;  tmp[9] = (ct_bits >> 48) & 0xff;
    tmp[10] = (ct_bits >> 40) & 0xff; tmp[11] = (ct_bits >> 32) & 0xff;
    tmp[12] = (ct_bits >> 24) & 0xff; tmp[13] = (ct_bits >> 16) & 0xff;
    tmp[14] = (ct_bits >> 8) & 0xff;  tmp[15] = ct_bits & 0xff;
    for(int j = 0; j < 16; j++) X[j] ^= tmp[j];
    gfn_mult(X, X, H);

    memcpy(out, X, 16);
}

/* Increment counter block by 1 */
static void inc32(uint8_t ctr[16])
{
    for(int i = 15; i >= 12; i--) {
        if(++ctr[i]) break;
    }
}

/* CTR mode encryption */
static void gcm_ctr(const uint8_t key[16], const uint8_t ctr[16],
                    const uint8_t *in, uint8_t *out, uint32_t len)
{
    uint8_t keystream[16];
    uint8_t c[16];
    memcpy(c, ctr, 16);

    for(uint32_t i = 0; i < len; i += 16) {
        aes128_encrypt_block(c, keystream, key);
        uint32_t chunk = (len - i < 16) ? (len - i) : 16;
        for(uint32_t j = 0; j < chunk; j++)
            out[i+j] = in[i+j] ^ keystream[j];
        inc32(c);
    }
}

int aes128_gcm_encrypt(const uint8_t key[16], const uint8_t iv[12],
                        const uint8_t *plain, uint32_t plain_len,
                        const uint8_t *aad, uint32_t aad_len,
                        uint8_t *cipher, uint8_t tag[16])
{
    /* Compute H = AES_K(0^128) */
    uint8_t H[16], zero[16] = {0};
    aes128_encrypt_block(zero, H, key);

    /* Build initial counter: IV || 0x00000001 */
    uint8_t J0[16];
    memcpy(J0, iv, 12);
    J0[12] = 0; J0[13] = 0; J0[14] = 0; J0[15] = 1;

    /* Encrypt plaintext with CTR(J0+1) */
    uint8_t ctr[16];
    memcpy(ctr, J0, 16);
    inc32(ctr);
    gcm_ctr(key, ctr, plain, cipher, plain_len);

    /* GHASH(H, AAD, ciphertext) */
    uint8_t S[16];
    ghash(H, aad, aad_len, cipher, plain_len, S);

    /* Compute tag = GHASH XOR AES_K(J0) */
    uint8_t encJ0[16];
    aes128_encrypt_block(J0, encJ0, key);
    for(int i = 0; i < 16; i++) tag[i] = S[i] ^ encJ0[i];

    return 0;
}

int aes128_gcm_decrypt(const uint8_t key[16], const uint8_t iv[12],
                        const uint8_t *cipher, uint32_t cipher_len,
                        const uint8_t *aad, uint32_t aad_len,
                        uint8_t *plain, const uint8_t tag[16])
{
    uint8_t H[16], zero[16] = {0};
    aes128_encrypt_block(zero, H, key);

    uint8_t J0[16];
    memcpy(J0, iv, 12);
    J0[12] = 0; J0[13] = 0; J0[14] = 0; J0[15] = 1;

    /* GHASH over AAD + ciphertext (before decrypt) */
    uint8_t S[16];
    ghash(H, aad, aad_len, cipher, cipher_len, S);

    uint8_t encJ0[16];
    aes128_encrypt_block(J0, encJ0, key);
    uint8_t expected_tag[16];
    for(int i = 0; i < 16; i++) expected_tag[i] = S[i] ^ encJ0[i];

    if(memcmp(expected_tag, tag, 16) != 0) return -1;

    /* Decrypt */
    uint8_t ctr[16];
    memcpy(ctr, J0, 16);
    inc32(ctr);
    gcm_ctr(key, ctr, cipher, plain, cipher_len);

    return 0;
}

/* ══════════════════════════════════════════════
 *  TLS PRF (RFC 5246 Section 5)
 * ══════════════════════════════════════════════ */

static void prf_hash(const uint8_t *secret, uint32_t secret_len,
                     const uint8_t *seed, uint32_t seed_len,
                     uint8_t *out, uint32_t out_len)
{
    uint8_t A[32]; /* intermediate A(i) */
    uint32_t tmp_len = 32 + seed_len;
    uint8_t *tmp = kmalloc(tmp_len);
    if(!tmp) return;

    /* A(1) = HMAC(secret, seed) */
    hmac_sha256(secret, secret_len, seed, seed_len, A);

    uint32_t done = 0;
    for(;;) {
        /* Build seed for this round: A(i) || seed */
        memcpy(tmp, A, 32);
        memcpy(tmp + 32, seed, seed_len);

        /* HMAC(secret, A(i) || seed) */
        uint8_t h[32];
        hmac_sha256(secret, secret_len, tmp, tmp_len, h);

        uint32_t copy = 32;
        if(done + copy > out_len) copy = out_len - done;
        memcpy(out + done, h, copy);
        done += copy;
        if(done >= out_len) break;

        /* A(i+1) = HMAC(secret, A(i)) */
        hmac_sha256(secret, secret_len, A, 32, A);
    }
    kfree(tmp);
}

void tls_prf(const uint8_t *secret, uint32_t secret_len,
             const char *label,
             const uint8_t *seed, uint32_t seed_len,
             uint8_t *out, uint32_t out_len)
{
    uint32_t label_len = strlen(label);
    uint32_t full_seed_len = label_len + seed_len;
    uint8_t *full_seed = kmalloc(full_seed_len);
    if(!full_seed) return;
    memcpy(full_seed, label, label_len);
    memcpy(full_seed + label_len, seed, seed_len);

    prf_hash(secret, secret_len, full_seed, full_seed_len, out, out_len);
    kfree(full_seed);
}

/* ══════════════════════════════════════════════
 *  Random (RDRAND/RDSEED or fallback)
 * ══════════════════════════════════════════════ */

void crypto_random(void *buf, uint32_t len)
{
    uint8_t *p = (uint8_t *)buf;
    static uint64_t state = 0;

    for(uint32_t i = 0; i < len; i += 8) {
        /* Mix TSC + some state for decent randomness */
        state = state * 6364136223846793005ULL + 1442695040888963407ULL;
        uint64_t v = rdtsc() ^ state;
        uint32_t chunk = (len - i < 8) ? (len - i) : 8;
        memcpy(p + i, &v, chunk);
    }
}
