/*
 * proto.c - BitTorrent peer wire protocol: handshake, message framing, and the
 * incremental push-model parser.
 *
 * The loop thread does the recv() and feeds the bytes here via proto_feed();
 * the parser is a resumable state machine that never blocks and never calls
 * recv() itself. A piece's payload is memcpy'd out of the fed buffer straight
 * into its destination arena slot, so the only copy on the hot path is that one
 * (unavoidable, since the bytes already live in the loop's recv staging buffer).
 *
 * Completed blocks are pushed to the owning loop's ready_ring and the engine's
 * shared ready eventfd is signalled so the single consumer wakes.
 */
#include "engine_internal.h"

#include <stdio.h>
#include <string.h>
#include <unistd.h>

/* pstrlen (19) followed by "BitTorrent protocol"; exactly 20 bytes, no NUL. */
static const uint8_t BT_PROTOCOL[20] = {
    19, 'B','i','t','T','o','r','r','e','n','t',' ','p','r','o','t','o','c','o','l'
};

static inline uint32_t be32(const uint8_t *p) {
    return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) |
           ((uint32_t)p[2] << 8) | (uint32_t)p[3];
}
static inline void put_be32(uint8_t *p, uint32_t v) {
    p[0] = (uint8_t)(v >> 24); p[1] = (uint8_t)(v >> 16);
    p[2] = (uint8_t)(v >> 8);  p[3] = (uint8_t)v;
}

/* ---- outgoing message construction --------------------------------- */

void proto_queue_handshake(conn *c) {
    uint8_t hs[HANDSHAKE_LEN];
    memcpy(hs, BT_PROTOCOL, 20);
    memset(hs + 20, 0, 8);                       /* reserved                  */
    hs[25] |= 0x10;                              /* BEP-10 extension protocol */
    hs[27] |= 0x04;                              /* BEP-6 Fast Extension      */
    memcpy(hs + 28, c->tor->info_hash, 20);
    memcpy(hs + 48, c->tor->peer_id, 20);
    outbuf_append(&c->out, hs, sizeof hs);
}

void proto_queue_msg(conn *c, uint8_t id) {
    uint8_t m[5];
    put_be32(m, 1);
    m[4] = id;
    outbuf_append(&c->out, m, sizeof m);
}

void proto_queue_request(conn *c, uint32_t piece, uint32_t begin, uint32_t length) {
    uint8_t m[17];
    put_be32(m, 13);
    m[4] = MSG_REQUEST;
    put_be32(m + 5, piece);
    put_be32(m + 9, begin);
    put_be32(m + 13, length);
    outbuf_append(&c->out, m, sizeof m);
}

void proto_queue_cancel(conn *c, uint32_t piece, uint32_t begin, uint32_t length) {
    uint8_t m[17];
    put_be32(m, 13);
    m[4] = MSG_CANCEL;
    put_be32(m + 5, piece);
    put_be32(m + 9, begin);
    put_be32(m + 13, length);
    outbuf_append(&c->out, m, sizeof m);
}

static void proto_queue_have_none(conn *c) {
    proto_queue_msg(c, MSG_HAVE_NONE);
}

static void proto_queue_ext_handshake(conn *c) {
    char payload[160];
    int n = snprintf(payload, sizeof payload,
                     "d1:md11:lt_donthavei%uee1:reqqi%ue1:v14:torrent-peer/0ee",
                     EXT_LT_DONTHAVE, c->lp->eng->cfg.max_pipeline);
    if (n <= 0 || (size_t)n >= sizeof payload) return;

    uint8_t hdr[6];
    put_be32(hdr, (uint32_t)n + 2);
    hdr[4] = MSG_EXTENDED;
    hdr[5] = 0;                                  /* extended handshake */
    outbuf_append(&c->out, hdr, sizeof hdr);
    outbuf_append(&c->out, payload, (size_t)n);
}

/* ---- fed-buffer reader --------------------------------------------- */

/* Copy up to (need - *got) bytes from the feed cursor [*p, end) into dst.
 * Advances the cursor. Returns 1 once *got reaches need, else 0 (need more). */
static inline int feed_take(const uint8_t **p, const uint8_t *end, void *dst,
                            size_t need, size_t *got) {
    size_t avail = (size_t)(end - *p);
    size_t n = need - *got;
    if (n > avail) n = avail;
    memcpy((uint8_t *)dst + *got, *p, n);
    *p += n;
    *got += n;
    return *got >= need;
}

static void signal_ready(conn *c) {
    uint64_t one = 1;
    ssize_t w = write(c->lp->eng->ready_efd, &one, sizeof one);
    (void)w; /* eventfd write only fails on overflow; harmless */
}

/* Endgame races a block across several peers; once one copy lands, cancel the
 * others so we don't pay to receive the same block twice. No-op outside endgame
 * (no two connections share an in-flight block then), so the hot path is clean. */
static void cancel_redundant_copies(conn *src, uint32_t piece, uint32_t begin,
                                    uint32_t len) {
    torrent *t = src->tor;
    if (!t->endgame) return;
    uint64_t key = block_key(t, piece, begin);
    for (conn *o = src->lp->conns; o; o = o->next) {
        if (o == src || o->tor != t || o->dead) continue;
        req_entry e;
        if (reqtab_take(&o->inflight, key, &e)) {
            proto_queue_cancel(o, piece, begin, len);
            atomic_fetch_sub_explicit(&o->outstanding, 1, memory_order_relaxed);
            o->lp->outstanding--;
        }
    }
}

static void drop_inflight_piece(conn *c, uint32_t piece) {
    uint32_t removed = reqtab_take_piece(&c->inflight, piece);
    if (removed == 0) return;
    atomic_fetch_sub_explicit(&c->outstanding, removed, memory_order_relaxed);
    c->lp->outstanding -= removed;
    c->tor->requested[piece] = 0;
    if (c->have_cur_piece && c->cur_piece == piece) c->have_cur_piece = 0;
}

static void handle_reject(conn *c, const uint8_t *payload, size_t len) {
    if (!c->fast_enabled || len < 12) {
        conn_set_error(c, PEER_ERR_PROTOCOL);
        return;
    }

    uint32_t piece = be32(payload);
    uint32_t begin = be32(payload + 4);
    if (piece >= c->tor->num_pieces) {
        conn_set_error(c, PEER_ERR_PROTOCOL);
        return;
    }

    req_entry e;
    if (reqtab_take(&c->inflight, block_key(c->tor, piece, begin), &e)) {
        atomic_fetch_sub_explicit(&c->outstanding, 1, memory_order_relaxed);
        c->lp->outstanding--;
        c->tor->requested[piece] = 0;
        if (c->have_cur_piece && c->cur_piece == piece) c->have_cur_piece = 0;
    }
}

static void handle_extended(conn *c, const uint8_t *payload, size_t len) {
    if (!c->ext_enabled || len < 1) return;

    uint8_t ext_id = payload[0];
    if (ext_id == 0) return;                     /* handshake: no fields needed */
    if (ext_id != EXT_LT_DONTHAVE || len < 5) return;

    uint32_t piece = be32(payload + 1);
    if (piece >= c->tor->num_pieces) return;
    clear_have_bit(c->have_bits, piece);
    clear_have_bit(c->allowed_fast_bits, piece);
    drop_inflight_piece(c, piece);
}

/* A fully-received non-piece control message (payload after the id). have_bits
 * is loop-thread-only, so the selector reads it from the same thread. */
static void handle_control(conn *c, uint8_t id, const uint8_t *payload, size_t len) {
    switch (id) {
    case MSG_UNCHOKE:
        c->unchoked = 1;
        atomic_store_explicit(&c->astate, PEER_STATE_RUNNING, memory_order_relaxed);
        break;
    case MSG_CHOKE:
        c->unchoked = 0;
        atomic_store_explicit(&c->astate, PEER_STATE_CHOKED, memory_order_relaxed);
        break;
    case MSG_BITFIELD: {
        size_t n = len < c->tor->bf_bytes ? len : c->tor->bf_bytes;
        memcpy(c->have_bits, payload, n);
        break;
    }
    case MSG_HAVE:
        if (len >= 4) {
            uint32_t idx = be32(payload);
            if (idx < c->tor->num_pieces) set_have_bit(c->have_bits, idx);
        }
        break;
    case MSG_SUGGEST:
        if (!c->fast_enabled || len < 4) conn_set_error(c, PEER_ERR_PROTOCOL);
        break;
    case MSG_HAVE_ALL:
        if (!c->fast_enabled) {
            conn_set_error(c, PEER_ERR_PROTOCOL);
            break;
        }
        set_all_have_bits(c->have_bits, c->tor->num_pieces);
        break;
    case MSG_HAVE_NONE:
        if (!c->fast_enabled) {
            conn_set_error(c, PEER_ERR_PROTOCOL);
            break;
        }
        memset(c->have_bits, 0, c->tor->bf_bytes);
        memset(c->allowed_fast_bits, 0, c->tor->bf_bytes);
        break;
    case MSG_REJECT:
        handle_reject(c, payload, len);
        break;
    case MSG_ALLOWED_FAST:
        if (!c->fast_enabled || len < 4) {
            conn_set_error(c, PEER_ERR_PROTOCOL);
            break;
        }
        {
            uint32_t idx = be32(payload);
            if (idx < c->tor->num_pieces) set_have_bit(c->allowed_fast_bits, idx);
        }
        break;
    case MSG_EXTENDED:
        handle_extended(c, payload, len);
        break;
    default:
        /* not-interested/port/fast-extension/etc: ignored by a leech. */
        break;
    }
}

void conn_set_error(conn *c, peer_error e) {
    atomic_store_explicit(&c->aerror, (int)e, memory_order_relaxed);
    atomic_store_explicit(&c->astate, PEER_STATE_ERROR, memory_order_relaxed);
}

/* ---- the parser ----------------------------------------------------- */

int proto_feed(conn *c, const uint8_t *data, size_t len) {
    reader *r = &c->rd;
    torrent *tor = c->tor;
    loop *lp = c->lp;
    const uint8_t *p = data;
    const uint8_t *end = data + len;
    int pushed = 0;

    while (p < end) {
        switch (r->state) {
        case RS_HANDSHAKE:
            if (!feed_take(&p, end, r->hs, HANDSHAKE_LEN, &r->hs_got)) goto out;
            if (memcmp(r->hs, BT_PROTOCOL, 20) != 0 ||
                memcmp(r->hs + 28, tor->info_hash, 20) != 0) {
                conn_set_error(c, PEER_ERR_HANDSHAKE);
                return -1;
            }
            c->fast_enabled = (r->hs[27] & 0x04) != 0;
            c->ext_enabled = (r->hs[25] & 0x10) != 0;
            if (c->fast_enabled) proto_queue_have_none(c);
            if (c->ext_enabled) proto_queue_ext_handshake(c);
            /* Express interest; wait for unchoke or Allowed Fast. */
            c->bt_established = 1;
            proto_queue_msg(c, MSG_INTERESTED);
            atomic_store_explicit(&c->astate, PEER_STATE_CHOKED, memory_order_relaxed);
            r->state = RS_LEN;
            r->len_got = 0;
            break;

        case RS_LEN:
            if (!feed_take(&p, end, r->lenb, 4, &r->len_got)) goto out;
            r->msg_len = be32(r->lenb);
            r->len_got = 0;
            if (r->msg_len == 0) break; /* keep-alive */
            r->id_got = 0;
            r->state = RS_ID;
            break;

        case RS_ID:
            if (!feed_take(&p, end, &r->msg_id, 1, &r->id_got)) goto out;
            if (r->msg_id == MSG_PIECE) {
                if (r->msg_len < 1 + PIECE_HDR_LEN ||
                    r->msg_len - 1 - PIECE_HDR_LEN > PEER_BLOCK_SIZE) {
                    conn_set_error(c, PEER_ERR_PROTOCOL);
                    return -1;
                }
                r->body_len = r->msg_len - 1 - PIECE_HDR_LEN;
                r->phdr_got = 0;
                r->state = RS_PIECE_HDR;
            } else {
                uint32_t payload = r->msg_len - 1;
                if (payload == 0) {
                    handle_control(c, r->msg_id, NULL, 0);
                    if (atomic_load_explicit(&c->aerror, memory_order_relaxed) != PEER_OK)
                        return -1;
                    r->state = RS_LEN;
                } else {
                    if (payload > r->other_cap) {
                        conn_set_error(c, PEER_ERR_PROTOCOL);
                        return -1;
                    }
                    r->other_total = payload;
                    r->other_got = 0;
                    r->state = RS_OTHER;
                }
            }
            break;

        case RS_PIECE_HDR: {
            if (!feed_take(&p, end, r->phdr, PIECE_HDR_LEN, &r->phdr_got)) goto out;
            r->cur_piece = be32(r->phdr);
            r->cur_begin = be32(r->phdr + 4);
            /* C2: only accept a block we actually requested; otherwise drain and
             * drop it without consuming a slot or disturbing credit. */
            req_entry e;
            if (!reqtab_take(&c->inflight, block_key(tor, r->cur_piece, r->cur_begin), &e)) {
                r->other_total = r->body_len; /* <= block size <= other_cap */
                r->other_got = 0;
                r->state = RS_PIECE_DROP;
                break;
            }
            uint64_t rtt = peer_now_ns() - e.issue_ns;
            if (c->rtt_min_ns == 0 || rtt < c->rtt_min_ns) c->rtt_min_ns = rtt;
            atomic_fetch_sub_explicit(&c->outstanding, 1, memory_order_relaxed);
            lp->outstanding--;

            if (!slot_ring_pop(&lp->free_ring, &r->cur_slot)) {
                /* Flow control guarantees a free slot here. */
                conn_set_error(c, PEER_ERR_PROTOCOL);
                return -1;
            }
            r->cur_slot_ptr = lp->arena + (uint64_t)r->cur_slot * PEER_BLOCK_SIZE;
            r->body_got = 0;
            r->state = RS_PIECE_BODY;
            break;
        }

        case RS_PIECE_BODY: {
            size_t avail = (size_t)(end - p);
            size_t n = r->body_len - r->body_got;
            if (n > avail) n = avail;
            memcpy(r->cur_slot_ptr + r->body_got, p, n);
            p += n;
            r->body_got += n;
            if (r->body_got < r->body_len) goto out;
            engine_block d = { tor->id, r->cur_piece, r->cur_begin, r->body_len,
                               (uint32_t)lp->index, r->cur_slot };
            /* ready_ring capacity >= num_slots, so this cannot fail. */
            desc_ring_push(&lp->ready_ring, d);
            torrent_mark_received(tor, r->cur_piece, r->cur_begin);
            cancel_redundant_copies(c, r->cur_piece, r->cur_begin, r->body_len);
            atomic_fetch_add_explicit(&c->bytes_received, r->body_len,
                                      memory_order_relaxed);
            atomic_fetch_add_explicit(&c->blocks_received, 1, memory_order_relaxed);
            pushed = 1;
            r->state = RS_LEN;
            break;
        }

        case RS_PIECE_DROP:
            if (!feed_take(&p, end, r->other, r->other_total, &r->other_got)) goto out;
            r->state = RS_LEN;
            break;

        case RS_OTHER:
            if (!feed_take(&p, end, r->other, r->other_total, &r->other_got)) goto out;
            handle_control(c, r->msg_id, r->other, r->other_total);
            if (atomic_load_explicit(&c->aerror, memory_order_relaxed) != PEER_OK)
                return -1;
            r->state = RS_LEN;
            break;
        }
    }

out:
    if (pushed) signal_ready(c);
    return 0;
}