#include "naut/mse.h"
#include "naut/hash.h"

#include <errno.h>
#include <openssl/bn.h>
#include <openssl/rand.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>

#define MSE_PAD_MAX 512
#define MSE_CRYPTO_RC4 2u

static const char DH_PRIME_HEX[] =
    "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC"
    "74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF2"
    "5F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A63A3621000000"
    "0000090563";

static void wr16(uint8_t *p, uint16_t value) {
    p[0] = (uint8_t)(value >> 8);
    p[1] = (uint8_t)value;
}

static void wr32(uint8_t *p, uint32_t value) {
    p[0] = (uint8_t)(value >> 24);
    p[1] = (uint8_t)(value >> 16);
    p[2] = (uint8_t)(value >> 8);
    p[3] = (uint8_t)value;
}

static uint16_t rd16(const uint8_t *p) {
    return ((uint16_t)p[0] << 8) | p[1];
}

static uint32_t rd32(const uint8_t *p) {
    return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) |
           ((uint32_t)p[2] << 8) | p[3];
}

static void hash_parts(const char label[4],
                       const uint8_t *first, size_t first_len,
                       const uint8_t *second, size_t second_len,
                       uint8_t out[20]) {
    naut_sha1_ctx sha;
    naut_sha1_init(&sha);
    naut_sha1_update(&sha, label, 4);
    naut_sha1_update(&sha, first, first_len);
    if (second && second_len) naut_sha1_update(&sha, second, second_len);
    naut_sha1_final(&sha, out);
}

static void init_rc4(const uint8_t secret[NAUT_MSE_DH_LEN],
                     const uint8_t info_hash[20],
                     naut_mse_stream *stream) {
    uint8_t key_a[20], key_b[20];
    hash_parts("keyA", secret, NAUT_MSE_DH_LEN, info_hash, 20, key_a);
    hash_parts("keyB", secret, NAUT_MSE_DH_LEN, info_hash, 20, key_b);
    naut_rc4_init(&stream->send, key_a, sizeof key_a, 1024);
    naut_rc4_init(&stream->recv, key_b, sizeof key_b, 1024);
    memset(key_a, 0, sizeof key_a);
    memset(key_b, 0, sizeof key_b);
}

/* ---- sans-IO handshake state machine ------------------------------------- */

enum {
    PH_RECV_PUBKEY,   /* waiting for the peer's 96-byte DH public key */
    PH_SYNC_VC,       /* scanning past PadB for the encrypted VC      */
    PH_RECV_SELECT,   /* crypto_select + len(PadD)                    */
    PH_RECV_PAD,      /* PadD bytes (discarded)                       */
    PH_RECV_HS,       /* the peer's encrypted BitTorrent handshake    */
};

struct naut_mse_handshake {
    int      phase;
    naut_err err;
    bool     done;

    uint8_t  info_hash[20];
    uint8_t  peer_id[NAUT_PEERID_LEN];
    uint64_t reserved;

    /* DH state retained until the shared secret is computed. */
    BN_CTX  *ctx;
    BIGNUM  *prime;
    BIGNUM  *priv;

    naut_mse_stream stream;
    uint8_t  expected_vc[8];
    size_t   vc_scanned;
    size_t   pad_remaining;
    uint8_t  remote_handshake[NAUT_HANDSHAKE_LEN];

    uint8_t  out[256];
    size_t   out_len, out_off;

    uint8_t  in[1024];
    size_t   in_len;
};

static void dh_free(naut_mse_handshake *h) {
    BN_CTX_free(h->ctx);   h->ctx = NULL;
    BN_free(h->prime);     h->prime = NULL;
    BN_clear_free(h->priv); h->priv = NULL;
}

/* Generate our private key and public value, writing the 96-byte public key
 * into the outgoing buffer. Retains prime/priv/ctx for dh_complete(). */
static naut_err dh_begin(naut_mse_handshake *h) {
    naut_err result = NAUT_ERR_IO;
    BIGNUM *generator = BN_new();
    BIGNUM *local = BN_new();
    h->ctx = BN_CTX_new();
    h->priv = BN_new();
    if (!generator || !local || !h->ctx || !h->priv ||
        !BN_hex2bn(&h->prime, DH_PRIME_HEX) || !BN_set_word(generator, 2))
        goto done;
    do {
        if (!BN_rand_range(h->priv, h->prime)) goto done;
    } while (BN_cmp(h->priv, generator) < 0);
    if (!BN_mod_exp(local, generator, h->priv, h->prime, h->ctx) ||
        BN_bn2binpad(local, h->out, NAUT_MSE_DH_LEN) != NAUT_MSE_DH_LEN)
        goto done;
    h->out_len = NAUT_MSE_DH_LEN;
    h->out_off = 0;
    result = NAUT_OK;
done:
    BN_free(generator);
    BN_free(local);
    if (result != NAUT_OK) dh_free(h);
    return result;
}

/* Validate the peer's public key and derive the shared secret. */
static naut_err dh_complete(naut_mse_handshake *h, const uint8_t remote_bytes[96],
                            uint8_t secret[NAUT_MSE_DH_LEN]) {
    naut_err result = NAUT_ERR_IO;
    BIGNUM *remote = BN_new();
    BIGNUM *shared = BN_new();
    BIGNUM *limit = BN_new();
    BIGNUM *two = BN_new();
    if (!remote || !shared || !limit || !two ||
        !BN_bin2bn(remote_bytes, NAUT_MSE_DH_LEN, remote) ||
        !BN_set_word(two, 2) || !BN_copy(limit, h->prime) ||
        !BN_sub_word(limit, 1))
        goto done;
    if (BN_cmp(remote, two) < 0 || BN_cmp(remote, limit) >= 0) {
        result = NAUT_ERR_PROTO;
        goto done;
    }
    if (!BN_mod_exp(shared, remote, h->priv, h->prime, h->ctx) ||
        BN_bn2binpad(shared, secret, NAUT_MSE_DH_LEN) != NAUT_MSE_DH_LEN)
        goto done;
    result = NAUT_OK;
done:
    BN_free(remote);
    BN_clear_free(shared);
    BN_free(limit);
    BN_free(two);
    return result;
}

naut_mse_handshake *naut_mse_handshake_begin(
    const uint8_t info_hash[20],
    const uint8_t peer_id[NAUT_PEERID_LEN],
    uint64_t reserved) {
    if (!info_hash || !peer_id) return NULL;
    naut_mse_handshake *h = calloc(1, sizeof(*h));
    if (!h) return NULL;
    memcpy(h->info_hash, info_hash, 20);
    memcpy(h->peer_id, peer_id, NAUT_PEERID_LEN);
    h->reserved = reserved;
    h->phase = PH_RECV_PUBKEY;
    if (dh_begin(h) != NAUT_OK) {
        naut_mse_handshake_free(h);
        return NULL;
    }
    return h;
}

void naut_mse_handshake_free(naut_mse_handshake *h) {
    if (!h) return;
    dh_free(h);
    /* keystream state is sensitive; scrub before release */
    memset(h, 0, sizeof(*h));
    free(h);
}

static void consume(naut_mse_handshake *h, size_t n) {
    memmove(h->in, h->in + n, h->in_len - n);
    h->in_len -= n;
}

/* Build req1/req2 + encrypted offer (VC, crypto_provide, PadC, IA) into out. */
static void build_request(naut_mse_handshake *h, const uint8_t secret[96]) {
    uint8_t req1[20], req2[20], req3[20];
    hash_parts("req1", secret, NAUT_MSE_DH_LEN, NULL, 0, req1);
    hash_parts("req2", h->info_hash, 20, NULL, 0, req2);
    hash_parts("req3", secret, NAUT_MSE_DH_LEN, NULL, 0, req3);
    for (size_t i = 0; i < sizeof req2; i++) req2[i] ^= req3[i];

    init_rc4(secret, h->info_hash, &h->stream);

    uint8_t *p = h->out;
    memcpy(p, req1, 20);
    memcpy(p + 20, req2, 20);
    p += 40;

    uint8_t *offer = p;            /* VC(8) crypto_provide(4) padlen(2) ialen(2) IA */
    memset(offer, 0, 8);
    wr32(offer + 8, MSE_CRYPTO_RC4);
    wr16(offer + 12, 0);
    wr16(offer + 14, NAUT_HANDSHAKE_LEN);
    naut_peer_handshake_build(offer + 16, h->info_hash, h->peer_id, h->reserved);
    size_t offer_len = 16 + NAUT_HANDSHAKE_LEN;
    naut_rc4_xor(&h->stream.send, offer, offer_len);

    h->out_len = 40 + offer_len;
    h->out_off = 0;

    /* expected_vc = our recv keystream applied to 8 zero bytes at position 0,
     * without advancing the real recv state (we resync on it). */
    naut_rc4 probe = h->stream.recv;
    uint8_t vc[8] = {0};
    naut_rc4_xor(&probe, vc, sizeof vc);
    memcpy(h->expected_vc, vc, sizeof vc);
    h->vc_scanned = 0;
}

static void advance(naut_mse_handshake *h) {
    for (;;) {
        switch (h->phase) {
        case PH_RECV_PUBKEY: {
            if (h->in_len < NAUT_MSE_DH_LEN) return;
            uint8_t secret[NAUT_MSE_DH_LEN];
            naut_err e = dh_complete(h, h->in, secret);
            if (e != NAUT_OK) { h->err = e; return; }
            consume(h, NAUT_MSE_DH_LEN);
            dh_free(h);                 /* DH no longer needed */
            build_request(h, secret);
            memset(secret, 0, sizeof secret);
            h->phase = PH_SYNC_VC;
            return;                     /* out now holds req+offer: NEED_WRITE */
        }
        case PH_SYNC_VC: {
            while (h->in_len >= sizeof h->expected_vc) {
                if (memcmp(h->in, h->expected_vc, sizeof h->expected_vc) == 0) {
                    uint8_t vc[8];
                    memcpy(vc, h->in, sizeof vc);
                    naut_rc4_xor(&h->stream.recv, vc, sizeof vc);
                    static const uint8_t zero8[8] = {0};
                    if (memcmp(vc, zero8, sizeof vc) != 0) {
                        h->err = NAUT_ERR_PROTO;
                        return;
                    }
                    consume(h, sizeof vc);
                    h->phase = PH_RECV_SELECT;
                    break;
                }
                consume(h, 1);
                if (++h->vc_scanned > MSE_PAD_MAX) {
                    h->err = NAUT_ERR_PROTO;
                    return;
                }
            }
            if (h->phase == PH_SYNC_VC) return;   /* need more bytes */
            continue;
        }
        case PH_RECV_SELECT: {
            if (h->in_len < 6) return;
            uint8_t hdr[6];
            memcpy(hdr, h->in, sizeof hdr);
            naut_rc4_xor(&h->stream.recv, hdr, sizeof hdr);
            consume(h, sizeof hdr);
            if (rd32(hdr) != MSE_CRYPTO_RC4) { h->err = NAUT_ERR_PROTO; return; }
            h->pad_remaining = rd16(hdr + 4);
            if (h->pad_remaining > MSE_PAD_MAX) { h->err = NAUT_ERR_PROTO; return; }
            h->phase = PH_RECV_PAD;
            continue;
        }
        case PH_RECV_PAD: {
            if (h->pad_remaining > 0) {
                size_t n = h->pad_remaining < h->in_len ? h->pad_remaining
                                                        : h->in_len;
                if (n == 0) return;
                naut_rc4_xor(&h->stream.recv, h->in, n);  /* advance keystream */
                consume(h, n);
                h->pad_remaining -= n;
                if (h->pad_remaining > 0) return;
            }
            h->phase = PH_RECV_HS;
            continue;
        }
        case PH_RECV_HS: {
            if (h->in_len < NAUT_HANDSHAKE_LEN) return;
            memcpy(h->remote_handshake, h->in, NAUT_HANDSHAKE_LEN);
            naut_rc4_xor(&h->stream.recv, h->remote_handshake, NAUT_HANDSHAKE_LEN);
            consume(h, NAUT_HANDSHAKE_LEN);
            uint8_t remote_hash[20], remote_id[20];
            if (!naut_peer_handshake_parse(h->remote_handshake, remote_hash,
                                           remote_id, NULL) ||
                memcmp(remote_hash, h->info_hash, 20) != 0) {
                h->err = NAUT_ERR_PROTO;
                return;
            }
            h->stream.active = true;
            h->done = true;
            return;
        }
        default:
            h->err = NAUT_ERR_PROTO;
            return;
        }
    }
}

naut_mse_hs_status naut_mse_handshake_status(const naut_mse_handshake *h) {
    if (!h || h->err != NAUT_OK) return NAUT_MSE_HS_ERROR;
    if (h->done) return NAUT_MSE_HS_DONE;
    if (h->out_off < h->out_len) return NAUT_MSE_HS_NEED_WRITE;
    return NAUT_MSE_HS_NEED_READ;
}

size_t naut_mse_handshake_pull(naut_mse_handshake *h, uint8_t *buf, size_t cap) {
    if (!h || !buf) return 0;
    size_t avail = h->out_len - h->out_off;
    size_t n = avail < cap ? avail : cap;
    if (n) {
        memcpy(buf, h->out + h->out_off, n);
        h->out_off += n;
        if (h->out_off == h->out_len) h->out_len = h->out_off = 0;
    }
    return n;
}

naut_mse_hs_status naut_mse_handshake_feed(naut_mse_handshake *h,
                                           const uint8_t *data, size_t len,
                                           size_t *consumed) {
    if (consumed) *consumed = 0;
    if (!h) return NAUT_MSE_HS_ERROR;
    if (h->err == NAUT_OK && !h->done && data && len) {
        size_t space = sizeof h->in - h->in_len;
        size_t take = len < space ? len : space;
        memcpy(h->in + h->in_len, data, take);
        h->in_len += take;
        if (consumed) *consumed = take;
        advance(h);
    }
    return naut_mse_handshake_status(h);
}

naut_err naut_mse_handshake_finish(naut_mse_handshake *h,
                                   naut_mse_stream *stream,
                                   uint8_t remote_handshake[NAUT_HANDSHAKE_LEN]) {
    if (!h || !stream || !remote_handshake) return NAUT_ERR_INVAL;
    if (h->err != NAUT_OK) return h->err;
    if (!h->done) return NAUT_ERR_AGAIN;
    *stream = h->stream;
    memcpy(remote_handshake, h->remote_handshake, NAUT_HANDSHAKE_LEN);
    return NAUT_OK;
}

/* ---- blocking I/O helpers + convenience wrapper -------------------------- */

static bool raw_send_all(int fd, const void *data, size_t len) {
    const uint8_t *p = data;
    while (len) {
        ssize_t n = send(fd, p, len, MSG_NOSIGNAL);
        if (n < 0) {
            if (errno == EINTR) continue;
            return false;
        }
        if (n == 0) return false;
        p += n;
        len -= (size_t)n;
    }
    return true;
}

static bool raw_recv_exact(int fd, void *data, size_t len) {
    uint8_t *p = data;
    while (len) {
        ssize_t n = recv(fd, p, len, 0);
        if (n < 0) {
            if (errno == EINTR) continue;
            return false;
        }
        if (n == 0) return false;
        p += n;
        len -= (size_t)n;
    }
    return true;
}

naut_err naut_mse_client_handshake(
    int fd,
    const uint8_t info_hash[20],
    const uint8_t peer_id[NAUT_PEERID_LEN],
    uint64_t reserved,
    naut_mse_stream *stream,
    uint8_t remote_handshake[NAUT_HANDSHAKE_LEN]) {
    if (fd < 0 || !info_hash || !peer_id || !stream || !remote_handshake)
        return NAUT_ERR_INVAL;
    memset(stream, 0, sizeof(*stream));

    naut_mse_handshake *h =
        naut_mse_handshake_begin(info_hash, peer_id, reserved);
    if (!h) return NAUT_ERR_NOMEM;

    naut_err rc = NAUT_ERR_PROTO;
    for (;;) {
        naut_mse_hs_status st = naut_mse_handshake_status(h);
        if (st == NAUT_MSE_HS_NEED_WRITE) {
            uint8_t buf[256];
            size_t n;
            bool ok = true;
            while ((n = naut_mse_handshake_pull(h, buf, sizeof buf)) > 0)
                if (!raw_send_all(fd, buf, n)) { ok = false; break; }
            if (!ok) { rc = NAUT_ERR_IO; break; }
        } else if (st == NAUT_MSE_HS_NEED_READ) {
            /* One byte at a time: the handshake is tiny and one-shot, and this
             * keeps the wrapper from over-reading into the payload stream. */
            uint8_t byte;
            if (!raw_recv_exact(fd, &byte, 1)) { rc = NAUT_ERR_IO; break; }
            naut_mse_handshake_feed(h, &byte, 1, NULL);
        } else if (st == NAUT_MSE_HS_DONE) {
            rc = naut_mse_handshake_finish(h, stream, remote_handshake);
            break;
        } else {
            rc = h->err != NAUT_OK ? h->err : NAUT_ERR_PROTO;
            break;
        }
    }
    naut_mse_handshake_free(h);
    return rc;
}

/* ---- post-handshake stream I/O ------------------------------------------- */

bool naut_mse_send_all(int fd, naut_mse_stream *stream,
                       const void *data, size_t len) {
    if (!stream || !stream->active) return raw_send_all(fd, data, len);
    const uint8_t *p = data;
    uint8_t block[16 * 1024];
    while (len) {
        size_t n = len < sizeof block ? len : sizeof block;
        memcpy(block, p, n);
        naut_rc4_xor(&stream->send, block, n);
        if (!raw_send_all(fd, block, n)) return false;
        p += n;
        len -= n;
    }
    return true;
}

ssize_t naut_mse_recv(int fd, naut_mse_stream *stream,
                      void *data, size_t len) {
    ssize_t n;
    do {
        n = recv(fd, data, len, 0);
    } while (n < 0 && errno == EINTR);
    if (n > 0 && stream && stream->active)
        naut_rc4_xor(&stream->recv, data, (size_t)n);
    return n;
}