/*
 * engine.c - Public engine ABI, object lifecycle, the torrent registry, and the
 * control-plane command fan-out.
 *
 * The engine owns a fixed pool of event loops (loop.c). Torrents are pinned to
 * the least-loaded loop at registration time ("affinity"); all of a torrent's
 * connections then live on that one loop, which keeps the per-torrent piece
 * state lock-free. Control calls (add peer, set priorities) are turned into
 * commands posted to the owning loop; the data plane (poll/release) talks to the
 * per-loop arenas and rings directly.
 */
#include "engine_internal.h"
#include "arena.h"

#include <poll.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#include <sys/epoll.h>
#include <sys/eventfd.h>

#define DEFAULT_SLOTS      4096u   /* 64 MiB arena per loop                  */
#define DEFAULT_PIPELINE   2048u   /* per-connection outstanding cap         */
#define DEFAULT_TIMEOUT_MS 30000u  /* re-request a block after this long     */
#define DEFAULT_CONNECT_MS 10000u  /* drop a peer stuck connecting/handshaking */
#define DEFAULT_MAX_LOOPS  8u
#define RECV_STAGING_CAP   (256u * 1024u)

/* ---- shared helpers -------------------------------------------------- */

/* Update the limit (control plane). burst is one second of credit, floored at
 * 1 MiB so a small limit can still admit whole blocks. */
void rate_set(rate_limiter *rl, uint64_t bytes_per_sec) {
    pthread_mutex_lock(&rl->lock);
    atomic_store_explicit(&rl->rate_bps, bytes_per_sec, memory_order_relaxed);
    rl->burst = bytes_per_sec > (1u << 20) ? bytes_per_sec : (1u << 20);
    rl->last_ns = peer_now_ns();
    if (rl->tokens > (double)rl->burst) rl->tokens = (double)rl->burst;
    pthread_mutex_unlock(&rl->lock);
}

/* Token-bucket throttle. Refills lazily: tokens accrue at rate_bps since the
 * last call, capped at burst. Returns false without consuming when starved. */
bool rate_try_consume(rate_limiter *rl, uint32_t bytes) {
    if (atomic_load_explicit(&rl->rate_bps, memory_order_relaxed) == 0)
        return true;   /* unlimited; no lock on the hot path */
    pthread_mutex_lock(&rl->lock);
    bool ok = true;
    if (atomic_load_explicit(&rl->rate_bps, memory_order_relaxed) == 0) {
        pthread_mutex_unlock(&rl->lock);   /* became unlimited */
        return true;
    }
    uint64_t now = peer_now_ns();
    if (rl->last_ns == 0) rl->last_ns = now;
    double accrued = (double)(now - rl->last_ns) * 1e-9 *
                     (double)atomic_load_explicit(&rl->rate_bps,
                                                  memory_order_relaxed);
    rl->last_ns = now;
    rl->tokens += accrued;
    if (rl->tokens > (double)rl->burst) rl->tokens = (double)rl->burst;
    if (rl->tokens >= (double)bytes)
        rl->tokens -= (double)bytes;
    else
        ok = false;
    pthread_mutex_unlock(&rl->lock);
    return ok;
}

uint64_t peer_now_ns(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (uint64_t)ts.tv_sec * 1000000000ull + (uint64_t)ts.tv_nsec;
}

uint64_t torrent_piece_len(const torrent *t, uint32_t piece) {
    if (piece + 1 == t->num_pieces) {
        uint64_t before = (uint64_t)piece * t->piece_length;
        return t->total_size - before;
    }
    return t->piece_length;
}

torrent *engine_find_torrent(engine *e, uint32_t id) {
    for (torrent *t = e->torrents; t; t = t->enext)
        if (t->id == id) return t;
    return NULL;
}

void loop_post(loop *lp, cmd *c) {
    pthread_mutex_lock(&lp->cmd_lock);
    c->next = NULL;
    if (lp->cmd_tail) lp->cmd_tail->next = c; else lp->cmd_head = c;
    lp->cmd_tail = c;
    pthread_mutex_unlock(&lp->cmd_lock);
    uint64_t one = 1;
    ssize_t w = write(lp->cmd_efd, &one, sizeof one);
    (void)w;
}

/* ---- lifecycle ------------------------------------------------------- */

static int loop_init(engine *e, loop *lp, uint32_t index) {
    lp->eng = e;
    lp->index = (int)index;
    lp->epfd = -1;
    lp->cmd_efd = -1;
    atomic_init(&lp->stop, 0);
    atomic_init(&lp->want_release_wake, 0);
    pthread_mutex_init(&lp->cmd_lock, NULL);

    lp->num_slots = e->cfg.slots_per_loop;
    lp->arena = arena_alloc(lp->num_slots, &lp->arena_bytes);
    if (!lp->arena) return -1;
    if (slot_ring_init(&lp->free_ring, lp->num_slots) != 0) return -1;
    if (desc_ring_init(&lp->ready_ring, lp->num_slots) != 0) return -1;
    for (uint32_t s = 0; s < lp->num_slots; s++) slot_ring_push(&lp->free_ring, s);

    lp->recvbuf_cap = RECV_STAGING_CAP;
    lp->recvbuf = malloc(lp->recvbuf_cap);
    if (!lp->recvbuf) return -1;

    lp->epfd = epoll_create1(0);
    lp->cmd_efd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC);
    if (lp->epfd < 0 || lp->cmd_efd < 0) return -1;

    struct epoll_event ev;
    memset(&ev, 0, sizeof ev);
    ev.events = EPOLLIN;
    ev.data.ptr = NULL;            /* NULL = the command/credit eventfd */
    epoll_ctl(lp->epfd, EPOLL_CTL_ADD, lp->cmd_efd, &ev);
    return 0;
}

engine *engine_create(const engine_config *cfg) {
    engine *e = calloc(1, sizeof *e);
    if (!e) return NULL;
    pthread_mutex_init(&e->lock, NULL);
    pthread_mutex_init(&e->dl_limit.lock, NULL);  /* rate_bps 0 => unlimited */
    e->ready_efd = -1;
    if (cfg) e->cfg = *cfg;

    if (e->cfg.loop_count == 0) {
        long nc = sysconf(_SC_NPROCESSORS_ONLN);
        if (nc < 1) nc = 1;
        if (nc > (long)DEFAULT_MAX_LOOPS) nc = DEFAULT_MAX_LOOPS;
        e->cfg.loop_count = (uint32_t)nc;
    }
    if (e->cfg.slots_per_loop == 0)     e->cfg.slots_per_loop = DEFAULT_SLOTS;
    if (e->cfg.max_pipeline == 0)       e->cfg.max_pipeline = DEFAULT_PIPELINE;
    if (e->cfg.max_pipeline > e->cfg.slots_per_loop)
        e->cfg.max_pipeline = e->cfg.slots_per_loop;
    if (e->cfg.request_timeout_ms == 0) e->cfg.request_timeout_ms = DEFAULT_TIMEOUT_MS;
    if (e->cfg.connect_timeout_ms == 0) e->cfg.connect_timeout_ms = DEFAULT_CONNECT_MS;
    /* cfg.fallback defaults to 0 (single attempt) -- left as-is. */
    e->nloops = e->cfg.loop_count;

    e->ready_efd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC);
    if (e->ready_efd < 0) goto fail;

    /* loop embeds rings whose head/tail are _Alignas(64), so loop is
     * over-aligned; calloc only guarantees max_align_t. Allocate the array with
     * the real alignment (sizeof(loop) is a multiple of it). */
    size_t align = _Alignof(loop);
    if (align < sizeof(void *)) align = sizeof(void *);
    e->loops = aligned_alloc(align, (size_t)e->nloops * sizeof(loop));
    if (!e->loops) goto fail;
    memset(e->loops, 0, (size_t)e->nloops * sizeof(loop));

    for (uint32_t i = 0; i < e->nloops; i++)
        if (loop_init(e, &e->loops[i], i) != 0) goto fail;

    for (uint32_t i = 0; i < e->nloops; i++) {
        if (pthread_create(&e->loops[i].thread, NULL, loop_run, &e->loops[i]) != 0)
            goto fail;
        e->loops[i].thread_started = 1;
    }
    return e;

fail:
    engine_destroy(e);
    return NULL;
}

void engine_destroy(engine *e) {
    if (!e) return;

    if (e->loops) {
        for (uint32_t i = 0; i < e->nloops; i++) {
            loop *lp = &e->loops[i];
            if (lp->thread_started) {
                atomic_store_explicit(&lp->stop, 1, memory_order_relaxed);
                if (lp->cmd_efd >= 0) {
                    uint64_t one = 1;
                    ssize_t w = write(lp->cmd_efd, &one, sizeof one);
                    (void)w;
                }
            }
        }
        for (uint32_t i = 0; i < e->nloops; i++) {
            loop *lp = &e->loops[i];
            if (lp->thread_started) {
                pthread_join(lp->thread, NULL);
                lp->thread_started = 0;
            }
        }
        /* Threads are gone: tear down loop-owned state without contention. */
        for (uint32_t i = 0; i < e->nloops; i++) {
            loop *lp = &e->loops[i];
            conn *c = lp->conns;
            while (c) { conn *nx = c->next; conn_destroy(c); c = nx; }
            cmd *cm = lp->cmd_head;
            while (cm) {
                cmd *nx = cm->next;
                if (cm->kind == CMD_SET_PRIORITIES) free(cm->pri);
                free(cm);
                cm = nx;
            }
            pthread_mutex_destroy(&lp->cmd_lock);
            if (lp->cmd_efd >= 0) close(lp->cmd_efd);
            if (lp->epfd >= 0) close(lp->epfd);
            free(lp->recvbuf);
            desc_ring_free(&lp->ready_ring);
            slot_ring_free(&lp->free_ring);
            if (lp->arena) arena_free(lp->arena);
        }
        free(e->loops);
    }

    torrent *t = e->torrents;
    while (t) {
        torrent *nx = t->enext;
        free(t->priority);
        free(t->requested);
        if (t->recv_bits) {
            for (uint32_t p = 0; p < t->num_pieces; p++) free(t->recv_bits[p]);
            free(t->recv_bits);
        }
        free(t);
        t = nx;
    }

    if (e->ready_efd >= 0) close(e->ready_efd);
    pthread_mutex_destroy(&e->lock);
    free(e);
}

/* ---- control plane --------------------------------------------------- */

int32_t engine_add_torrent(engine *e, const uint8_t info_hash[20],
                           const uint8_t peer_id[20], uint64_t piece_length,
                           uint64_t total_size, uint32_t num_pieces) {
    if (!e || num_pieces == 0 || piece_length == 0 || total_size == 0) return -1;

    torrent *t = calloc(1, sizeof *t);
    if (!t) return -1;
    t->eng = e;
    memcpy(t->info_hash, info_hash, 20);
    memcpy(t->peer_id, peer_id, 20);
    t->piece_length = piece_length;
    t->total_size = total_size;
    t->num_pieces = num_pieces;
    t->bpp = (uint32_t)((piece_length + PEER_BLOCK_SIZE - 1) / PEER_BLOCK_SIZE);
    t->bf_bytes = (num_pieces + 7) / 8;
    t->priority = calloc(num_pieces, 1);
    t->requested = calloc(num_pieces, 1);
    t->recv_bits = calloc(num_pieces, sizeof(*t->recv_bits));
    if (!t->priority || !t->requested || !t->recv_bits) {
        free(t->priority);
        free(t->requested);
        free(t->recv_bits);
        free(t);
        return -1;
    }

    pthread_mutex_lock(&e->lock);
    loop *best = &e->loops[0];
    for (uint32_t i = 1; i < e->nloops; i++)
        if (e->loops[i].load < best->load) best = &e->loops[i];
    t->lp = best;
    best->load++;
    t->id = e->next_torrent_id++;
    t->enext = e->torrents;
    e->torrents = t;
    t->next = best->tors;          /* bookkeeping; loop thread never walks it */
    best->tors = t;
    pthread_mutex_unlock(&e->lock);

    return (int32_t)t->id;
}

static torrent *find_locked(engine *e, uint32_t id) {
    pthread_mutex_lock(&e->lock);
    torrent *t = engine_find_torrent(e, id);
    pthread_mutex_unlock(&e->lock);
    return t;
}

int engine_add_peer(engine *e, uint32_t torrent_id, const char *ip, uint16_t port) {
    if (!e || !ip) return -1;
    torrent *t = find_locked(e, torrent_id);
    if (!t) return -1;

    cmd *c = calloc(1, sizeof *c);
    if (!c) return -1;
    c->kind = CMD_ADD_PEER;
    c->tor = t;
    strncpy(c->ip, ip, sizeof c->ip - 1);
    c->port = port;
    loop_post(t->lp, c);
    return 0;
}

int engine_set_priorities(engine *e, uint32_t torrent_id,
                          const uint8_t *priorities, uint32_t count) {
    if (!e || !priorities) return -1;
    torrent *t = find_locked(e, torrent_id);
    if (!t || count != t->num_pieces) return -1;

    uint8_t *copy = malloc(count);
    if (!copy) return -1;
    memcpy(copy, priorities, count);

    cmd *c = calloc(1, sizeof *c);
    if (!c) { free(copy); return -1; }
    c->kind = CMD_SET_PRIORITIES;
    c->tor = t;
    c->pri = copy;
    c->pri_count = count;
    loop_post(t->lp, c);
    return 0;
}

int engine_set_priority(engine *e, uint32_t torrent_id, uint32_t piece,
                        uint8_t value) {
    if (!e) return -1;
    torrent *t = find_locked(e, torrent_id);
    if (!t || piece >= t->num_pieces) return -1;

    cmd *c = calloc(1, sizeof *c);
    if (!c) return -1;
    c->kind = CMD_SET_PRIORITY;
    c->tor = t;
    c->piece = piece;
    c->value = value;
    loop_post(t->lp, c);
    return 0;
}

int engine_request_piece(engine *e, uint32_t torrent_id, uint32_t piece) {
    if (!e) return -1;
    torrent *t = find_locked(e, torrent_id);
    if (!t || piece >= t->num_pieces) return -1;

    cmd *c = calloc(1, sizeof *c);
    if (!c) return -1;
    c->kind = CMD_REQUEST_PIECE;
    c->tor = t;
    c->piece = piece;
    loop_post(t->lp, c);
    return 0;
}

void engine_set_download_rate(engine *e, uint64_t bytes_per_sec) {
    if (e) rate_set(&e->dl_limit, bytes_per_sec);
}

/* ---- data plane ------------------------------------------------------ */

uint32_t engine_poll_ready(engine *e, engine_block *out, uint32_t max) {
    uint32_t n = 0;
    for (uint32_t i = 0; i < e->nloops && n < max; i++) {
        desc_ring *r = &e->loops[i].ready_ring;
        while (n < max && desc_ring_pop(r, &out[n])) n++;
    }
    return n;
}

void engine_release_slot(engine *e, uint32_t loop, uint32_t slot) {
    if (!e || loop >= e->nloops) return;
    slot_ring_push(&e->loops[loop].free_ring, slot);
    /* If the loop parked out of credit, wake it so the returned slot becomes a
     * request immediately. Coalesced: only the first release after starvation
     * pays the eventfd write. */
    if (atomic_exchange_explicit(&e->loops[loop].want_release_wake, 0,
                                 memory_order_relaxed)) {
        uint64_t one = 1;
        ssize_t w = write(e->loops[loop].cmd_efd, &one, sizeof one);
        (void)w;
    }
}

int engine_wait(engine *e, int timeout_ms) {
    struct pollfd pfd = { .fd = e->ready_efd, .events = POLLIN, .revents = 0 };
    int r = poll(&pfd, 1, timeout_ms);
    if (r < 0) return -1;
    if (r == 0) return 0;
    uint64_t drain;
    ssize_t rd = read(e->ready_efd, &drain, sizeof drain);
    (void)rd;
    return 1;
}

void *engine_arena_base(engine *e, uint32_t loop) {
    if (!e || loop >= e->nloops) return NULL;
    return e->loops[loop].arena;
}

uint64_t engine_arena_bytes(engine *e, uint32_t loop) {
    if (!e || loop >= e->nloops) return 0;
    return e->loops[loop].arena_bytes;
}

uint32_t engine_loop_count(engine *e) { return e ? e->nloops : 0; }

void engine_torrent_status(engine *e, uint32_t torrent_id, torrent_status *out) {
    memset(out, 0, sizeof *out);
    out->state = PEER_STATE_IDLE;
    if (!e) return;
    torrent *t = find_locked(e, torrent_id);
    if (!t) return;

    loop *lp = t->lp;
    int any_running = 0, max_state = PEER_STATE_IDLE, err = PEER_OK;
    uint32_t peers = 0, connected = 0, failed = 0, outst = 0, ptarget = 0;
    uint64_t bytes = 0, blocks = 0;
    double rate = 0.0, rttmin = 0.0;

    for (conn *c = lp->conns; c; c = c->next) {
        if (c->tor != t) continue;
        peers++;
        int st = atomic_load_explicit(&c->astate, memory_order_relaxed);
        int er = atomic_load_explicit(&c->aerror, memory_order_relaxed);
        if (er != PEER_OK) err = er;
        if (st == PEER_STATE_ERROR) failed++;
        if (st == PEER_STATE_RUNNING || st == PEER_STATE_CHOKED) connected++;
        if (st == PEER_STATE_RUNNING) any_running = 1;
        if (st > max_state) max_state = st;
        bytes  += atomic_load_explicit(&c->bytes_received, memory_order_relaxed);
        blocks += atomic_load_explicit(&c->blocks_received, memory_order_relaxed);
        outst  += atomic_load_explicit(&c->outstanding, memory_order_relaxed);
        ptarget += c->pipeline_target;
        rate += c->rate_bps;
        double rm = (double)c->rtt_min_ns / 1e6;
        if (c->rtt_min_ns && (rttmin == 0.0 || rm < rttmin)) rttmin = rm;
    }

    out->state = peers ? (any_running ? PEER_STATE_RUNNING : max_state)
                       : PEER_STATE_IDLE;
    out->error = err;
    out->bytes_received = bytes;
    out->blocks_received = blocks;
    out->peers = peers;
    out->peers_connected = connected;
    out->peers_failed = failed;
    out->outstanding = outst;
    out->free_slots = slot_ring_count(&lp->free_ring);
    out->pipeline_target = ptarget;
    out->rate_bps = rate;
    out->rtt_min_ms = rttmin;
}

static const char *peer_state_name(int state) {
    switch (state) {
    case PEER_STATE_IDLE:       return "idle";
    case PEER_STATE_CONNECTING: return "connecting";
    case PEER_STATE_HANDSHAKE:  return "handshake";
    case PEER_STATE_CHOKED:     return "choked";
    case PEER_STATE_RUNNING:    return "running";
    case PEER_STATE_STOPPED:    return "stopped";
    case PEER_STATE_ERROR:      return "error";
    default:                    return "?";
    }
}

void engine_dump_torrent(engine *e, uint32_t torrent_id, FILE *out) {
    if (!e || !out) return;
    torrent *t = find_locked(e, torrent_id);
    if (!t) { fprintf(out, "engine: torrent %u not found\n", torrent_id); return; }
    loop *lp = t->lp;

    /* Tally availability (connected peers advertising the piece) and in-flight
     * block requests per piece in a single pass over the loop's connections,
     * then walk the wanted pieces once. Avoids an O(pieces * conns) scan. */
    uint32_t *avail = calloc(t->num_pieces, sizeof *avail);
    uint32_t *inflight = calloc(t->num_pieces, sizeof *inflight);
    if (!avail || !inflight) {
        free(avail); free(inflight);
        fprintf(out, "engine: out of memory rendering dump for torrent %u\n",
                torrent_id);
        return;
    }

    fprintf(out, "=== engine dump: torrent %u (%u pieces, %llu B/piece) ===\n",
            torrent_id, t->num_pieces, (unsigned long long)t->piece_length);
    fprintf(out, "loop %d: outstanding=%u free_slots=%u\n",
            lp->index, lp->outstanding, slot_ring_count(&lp->free_ring));

    fprintf(out, "connections:\n");
    uint32_t conns = 0, connected = 0;
    for (conn *c = lp->conns; c; c = c->next) {
        if (c->tor != t) continue;
        conns++;
        int st = atomic_load_explicit(&c->astate, memory_order_relaxed);
        bool up = (st == PEER_STATE_RUNNING || st == PEER_STATE_CHOKED) && !c->dead;
        if (up) connected++;

        uint32_t have = 0;
        if (c->have_bits)
            for (uint32_t i = 0; i < t->num_pieces; i++) {
                if (!have_bit(c->have_bits, i)) continue;
                have++;
                if (up) avail[i]++;
            }
        if (c->inflight.slots) {
            uint32_t cap = c->inflight.mask + 1;
            for (uint32_t i = 0; i < cap; i++) {
                if (c->inflight.slots[i].key == REQ_EMPTY) continue;
                uint32_t p = c->inflight.slots[i].piece;
                if (p < t->num_pieces) inflight[p]++;
            }
        }

        uint32_t out_reqs = atomic_load_explicit(&c->outstanding,
                                                 memory_order_relaxed);
        uint64_t rx = atomic_load_explicit(&c->bytes_received,
                                           memory_order_relaxed);
        fprintf(out,
                "  %s:%u state=%s%s unchoked=%d cur_piece=%s out=%u "
                "requeue=%u rate=%.1fKB/s have=%u/%u rx=%lluB\n",
                c->ip, c->port, peer_state_name(st), c->dead ? "(dead)" : "",
                c->unchoked, c->have_cur_piece ? "" : "-",
                out_reqs, c->requeue_count, c->rate_bps / 1024.0,
                have, t->num_pieces, (unsigned long long)rx);
        if (c->have_cur_piece)
            fprintf(out, "       (working piece %u)\n", c->cur_piece);
    }
    fprintf(out, "peers: %u total, %u connected\n", conns, connected);

    /* Per-piece breakdown of everything still wanted. At the tail of a download
     * this is the handful of pieces that refuse to finish. */
    uint32_t wanted = 0, claimed = 0, starved = 0;
    for (uint32_t i = 0; i < t->num_pieces; i++) {
        if (t->priority[i] == 0) continue;
        wanted++;
        if (t->requested[i]) claimed++;
        if (avail[i] == 0) starved++;
    }
    fprintf(out,
            "wanted pieces: %u (claimed=%u, no-connected-peer-has-it=%u)\n",
            wanted, claimed, starved);

    uint32_t shown = 0;
    const uint32_t limit = 512;
    for (uint32_t i = 0; i < t->num_pieces; i++) {
        if (t->priority[i] == 0) continue;
        if (shown++ >= limit) continue;
        fprintf(out,
                "  piece %u pri=%u claimed=%u avail=%u inflight=%u%s\n",
                i, t->priority[i], t->requested[i], avail[i], inflight[i],
                (t->requested[i] && inflight[i] == 0)
                    ? "  <-- claimed but no requests in flight"
                    : (avail[i] == 0 ? "  <-- no connected peer has it" : ""));
    }
    if (wanted > limit)
        fprintf(out, "  ... %u more wanted pieces not shown\n", wanted - limit);

    free(avail);
    free(inflight);
}