// media_download.c — скачивание блоков медиа
#include "media_download.h"
#include "media_delivery.h"
#include "media_delivery_proto.h"
#include "../media_delivery/media_index.h"
#include "../utun_instance.h"
#include "../routing_layer/etcp_router.h"
#include "../routing_layer/conn_mgr.h"
#include "../routing_layer/topo_group.h"
#include "../transport_layer/etcp_api.h"
#include "../transport_layer/secure_channel.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"
#include "../lib/u_async.h"
#include "../lib/platform_compat.h"
#include <string.h>
#include <stdio.h>
#include <openssl/evp.h>

#define MDL_ID "media_download"

/* ── internal download state ── */

struct media_download {
    struct ll_entry ll;         // индекс по media_id (16 байт)
    uint8_t  media_id[16];
    uint64_t group_id;
    char     dest_path[1024];
    char     media_base[512];
    int      num_blocks;
    uint8_t* block_ids;         // num_blocks * 16
    uint8_t* block_sigs;        // num_blocks * 64
    uint8_t  content_hash[32];
    int64_t  file_size;
    int64_t  block_size;
    int      blocks_received;
    int      blocks_validated;
    int      num_peers;
    struct   media_download_peer peers[10];
    uint8_t  active;
    uint8_t  assembled;
    int      err;

    /* supernode list */
    uint64_t super_nodes[10];
    int      super_count;
    int      super_current;

    void*    query_timer;
    void*    timeout_timer;
    void     (*done_cb)(void* arg, int err);
    void*    done_arg;
};

/* ── helpers ── */

static struct media_download* md_dl_find(struct UTUN_INSTANCE* inst, const uint8_t* media_id) {
    struct media_delivery_ctx* md = &inst->md;
    if (!md->downloads) return NULL;
    struct ll_entry* e = queue_find_data_by_index(md->downloads, media_id);
    return e ? (struct media_download*)e->data : NULL;
}

static void md_dl_free(struct media_download* dl) {
    if (!dl) return;
    if (dl->query_timer)   { uasync_cancel_timeout(NULL, dl->query_timer); dl->query_timer = NULL; }
    if (dl->timeout_timer) { uasync_cancel_timeout(NULL, dl->timeout_timer); dl->timeout_timer = NULL; }
    if (dl->block_ids) { u_free(dl->block_ids); dl->block_ids = NULL; }
    if (dl->block_sigs) { u_free(dl->block_sigs); dl->block_sigs = NULL; }
}

/* ── send helpers ── */

static int md_dl_send(struct UTUN_INSTANCE* inst, uint64_t dst, const uint8_t* data, size_t len) {
    struct ll_entry* e = queue_entry_new(0);
    if (!e) return -1;
    e->dgram = u_malloc(len);
    if (!e->dgram) { queue_entry_free(e); return -1; }
    memcpy(e->dgram, data, len); e->len = (uint16_t)len;
    int rc = etcp_route_send(inst, 0, dst, e, 1);
    if (rc != 0) { u_free(e->dgram); queue_entry_free(e); }
    return rc;
}

/* ── build block request ── */

static int md_dl_send_block_req(struct UTUN_INSTANCE* inst, uint64_t dst,
                                 struct media_download* dl, int bi) {
    struct media_pkt_block_req req;
    memset(&req, 0, sizeof(req));
    req.subcmd = MEDIA_SUBCMD_BLOCK_REQ;
    memcpy(req.media_id, dl->media_id, 16);
    memcpy(req.block_id, dl->block_ids + bi * 16, 16);
    req.chunk = (uint32_t)bi;
    req.offset = 0;

    DEBUG_INFO(DEBUG_CATEGORY_DEBUG, "%s: BLOCK_REQ to 0x%016llx block=%d",
               MDL_ID, (unsigned long long)dst, bi);
    return md_dl_send(inst, dst, (const uint8_t*)&req, sizeof(req));
}

/* ── send HAVE_BLOCK to supernode ── */

static void md_dl_send_have_block(struct UTUN_INSTANCE* inst, struct media_download* dl, int bi) {
    uint64_t super = dl->super_count > 0 ? dl->super_nodes[dl->super_current] : 0;
    if (!super) return;

    struct media_pkt_have_block hb;
    memset(&hb, 0, sizeof(hb));
    hb.subcmd = MEDIA_SUBCMD_HAVE_BLOCK;
    hb.group_id = dl->group_id;
    memcpy(hb.media_id, dl->media_id, 16);
    memcpy(hb.block_id, dl->block_ids + bi * 16, 16);
    hb.chunk = (uint32_t)bi;
    hb.timestamp = (int64_t)time(NULL);

    /* sign: block_id + chunk + timestamp + my_node_id */
    uint8_t smsg[64]; size_t soff = 0;
    memcpy(smsg + soff, dl->block_ids + bi * 16, 16); soff += 16;
    memcpy(smsg + soff, &hb.chunk, 4); soff += 4;
    memcpy(smsg + soff, &hb.timestamp, 8); soff += 8;
    uint64_t self = inst->node_id; memcpy(smsg + soff, &self, 8); soff += 8;
    sc_ed25519_sign(inst->my_ed25519_privkey, smsg, soff, hb.node_sign);

    DEBUG_INFO(DEBUG_CATEGORY_DEBUG, "%s: HAVE_BLOCK to super 0x%016llx block=%d",
               MDL_ID, (unsigned long long)super, bi);
    md_dl_send(inst, super, (const uint8_t*)&hb, sizeof(hb));
}

/* ── conn_mgr callback ── */

static void md_dl_conn_cb(int result, uint64_t node_id, void* arg) {
    struct media_download* dl = (struct media_download*)arg;
    if (result != CONN_MGR_OK) {
        DEBUG_WARN(DEBUG_CATEGORY_GENERAL, "%s: connect to 0x%016llx failed rc=%d",
                   MDL_ID, (unsigned long long)node_id, result);
        return;
    }
    /* send block requests to this peer */
    for (int pi = 0; pi < dl->num_peers; pi++) {
        if (dl->peers[pi].node_id == node_id) {
            dl->peers[pi].connected = 1;
            for (int bi = 0; bi < dl->peers[pi].num_blocks; bi++) {
                if (!dl->peers[pi].blocks[bi].started) {
                    dl->peers[pi].blocks[bi].started = 1;
                    md_dl_send_block_req((struct UTUN_INSTANCE*)dl->done_arg, node_id, dl, bi);
                }
            }
        }
    }
}

/* ── query supernode, collect supernodes ── */

static void md_dl_collect_supernodes(struct UTUN_INSTANCE* inst, struct media_download* dl) {
    dl->super_count = 0;
    struct ll_entry* gle = inst->topo_groups ? inst->topo_groups->group_list->head : NULL;
    while (gle) {
        struct TOPO_GROUP* g = (struct TOPO_GROUP*)gle->data;
        if (g->group_type != TOPO_GROUP_TYPE_CHAT || !g->channel_id[0]) { gle = gle->next; continue; }
        char sql[256];
        snprintf(sql, sizeof(sql),
            "SELECT node_id FROM peers_%s WHERE node_type=4 AND node_id!=%llu",
            g->channel_id, (unsigned long long)inst->node_id);
        sqlite3_stmt* st = NULL;
        if (sqlite3_prepare_v2(inst->topo_sqlite_db, sql, -1, &st, NULL) == SQLITE_OK) {
            while (sqlite3_step(st) == SQLITE_ROW && dl->super_count < 10) {
                dl->super_nodes[dl->super_count] = (uint64_t)sqlite3_column_int64(st, 0);
                dl->super_count++;
            }
            sqlite3_finalize(st);
        }
        gle = gle->next;
    }
    dl->super_current = 0;
    if (dl->super_count == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "%s: no supernodes found", MDL_ID);
    }
}

/* ── send query to current supernode ── */

static void md_dl_send_query(struct UTUN_INSTANCE* inst, struct media_download* dl) {
    if (dl->super_current >= dl->super_count) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "%s: all supernodes exhausted", MDL_ID);
        dl->err = -1;
        if (dl->done_cb) dl->done_cb(dl->done_arg, dl->err);
        return;
    }
    uint64_t super = dl->super_nodes[dl->super_current];

    size_t pkt_len = MEDIA_QUERY_HDR_SIZE + (size_t)dl->num_blocks * 16;
    uint8_t* pkt = u_malloc(pkt_len);
    if (!pkt) return;
    struct media_pkt_query* q = (struct media_pkt_query*)pkt;
    memset(q, 0, sizeof(*q));
    q->subcmd = MEDIA_SUBCMD_QUERY;
    q->group_id = dl->group_id;
    memcpy(q->media_id, dl->media_id, 16);
    q->num_blocks = (uint16_t)dl->num_blocks;
    memcpy(pkt + MEDIA_QUERY_HDR_SIZE, dl->block_ids, (size_t)dl->num_blocks * 16);

    DEBUG_INFO(DEBUG_CATEGORY_DEBUG, "%s: QUERY to super 0x%016llx blocks=%d",
               MDL_ID, (unsigned long long)super, dl->num_blocks);
    md_dl_send(inst, super, pkt, pkt_len);
    u_free(pkt);
}

/* ── handle QUERY_RESP ── */

void media_download_handle_query_resp(struct UTUN_INSTANCE* inst,
                                      const uint8_t* data, size_t len) {
    if (!data || len < MEDIA_QUERY_RESP_HDR_SIZE) return;
    struct media_pkt_query_resp* r = (struct media_pkt_query_resp*)data;
    struct media_pkt_query_resp_entry* entries = (struct media_pkt_query_resp_entry*)(data + MEDIA_QUERY_RESP_HDR_SIZE);
    int ne = r->num_entries;
    if (ne > 100) ne = 100;

    /* find the download by media_id from first entry's block_id */
    if (ne < 1 || !inst->md.downloads) return;
    const uint8_t* bid = entries[0].block_id;
    struct media_download* dl = md_dl_find(inst, bid);
    if (!dl) {
        /* search all blocks */
        for (int i = 0; i < ne; i++) { dl = md_dl_find(inst, entries[i].block_id); if (dl) break; }
        if (!dl) return;
    }

    dl->num_peers = 0;

    /* group entries by node_id */
    for (int i = 0; i < ne && dl->num_peers < 10; i++) {
        uint64_t nid = entries[i].node_id;
        int pi = -1;
        for (int j = 0; j < dl->num_peers; j++) {
            if (dl->peers[j].node_id == nid) { pi = j; break; }
        }
        if (pi < 0) { pi = dl->num_peers; dl->num_peers++; }
        int bi = pi;
        struct media_download_peer* peer = &dl->peers[bi];
        peer->node_id = nid;
        int nb = peer->num_blocks;
        if (nb < MD_MAX_BLOCKS_PER_PEER) {
            memcpy(peer->blocks[nb].block_id, entries[i].block_id, 16);
            peer->blocks[nb].started = 0;
            peer->blocks[nb].received = 0;
            peer->blocks[nb].validated = 0;
            peer->num_blocks++;
        }
    }

    DEBUG_INFO(DEBUG_CATEGORY_DEBUG, "%s: QUERY_RESP entries=%d peers=%d",
               MDL_ID, ne, dl->num_peers);

    /* connect to peers */
    struct ll_entry* gle = inst->topo_groups->group_list->head;
    struct TOPO_GROUP* grp = NULL;
    while (gle) {
        struct TOPO_GROUP* g = (struct TOPO_GROUP*)gle->data;
        if (g->conn_mgr) { grp = g; break; }
        gle = gle->next;
    }
    if (!grp || !grp->conn_mgr) return;

    for (int pi = 0; pi < dl->num_peers; pi++) {
        conn_mgr_connect_node(grp->conn_mgr, dl->peers[pi].node_id, 0, md_dl_conn_cb, dl);
    }
}

/* ── handle incoming BLOCK_CHUNK ── */

void media_download_handle_chunk(struct UTUN_INSTANCE* inst,
                                  const uint8_t* data, size_t len) {
    if (!data || len < MEDIA_BLOCK_CHUNK_HDR_SIZE) return;

    const uint8_t* d = data;
    struct media_pkt_block_chunk* ch = (struct media_pkt_block_chunk*)data;
    struct media_download* dl = md_dl_find(inst, ch->media_id);
    if (!dl || !dl->active) return;

    /* find block index */
    int bi = -1;
    for (int i = 0; i < dl->num_blocks; i++) {
        if (memcmp(dl->block_ids + i * 16, ch->block_id, 16) == 0) { bi = i; break; }
    }
    if (bi < 0) return;

    /* write chunk to temp file */
    char tmp[2048];
    snprintf(tmp, sizeof(tmp), "%s.chunk_%d", dl->dest_path, bi);
    FILE* f = fopen(tmp, "ab");
    if (f) {
        size_t wlen = ch->data_len;
        if (len >= MEDIA_BLOCK_CHUNK_HDR_SIZE + wlen) {
            fwrite(d + MEDIA_BLOCK_CHUNK_HDR_SIZE, 1, wlen, f);
        }
        fclose(f);
    }
}

/* ── handle incoming BLOCK_DONE ── */

void media_download_handle_done(struct UTUN_INSTANCE* inst,
                                 const uint8_t* data, size_t len) {
    if (!data || len < MEDIA_BLOCK_DONE_SIZE) return;

    const uint8_t* d = data;
    struct media_pkt_block_done* bd = (struct media_pkt_block_done*)data;
    struct media_download* dl = md_dl_find(inst, bd->media_id);
    if (!dl || !dl->active) return;

    int bi = -1;
    for (int i = 0; i < dl->num_blocks; i++) {
        if (memcmp(dl->block_ids + i * 16, bd->block_id, 16) == 0) { bi = i; break; }
    }
    if (bi < 0) return;

    /* verify signature */
    EVP_PKEY* pkey = NULL;
    EVP_MD_CTX* vctx = EVP_MD_CTX_new();
    int sig_ok = 0;
    if (vctx) {
        /* read block data from temp file */
        char tmp[2048];
        snprintf(tmp, sizeof(tmp), "%s.chunk_%d", dl->dest_path, bi);
        FILE* f = fopen(tmp, "rb");
        if (f) {
            fseeko(f, 0, SEEK_END);
            off_t fsz = ftello(f);
            fseeko(f, 0, SEEK_SET);
            uint8_t* buf = u_malloc((size_t)fsz);
            if (buf) {
                size_t rd = fread(buf, 1, (size_t)fsz, f);
                uint64_t node_id = inst->node_id;
                uint8_t smsg[128]; size_t soff = 0;
                memcpy(smsg + soff, buf, rd); soff += rd;
                memcpy(smsg + soff, &node_id, 8); soff += 8;
                EVP_DigestVerifyInit(vctx, NULL, EVP_sha256(), NULL, pkey);
                /* fallback: compare raw block_sig */
                sig_ok = (rd == (size_t)bd->total_size && rd > 0) ? 1 : 0;
                if (sig_ok && rd == (size_t)bd->total_size) {
                    /* basic check: block_sig should match pre-computed */
                    if (memcmp(bd->block_sig, dl->block_sigs + bi * 64, 64) == 0) sig_ok = 1;
                    else sig_ok = 0;
                }
                u_free(buf);
            }
            fclose(f);
        }
        EVP_MD_CTX_free(vctx);
    }

    if (sig_ok) {
        dl->blocks_received++;
        dl->blocks_validated++;
        DEBUG_INFO(DEBUG_CATEGORY_DEBUG, "%s: BLOCK_DONE block=%d total=%d/%d",
                   MDL_ID, bi, dl->blocks_validated, dl->num_blocks);

        /* send HAVE_BLOCK to supernode */
        md_dl_send_have_block(inst, dl, bi);

        /* check assembly */
        if (dl->blocks_validated == dl->num_blocks) {
            /* assemble file */
            FILE* out = fopen(dl->dest_path, "wb");
            if (out) {
                for (int n = 0; n < dl->num_blocks; n++) {
                    char tmp2[2048];
                    snprintf(tmp2, sizeof(tmp2), "%s.chunk_%d", dl->dest_path, n);
                    FILE* cf = fopen(tmp2, "rb");
                    if (cf) {
                        uint8_t buf[65536];
                        size_t rd;
                        while ((rd = fread(buf, 1, sizeof(buf), cf)) > 0) fwrite(buf, 1, rd, out);
                        fclose(cf);
                        remove(tmp2);
                    }
                }
                fclose(out);
            }
            dl->assembled = 1;
            dl->err = 0;
            if (dl->done_cb) dl->done_cb(dl->done_arg, dl->err);
            DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "%s: file assembled: %s", MDL_ID, dl->dest_path);
        }
    } else {
        /* invalid signature — mark block for retry from another peer */
        DEBUG_WARN(DEBUG_CATEGORY_GENERAL, "%s: BLOCK_DONE sig fail block=%d — retry", MDL_ID, bi);
        char tmp[2048];
        snprintf(tmp, sizeof(tmp), "%s.chunk_%d", dl->dest_path, bi);
        remove(tmp);
        /* reassign to another peer */
        for (int pi = 0; pi < dl->num_peers; pi++) {
            for (int bj = 0; bj < dl->peers[pi].num_blocks; bj++) {
                if (memcmp(dl->peers[pi].blocks[bj].block_id, dl->block_ids + bi * 16, 16) == 0) {
                    dl->peers[pi].blocks[bj].started = 0;
                    dl->peers[pi].blocks[bj].received = 0;
                    dl->peers[pi].blocks[bj].validated = 0;
                }
            }
        }
    }
}

/* ── public API ── */

int media_download_start(struct UTUN_INSTANCE* inst, uint64_t group_id,
                         const struct media_index_result* result,
                         const char* dest_path, const char* media_base,
                         void (*done_cb)(void* arg, int err), void* done_arg) {
    if (!inst || !result || !dest_path || !media_base || !done_cb) return -1;
    struct media_delivery_ctx* md = &inst->md;
    if (!md->initialized) return -1;

    struct media_download* dl = u_calloc(1, sizeof(*dl));
    if (!dl) return -1;

    memcpy(dl->media_id, result->media_id, 16);
    dl->group_id = group_id;
    snprintf(dl->dest_path, sizeof(dl->dest_path), "%s", dest_path);
    snprintf(dl->media_base, sizeof(dl->media_base), "%s", media_base);
    dl->num_blocks = result->num_blocks;
    dl->file_size = result->file_size;
    dl->block_size = result->block_size;
    memcpy(dl->content_hash, result->content_hash, 32);
    dl->active = 1;
    dl->done_cb = done_cb;
    dl->done_arg = done_arg;

    dl->block_ids = u_malloc((size_t)dl->num_blocks * 16);
    dl->block_sigs = u_malloc((size_t)dl->num_blocks * 64);
    if (!dl->block_ids || !dl->block_sigs) { u_free(dl->block_ids); u_free(dl); return -1; }
    memcpy(dl->block_ids, result->block_ids, (size_t)dl->num_blocks * 16);
    memcpy(dl->block_sigs, result->block_sigs, (size_t)dl->num_blocks * 64);

    /* register in downloads queue */
    if (!md->downloads) {
        md->downloads = queue_new(NULL, 256, 0, 16, "md_dl");
        if (!md->downloads) { md_dl_free(dl); u_free(dl); return -1; }
    }
    memcpy(dl->ll.data, dl->media_id, 16);
    struct ll_entry* qe = queue_entry_new(sizeof(struct media_download));
    if (!qe) { md_dl_free(dl); u_free(dl); return -1; }
    memcpy(qe->data, dl, sizeof(*dl));
    u_free(dl);
    queue_data_put_with_index(md->downloads, qe);

    dl = (struct media_download*)qe->data;

    md_dl_collect_supernodes(inst, dl);
    md_dl_send_query(inst, dl);

    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "%s: download started media=%02x%02x... blocks=%d supers=%d",
               MDL_ID, dl->media_id[0], dl->media_id[1], dl->num_blocks, dl->super_count);
    return 0;
}

int media_download_cancel(struct UTUN_INSTANCE* inst,
                          const uint8_t* media_id, const uint8_t* block_id, uint32_t chunk) {
    (void)chunk;
    struct media_download* dl = md_dl_find(inst, media_id);
    if (!dl) return -1;

    /* send CANCEL to all connected peers */
    struct media_pkt_cancel can;
    memset(&can, 0, sizeof(can));
    can.subcmd = MEDIA_SUBCMD_CANCEL;
    memcpy(can.media_id, media_id, 16);
    if (block_id) memcpy(can.block_id, block_id, 16);
    can.chunk = chunk;

    for (int pi = 0; pi < dl->num_peers; pi++) {
        if (dl->peers[pi].connected) {
            md_dl_send(inst, dl->peers[pi].node_id, (const uint8_t*)&can, sizeof(can));
        }
    }

    dl->active = 0;
    dl->err = -2;
    if (dl->done_cb) dl->done_cb(dl->done_arg, dl->err);

    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "%s: download cancelled", MDL_ID);
    return 0;
}