#include "naut/piece.h"
#include "naut/bitfield.h"
#include "naut/hash.h"
#include "naut/log.h"

#include <stdlib.h>
#include <string.h>

#define BLK NAUT_BLOCK          /* 16 KiB */
#define ENDGAME_BLOCKS 8        /* switch to endgame when this few remain */
#define ENDGAME_COPIES 2        /* at most two peers race a missing block */

/* per-piece in-progress state, lazily allocated and freed on completion */
typedef struct {
    uint8_t *recv_bits;   /* received block bitmap */
    uint8_t *req_count;   /* outstanding requests per block */
    uint8_t *buf;         /* assembly buffer, piece_size bytes */
    uint32_t nblocks;
    uint32_t nrecv;
    bool verifying;
    naut_job verify_job;
    struct naut_download *download;
    uint32_t piece;
    uint8_t digest[NAUT_SHA1_LEN];
} pstate;

struct naut_download {
    const naut_metainfo *mi;
    naut_storage *st;

    uint32_t num_pieces;
    uint64_t piece_len;
    uint64_t total;

    naut_bitfield have;
    uint32_t  *avail;       /* [num_pieces] swarm availability count */
    pstate   **ps;          /* [num_pieces] in-progress state or NULL */

    uint32_t cur_piece;     /* sequential cursor for next_request() */

    uint64_t total_blocks, recv_blocks;
    uint32_t pieces_done;
    uint64_t bytes_done;
    bool     endgame;
    naut_worker_pool *workers;

    size_t    num_files;
    uint32_t *file_first, *file_last, *file_remain;
    bool     *file_done;
    naut_file_complete_cb file_cb;
    void                 *file_cb_ctx;
};

static bool bget(const uint8_t *a, uint32_t i) { return (a[i>>3] >> (i&7)) & 1; }
static void bset(uint8_t *a, uint32_t i) { a[i>>3] |= (uint8_t)(1u << (i&7)); }

static uint64_t piece_size(const naut_download *d, uint32_t p) {
    if (p + 1 < d->num_pieces) return d->piece_len;
    return d->total - (uint64_t)p * d->piece_len;
}
static uint32_t nblocks(const naut_download *d, uint32_t p) {
    return (uint32_t)((piece_size(d, p) + BLK - 1) / BLK);
}
static uint32_t block_len(const naut_download *d, uint32_t p, uint32_t b) {
    uint64_t rem = piece_size(d, p) - (uint64_t)b * BLK;
    return rem < BLK ? (uint32_t)rem : BLK;
}

static pstate *ensure_ps(naut_download *d, uint32_t p) {
    if (d->ps[p]) return d->ps[p];
    pstate *s = calloc(1, sizeof(*s));
    if (!s) return NULL;
    s->nblocks = nblocks(d, p);
    s->download = d;
    s->piece = p;
    size_t bm = (s->nblocks + 7) / 8;
    s->recv_bits = calloc(1, bm);
    s->req_count = calloc(s->nblocks, sizeof(uint8_t));
    size_t alloc_size =
        (size_t)NAUT_ALIGN_UP(piece_size(d, p), NAUT_PAGE);
    s->buf = aligned_alloc(NAUT_PAGE, alloc_size);
    if (!s->recv_bits || !s->req_count || !s->buf) {
        free(s->recv_bits); free(s->req_count); free(s->buf); free(s);
        return NULL;
    }
    d->ps[p] = s;
    return s;
}
static void free_ps(naut_download *d, uint32_t p) {
    pstate *s = d->ps[p];
    if (!s) return;
    free(s->recv_bits); free(s->req_count); free(s->buf); free(s);
    d->ps[p] = NULL;
}

naut_download *naut_download_create(const naut_metainfo *mi, naut_storage *st) {
    if (!mi->has_v1 || mi->num_pieces == 0 || mi->piece_length <= 0 ||
        mi->total_length <= 0) {
        NAUT_ERROR("download: needs a v1/hybrid torrent (SHA-1 pieces)");
        return NULL;
    }
    uint64_t total = (uint64_t)mi->total_length;
    uint64_t piece_len = (uint64_t)mi->piece_length;
    uint64_t expected_pieces = 1 + (total - 1) / piece_len;
    if (expected_pieces != mi->num_pieces || piece_len > UINT32_MAX * (uint64_t)BLK) {
        NAUT_ERROR("download: inconsistent piece geometry");
        return NULL;
    }
    naut_download *d = calloc(1, sizeof(*d));
    if (!d) return NULL;
    d->mi = mi; d->st = st;
    d->num_pieces = mi->num_pieces;
    d->piece_len = (uint64_t)mi->piece_length;
    d->total = (uint64_t)mi->total_length;
    d->avail = calloc(d->num_pieces, sizeof(uint32_t));
    d->ps    = calloc(d->num_pieces, sizeof(pstate *));
    if (!d->avail || !d->ps || naut_bitfield_init(&d->have, d->num_pieces) != NAUT_OK) {
        naut_download_destroy(d); return NULL;
    }
    for (uint32_t p = 0; p < d->num_pieces; p++) d->total_blocks += nblocks(d, p);

    d->num_files   = mi->num_files;
    d->file_first  = calloc(mi->num_files, sizeof(uint32_t));
    d->file_last   = calloc(mi->num_files, sizeof(uint32_t));
    d->file_remain = calloc(mi->num_files, sizeof(uint32_t));
    d->file_done   = calloc(mi->num_files, sizeof(bool));
    if (mi->num_files && (!d->file_first || !d->file_last || !d->file_remain || !d->file_done)) {
        naut_download_destroy(d); return NULL;
    }
    uint64_t off = 0;
    for (size_t f = 0; f < mi->num_files; f++) {
        uint64_t flen = (uint64_t)mi->files[f].length;
        if (flen == 0) { d->file_first[f] = 1; d->file_last[f] = 0; d->file_done[f] = true; }
        else {
            d->file_first[f]  = (uint32_t)(off / d->piece_len);
            d->file_last[f]   = (uint32_t)((off + flen - 1) / d->piece_len);
            d->file_remain[f] = d->file_last[f] - d->file_first[f] + 1;
        }
        off += flen;
    }
    return d;
}

void naut_download_destroy(naut_download *d) {
    if (!d) return;
    if (d->ps) for (uint32_t p = 0; p < d->num_pieces; p++) free_ps(d, p);
    free(d->ps); free(d->avail);
    free(d->file_first); free(d->file_last); free(d->file_remain); free(d->file_done);
    naut_bitfield_free(&d->have);
    free(d);
}

void naut_download_set_worker_pool(naut_download *d, naut_worker_pool *pool) {
    if (d) d->workers = pool;
}

void naut_download_set_file_cb(naut_download *d, naut_file_complete_cb cb, void *ctx) {
    d->file_cb = cb; d->file_cb_ctx = ctx;
}
bool naut_download_file_complete(const naut_download *d, uint32_t f) {
    return f < d->num_files && d->file_done[f];
}

static void notify_files(naut_download *d, uint32_t p) {
    size_t lo = 0, hi = d->num_files;
    while (lo < hi) { size_t mid = (lo + hi) / 2;
        if (d->file_last[mid] < p) lo = mid + 1; else hi = mid; }
    for (size_t f = lo; f < d->num_files && d->file_first[f] <= p; f++) {
        if (d->file_done[f]) continue;
        if (--d->file_remain[f] == 0) {
            d->file_done[f] = true;
            if (d->file_cb) d->file_cb(d->file_cb_ctx, (uint32_t)f, d->mi->files[f].path);
        }
    }
}

/* --- availability -------------------------------------------------------- */
void naut_download_inc_avail(naut_download *d, uint32_t p) {
    if (p < d->num_pieces) d->avail[p]++;
}
void naut_download_add_bitfield(naut_download *d, const naut_bitfield *bf) {
    uint32_t limit = (uint32_t)NAUT_MIN((size_t)d->num_pieces, bf->nbits);
    for (uint32_t p = 0; p < limit; p++)
        if (naut_bitfield_test(bf, p)) d->avail[p]++;
}
void naut_download_remove_bitfield(naut_download *d, const naut_bitfield *bf) {
    uint32_t limit = (uint32_t)NAUT_MIN((size_t)d->num_pieces, bf->nbits);
    for (uint32_t p = 0; p < limit; p++)
        if (naut_bitfield_test(bf, p) && d->avail[p]) d->avail[p]--;
}

/* --- request selection --------------------------------------------------- */
/* Find the first missing block with no outstanding request. */
static uint32_t first_unreq(const pstate *s) {
    if (s->verifying) return UINT32_MAX;
    for (uint32_t b = 0; b < s->nblocks; b++)
        if (!bget(s->recv_bits, b) && s->req_count[b] == 0) return b;
    return UINT32_MAX;
}

static bool hand_out(naut_download *d, uint32_t p, uint32_t b,
                     uint32_t *index, uint32_t *begin, uint32_t *length) {
    d->ps[p]->req_count[b]++;
    *index = p; *begin = b * BLK; *length = block_len(d, p, b);
    return true;
}

bool naut_download_pick_for_peer(naut_download *d, const naut_bitfield *peer_have,
                                 naut_request_active_cb peer_has_request, void *ctx,
                                 uint32_t *index, uint32_t *begin, uint32_t *length) {
    d->endgame = (d->total_blocks - d->recv_blocks) <= ENDGAME_BLOCKS;

    /* pass 1: finish an in-progress piece the peer has (reduces fragmentation) */
    for (uint32_t p = 0; p < d->num_pieces; p++) {
        if (naut_bitfield_test(&d->have, p) || !d->ps[p]) continue;
        if (p >= peer_have->nbits || !naut_bitfield_test(peer_have, p)) continue;
        uint32_t b = first_unreq(d->ps[p]);
        if (b != UINT32_MAX) return hand_out(d, p, b, index, begin, length);
    }
    /* pass 2: start the rarest new piece the peer has */
    uint32_t best = UINT32_MAX, best_av = UINT32_MAX;
    for (uint32_t p = 0; p < d->num_pieces; p++) {
        if (naut_bitfield_test(&d->have, p) || d->ps[p]) continue;
        if (p >= peer_have->nbits || !naut_bitfield_test(peer_have, p) ||
            d->avail[p] == 0) continue;
        if (d->avail[p] < best_av) { best = p; best_av = d->avail[p]; }
    }
    if (best != UINT32_MAX) {
        if (!ensure_ps(d, best)) return false;
        return hand_out(d, best, 0, index, begin, length);
    }
    /* pass 3: endgame — race each missing block on at most two distinct peers */
    if (d->endgame) {
        for (uint8_t copies = 1; copies < ENDGAME_COPIES; copies++) {
            for (uint32_t p = 0; p < d->num_pieces; p++) {
                if (naut_bitfield_test(&d->have, p) ||
                    p >= peer_have->nbits ||
                    !naut_bitfield_test(peer_have, p)) continue;
                if (!ensure_ps(d, p)) continue;
                pstate *s = d->ps[p];
                for (uint32_t b = 0; b < s->nblocks; b++) {
                    uint32_t block_begin = b * BLK;
                    if (bget(s->recv_bits, b) || s->req_count[b] != copies) continue;
                    if (peer_has_request &&
                        peer_has_request(ctx, p, block_begin)) continue;
                    return hand_out(d, p, b, index, begin, length);
                }
            }
        }
    }
    return false;
}

bool naut_download_pick(naut_download *d, const naut_bitfield *peer_have,
                        uint32_t *index, uint32_t *begin, uint32_t *length) {
    return naut_download_pick_for_peer(d, peer_have, NULL, NULL,
                                       index, begin, length);
}

void naut_download_unrequest(naut_download *d, uint32_t index, uint32_t begin) {
    if (index >= d->num_pieces || !d->ps[index]) return;
    uint32_t b = begin / BLK;
    if (b < d->ps[index]->nblocks && !bget(d->ps[index]->recv_bits, b) &&
        d->ps[index]->req_count[b] != 0)
        d->ps[index]->req_count[b]--;
}

/* sequential single-peer convenience (Phase 3 leecher + tests) */
bool naut_download_next_request(naut_download *d,
                                uint32_t *index, uint32_t *begin, uint32_t *length) {
    while (d->cur_piece < d->num_pieces) {
        if (naut_bitfield_test(&d->have, d->cur_piece)) { d->cur_piece++; continue; }
        pstate *s = ensure_ps(d, d->cur_piece);
        if (!s) return false;
        uint32_t b = first_unreq(s);
        if (b != UINT32_MAX) return hand_out(d, d->cur_piece, b, index, begin, length);
        d->cur_piece++;
    }
    return false;
}

/* --- block ingest -------------------------------------------------------- */
static naut_err finish_verified(naut_download *d, uint32_t p,
                                const uint8_t digest[NAUT_SHA1_LEN],
                                bool *done) {
    pstate *s = d->ps[p];
    uint64_t ps = piece_size(d, p);
    if (memcmp(digest,
               d->mi->piece_hashes + (size_t)p * NAUT_SHA1_LEN,
               NAUT_SHA1_LEN) != 0) {
        NAUT_WARN("piece %u failed SHA-1; discarding for re-download", p);
        d->recv_blocks -= s->nrecv;          /* roll back so it can be refetched */
        free_ps(d, p);
        return NAUT_ERR_PROTO;
    }
    naut_err e = naut_storage_write(d->st, (int64_t)p * (int64_t)d->piece_len, s->buf, ps);
    if (e != NAUT_OK) return e;
    naut_bitfield_set(&d->have, p);
    d->pieces_done++;
    d->bytes_done += ps;
    free_ps(d, p);
    *done = true;
    notify_files(d, p);
    return NAUT_OK;
}

static void verify_job_run(naut_job *job) {
    pstate *state = job->context;
    naut_sha1(state->buf, piece_size(state->download, state->piece),
              state->digest);
    job->result = NAUT_OK;
}

naut_err naut_download_poll(naut_download *d, uint32_t *pieces_completed) {
    if (!d) return NAUT_ERR_INVAL;
    if (pieces_completed) *pieces_completed = 0;
    if (!d->workers) return NAUT_OK;
    naut_job *job;
    naut_err result = NAUT_OK;
    while (naut_worker_complete(d->workers, &job)) {
        pstate *state = job->context;
        uint32_t piece = state->piece;
        if (state->download != d || piece >= d->num_pieces ||
            d->ps[piece] != state || !state->verifying) {
            result = NAUT_ERR_PROTO;
            continue;
        }
        bool done = false;
        naut_err e = finish_verified(d, piece, state->digest, &done);
        if (e == NAUT_ERR_PROTO) {
            /* Hash mismatch already reset the piece for re-download. The
             * worker cannot attribute corruption to one peer, so keep the
             * torrent alive and let the picker request it again. */
            continue;
        }
        if (e != NAUT_OK) {
            result = e;
            continue;
        }
        if (done && pieces_completed) (*pieces_completed)++;
    }
    return result;
}

naut_err naut_download_on_block(naut_download *d, uint32_t index, uint32_t begin,
                                const uint8_t *data, uint32_t len, bool *piece_done) {
    *piece_done = false;
    if (index >= d->num_pieces) return NAUT_ERR_RANGE;
    if (naut_bitfield_test(&d->have, index)) return NAUT_OK;       /* already complete */
    if (begin % BLK != 0) return NAUT_ERR_PROTO;
    uint32_t b = begin / BLK;
    if (b >= nblocks(d, index) || len != block_len(d, index, b)) return NAUT_ERR_PROTO;

    pstate *s = ensure_ps(d, index);
    if (!s) return NAUT_ERR_NOMEM;
    if (bget(s->recv_bits, b)) return NAUT_OK;                     /* duplicate, ignore */

    memcpy(s->buf + begin, data, len);
    bset(s->recv_bits, b);
    s->req_count[b] = 0;
    s->nrecv++;
    d->recv_blocks++;
    if (s->nrecv == s->nblocks) {
        if (d->workers) {
            s->verifying = true;
            s->verify_job.run = verify_job_run;
            s->verify_job.context = s;
            s->verify_job.result = NAUT_ERR_AGAIN;
            if (naut_worker_submit(d->workers, &s->verify_job))
                return NAUT_OK;
            s->verifying = false;
        }
        uint8_t digest[NAUT_SHA1_LEN];
        naut_sha1(s->buf, piece_size(d, index), digest);
        return finish_verified(d, index, digest, piece_done);
    }
    return NAUT_OK;
}

bool     naut_download_complete(const naut_download *d) { return d->pieces_done == d->num_pieces; }
bool     naut_download_have(const naut_download *d, uint32_t p) { return naut_bitfield_test(&d->have, p); }
bool     naut_download_in_endgame(const naut_download *d) { return d->endgame; }
uint32_t naut_download_num_pieces(const naut_download *d) { return d->num_pieces; }
uint32_t naut_download_pieces_done(const naut_download *d) { return d->pieces_done; }
uint64_t naut_download_bytes_done(const naut_download *d) { return d->bytes_done; }