/* zip.c — zip 容器读取器实现，见 zip.h 的用法/线程/安全说明。 */
#define _POSIX_C_SOURCE 200809L
#define _FILE_OFFSET_BITS 64

#include "zip.h"

#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <zlib.h>

#define MZ_SIG_LOCAL    0x04034b50u
#define MZ_SIG_CD       0x02014b50u
#define MZ_SIG_EOCD     0x06054b50u
#define MZ_LOCAL_MIN    30u
#define MZ_CD_MIN       46u
#define MZ_EOCD_MIN     22u
#define MZ_COMMENT_MAX  65535u
#define MZ_IN_CHUNK     65536u
#define MZ_PREFIX_IN_CHUNK 4096u   /* prefix 读的栈缓冲（M32，见 inflate_entry） */
#define MZ_CRC_CHUNK    (1u << 24)

struct mz_archive {
    FILE *fp;
    uint64_t file_size;
    int count;
    mz_entry *entries;
    char *names;
};

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

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

static int u64_add(uint64_t a, uint64_t b, uint64_t *out)
{
    if (a > UINT64_MAX - b) return 0;
    *out = a + b;
    return 1;
}

static int seek_to(FILE *fp, uint64_t off)
{
    off_t o = (off_t)off;
    if ((uint64_t)o != off) return MZ_ERR_UNSUPPORTED;
    return fseeko(fp, o, SEEK_SET) == 0 ? MZ_OK : MZ_ERR_IO;
}

static int read_exact(FILE *fp, void *buf, size_t len)
{
    if (fread(buf, 1, len, fp) == len) return MZ_OK;
    return ferror(fp) ? MZ_ERR_IO : MZ_ERR_TRUNCATED;
}

static int read_at(FILE *fp, uint64_t off, void *buf, size_t len)
{
    int rc = seek_to(fp, off);
    if (rc != MZ_OK) return rc;
    return read_exact(fp, buf, len);
}

static int get_file_size(FILE *fp, uint64_t *size)
{
    if (fseeko(fp, 0, SEEK_END) != 0) return MZ_ERR_IO;
    off_t end = ftello(fp);
    if (end < 0) return MZ_ERR_IO;
    *size = (uint64_t)end;
    return MZ_OK;
}

/* 从文件尾向前找 EOCD 签名；只接受注释长度不越过文件尾的候选。 */
static int find_eocd(FILE *fp, uint64_t file_size, uint64_t *eocd_off,
                     uint8_t eocd[MZ_EOCD_MIN])
{
    if (file_size < MZ_EOCD_MIN) return MZ_ERR_TRUNCATED;

    uint64_t tail_len = file_size;
    if (tail_len > (uint64_t)MZ_EOCD_MIN + MZ_COMMENT_MAX)
        tail_len = (uint64_t)MZ_EOCD_MIN + MZ_COMMENT_MAX;
    uint64_t tail_off = file_size - tail_len;

    uint8_t *tail = malloc((size_t)tail_len);
    if (!tail) return MZ_ERR_MEM;
    int rc = read_at(fp, tail_off, tail, (size_t)tail_len);
    if (rc == MZ_OK) {
        rc = MZ_ERR_FORMAT;
        for (uint64_t i = tail_len - MZ_EOCD_MIN + 1; i-- > 0; ) {
            if (rd32(tail + i) != MZ_SIG_EOCD) continue;
            uint64_t pos = tail_off + i;
            uint64_t end = 0;
            if (!u64_add(pos, MZ_EOCD_MIN, &end)) continue;
            if (!u64_add(end, rd16(tail + i + 20), &end)) continue;
            if (end > file_size) continue;
            memcpy(eocd, tail + i, MZ_EOCD_MIN);
            *eocd_off = pos;
            rc = MZ_OK;
            break;
        }
    }
    free(tail);
    return rc;
}

/* 解析 central directory；名字先指向 cd 缓冲，成功后在 names_out 里做一份连续副本。 */
static int parse_cd(const uint8_t *cd, uint32_t cd_size, int declared,
                    mz_entry *entries, char **names_out)
{
    const uint8_t *p = cd;
    const uint8_t *end = cd + cd_size;
    uint64_t names_total = 0;
    int n = 0;
    int rc = MZ_OK;

    while (n < declared) {
        if ((size_t)(end - p) < MZ_CD_MIN || rd32(p) != MZ_SIG_CD) {
            rc = MZ_ERR_FORMAT;
            break;
        }
        uint16_t flags = rd16(p + 8);
        uint16_t method = rd16(p + 10);
        uint32_t crc = rd32(p + 16);
        uint32_t comp = rd32(p + 20);
        uint32_t uncomp = rd32(p + 24);
        uint16_t nlen = rd16(p + 28);
        uint16_t elen = rd16(p + 30);
        uint16_t clen = rd16(p + 32);
        uint16_t disk = rd16(p + 34);
        uint32_t loff = rd32(p + 42);
        uint64_t rec = (uint64_t)MZ_CD_MIN + nlen + elen + clen;

        if (rec > (uint64_t)(end - p)) {
            rc = MZ_ERR_FORMAT;
            break;
        }
        if (disk != 0) {
            rc = MZ_ERR_UNSUPPORTED;
            break;
        }
        if (comp == 0xFFFFFFFFu || uncomp == 0xFFFFFFFFu || loff == 0xFFFFFFFFu) {
            rc = MZ_ERR_UNSUPPORTED;
            break;
        }

        mz_entry *e = &entries[n];
        e->name = (const char *)(p + MZ_CD_MIN);
        e->name_len = nlen;
        e->method = method;
        e->flags = flags;
        e->crc = crc;
        e->comp_size = comp;
        e->uncomp_size = uncomp;
        e->local_offset = loff;

        names_total += (uint64_t)nlen + 1;
        n++;
        p += (size_t)rec;
    }

    if (rc == MZ_OK) {
        if (n != declared) {
            rc = MZ_ERR_FORMAT;
        } else if (names_total > (uint64_t)SIZE_MAX) {
            rc = MZ_ERR_MEM;
        } else {
            char *arena = malloc((size_t)names_total);
            if (!arena) {
                rc = MZ_ERR_MEM;
            } else {
                char *w = arena;
                for (int i = 0; i < n; i++) {
                    memcpy(w, entries[i].name, entries[i].name_len);
                    w[entries[i].name_len] = '\0';
                    entries[i].name = w;
                    w += entries[i].name_len + 1;
                }
                *names_out = arena;
            }
        }
    }
    return rc;
}

static int load_cd(mz_archive *z, uint64_t cd_off, uint32_t cd_size, int declared)
{
    uint8_t *cd = malloc(cd_size);
    if (!cd) return MZ_ERR_MEM;
    z->entries = calloc((size_t)declared, sizeof *z->entries);
    if (!z->entries) {
        free(cd);
        return MZ_ERR_MEM;
    }

    int rc = read_at(z->fp, cd_off, cd, cd_size);
    if (rc == MZ_OK) rc = parse_cd(cd, cd_size, declared, z->entries, &z->names);
    if (rc == MZ_OK) {
        z->count = declared;
        for (int i = 0; i < declared; i++) {
            uint64_t end = 0;
            if (!u64_add(z->entries[i].local_offset, MZ_LOCAL_MIN, &end) ||
                end > z->file_size) {
                rc = MZ_ERR_TRUNCATED;
                break;
            }
        }
    }
    free(cd);
    return rc;
}

int mz_open(const char *path, mz_archive **out)
{
    if (!out) return MZ_ERR_ARG;
    *out = NULL;
    if (!path) return MZ_ERR_ARG;

    mz_archive *z = calloc(1, sizeof *z);
    if (!z) return MZ_ERR_MEM;
    z->fp = fopen(path, "rb");
    if (!z->fp) {
        free(z);
        return MZ_ERR_IO;
    }

    int rc = get_file_size(z->fp, &z->file_size);
    uint8_t eocd[MZ_EOCD_MIN];
    uint64_t eocd_off = 0;
    uint16_t n_total = 0;
    uint32_t cd_size = 0;
    uint32_t cd_off = 0;

    if (rc == MZ_OK) rc = find_eocd(z->fp, z->file_size, &eocd_off, eocd);
    if (rc == MZ_OK) {
        uint16_t disk_no = rd16(eocd + 4);
        uint16_t cd_disk = rd16(eocd + 6);
        uint16_t n_disk = rd16(eocd + 8);
        n_total = rd16(eocd + 10);
        cd_size = rd32(eocd + 12);
        cd_off = rd32(eocd + 16);

        if (disk_no != 0 || cd_disk != 0) {
            rc = MZ_ERR_UNSUPPORTED;
        } else if (n_total == 0xFFFFu || n_disk == 0xFFFFu ||
                   cd_size == 0xFFFFFFFFu || cd_off == 0xFFFFFFFFu) {
            rc = MZ_ERR_UNSUPPORTED;  /* zip64：EOCD 字段饱和 */
        } else if (n_disk != n_total) {
            rc = MZ_ERR_FORMAT;
        } else {
            uint64_t cd_end = 0;
            if (!u64_add(cd_off, cd_size, &cd_end) || cd_end > eocd_off) {
                rc = MZ_ERR_TRUNCATED;
            } else if (n_total > 0 && (uint64_t)cd_size < (uint64_t)n_total * MZ_CD_MIN) {
                rc = MZ_ERR_FORMAT;
            } else if ((uint64_t)cd_size > (uint64_t)SIZE_MAX) {
                rc = MZ_ERR_UNSUPPORTED;
            }
        }
    }
    if (rc == MZ_OK && n_total > 0) rc = load_cd(z, cd_off, cd_size, n_total);

    if (rc != MZ_OK) {
        free(z->entries);
        free(z->names);
        fclose(z->fp);
        free(z);
        return rc;
    }
    *out = z;
    return MZ_OK;
}

void mz_close(mz_archive *z)
{
    if (!z) return;
    if (z->fp) fclose(z->fp);
    free(z->entries);
    free(z->names);
    free(z);
}

int mz_count(const mz_archive *z)
{
    return z ? z->count : 0;
}

const mz_entry *mz_entry_at(const mz_archive *z, int index)
{
    if (!z || index < 0 || index >= z->count) return NULL;
    return &z->entries[index];
}

const mz_entry *mz_find(const mz_archive *z, const char *name)
{
    if (!z || !name) return NULL;
    for (int i = 0; i < z->count; i++) {
        if (strcmp(z->entries[i].name, name) == 0) return &z->entries[i];
    }
    return NULL;
}

int mz_entry_is_dir(const mz_entry *e)
{
    return e && e->name_len > 0 && e->name[e->name_len - 1] == '/';
}

static uint32_t crc_buf(const uint8_t *p, size_t n)
{
    uLong c = crc32(0L, Z_NULL, 0);
    while (n > 0) {
        uInt chunk = n > MZ_CRC_CHUNK ? MZ_CRC_CHUNK : (uInt)n;
        c = crc32(c, p, chunk);
        p += chunk;
        n -= chunk;
    }
    return (uint32_t)c;
}

/* 从 data_off 流式 inflate comp_size 字节到 buf（容量 buf_cap = 声明解压尺寸）。
 * 输出多于声明尺寸或流提前结束都算 FORMAT/TRUNCATED，不整块读取压缩数据。
 * prefix=1（M32）：装满 buf_cap 就停——只取前缀，不把整条目解压完；此时不判 CRC
 * （调用方没读完整条目，见 mz_read_prefix）。prefix=0 的语义与旧实现逐字一致。 */
static int inflate_entry(FILE *fp, uint64_t data_off, uint64_t comp_size,
                         uint8_t *buf, size_t buf_cap, size_t *out_len, int prefix)
{
    z_stream z;
    memset(&z, 0, sizeof z);
    if (inflateInit2(&z, -MAX_WBITS) != Z_OK) return MZ_ERR_MEM;

    /* prefix 模式最多只要 buf_cap 字节（4KB 文件头），输入侧用 4KB 栈缓冲：
     * 不为一次前缀读分配 64KB（500 页排版 = 500 次 64KB 分配/释放）。 */
    uint8_t in_stack[MZ_PREFIX_IN_CHUNK];
    size_t in_chunk = prefix ? sizeof in_stack : MZ_IN_CHUNK;
    uint8_t *in = prefix ? in_stack : (uint8_t *)malloc(MZ_IN_CHUNK);
    if (!in) {
        inflateEnd(&z);
        return MZ_ERR_MEM;
    }

    int rc = seek_to(fp, data_off);
    size_t produced = 0;
    uint64_t consumed = 0;
    uint8_t scratch = 0;
    int stalls = 0;

    while (rc == MZ_OK) {
        if (prefix && produced == buf_cap) break;   /* 前缀够了：不再解压 */
        if (z.avail_in == 0) {
            if (consumed >= comp_size) {
                rc = MZ_ERR_TRUNCATED;
                break;
            }
            uint64_t remain = comp_size - consumed;
            size_t want = remain > in_chunk ? in_chunk : (size_t)remain;
            rc = read_exact(fp, in, want);
            if (rc != MZ_OK) break;
            z.next_in = in;
            z.avail_in = (uInt)want;
            consumed += want;
        }

        uInt avail;
        if (produced < buf_cap) {
            size_t room = buf_cap - produced;
            z.next_out = buf + produced;
            avail = room > UINT_MAX ? UINT_MAX : (uInt)room;
        } else {
            z.next_out = &scratch;  /* 声明尺寸之外的输出 → 直接判坏流 */
            avail = 1;
        }
        z.avail_out = avail;

        uInt in_before = z.avail_in;
        int zr = inflate(&z, Z_NO_FLUSH);
        size_t made = (size_t)avail - (size_t)z.avail_out;

        if (made > buf_cap - produced) {
            rc = MZ_ERR_FORMAT;
            break;
        }
        produced += made;

        if (zr == Z_STREAM_END) break;
        if (zr == Z_OK) {
            if (made == 0 && z.avail_in == in_before) {
                if (++stalls > 8 || (consumed >= comp_size && z.avail_in == 0)) {
                    rc = MZ_ERR_TRUNCATED;
                    break;
                }
            } else {
                stalls = 0;
            }
            continue;
        }
        if (zr == Z_BUF_ERROR && z.avail_in == 0 && consumed >= comp_size) {
            rc = MZ_ERR_TRUNCATED;
            break;
        }
        rc = MZ_ERR_FORMAT;
        break;
    }

    inflateEnd(&z);
    if (!prefix) free(in);
    if (rc == MZ_OK) *out_len = produced;
    return rc;
}

int mz_read(mz_archive *z, const mz_entry *e, void *buf, size_t buf_size,
            size_t *out_size)
{
    if (out_size) *out_size = 0;
    if (!z || !z->fp || !e) return MZ_ERR_ARG;
    if (!buf && buf_size > 0) return MZ_ERR_ARG;
    if (e->flags & 1u) return MZ_ERR_UNSUPPORTED;
    if (e->method != 0 && e->method != 8) return MZ_ERR_UNSUPPORTED;
    if (e->uncomp_size > (uint64_t)SIZE_MAX) return MZ_ERR_UNSUPPORTED;

    size_t need = (size_t)e->uncomp_size;
    if (need > buf_size) return MZ_ERR_BUF;

    uint8_t lh[MZ_LOCAL_MIN];
    int rc = read_at(z->fp, e->local_offset, lh, sizeof lh);
    if (rc != MZ_OK) return rc;
    if (rd32(lh) != MZ_SIG_LOCAL) return MZ_ERR_FORMAT;
    if (rd16(lh + 8) != e->method) return MZ_ERR_FORMAT;

    uint64_t data_off = 0;
    uint64_t data_end = 0;
    if (!u64_add(e->local_offset, MZ_LOCAL_MIN, &data_off) ||
        !u64_add(data_off, rd16(lh + 26), &data_off) ||
        !u64_add(data_off, rd16(lh + 28), &data_off) ||
        !u64_add(data_off, e->comp_size, &data_end) ||
        data_end > z->file_size) {
        return MZ_ERR_TRUNCATED;
    }

    if (e->method == 0) {
        if (e->comp_size != e->uncomp_size) return MZ_ERR_FORMAT;
        if (need > 0) {
            rc = read_at(z->fp, data_off, buf, need);
            if (rc != MZ_OK) return rc;
        }
        if (crc_buf((const uint8_t *)buf, need) != e->crc) return MZ_ERR_CRC;
    } else {
        size_t made = 0;
        rc = inflate_entry(z->fp, data_off, e->comp_size, (uint8_t *)buf, need, &made, 0);
        if (rc != MZ_OK) return rc;
        if (made != need) return MZ_ERR_FORMAT;
        if (crc_buf((const uint8_t *)buf, made) != e->crc) return MZ_ERR_CRC;
    }

    if (out_size) *out_size = need;
    return MZ_OK;
}

int mz_read_alloc(mz_archive *z, const mz_entry *e, void **out_buf,
                  size_t *out_size)
{
    if (!out_buf || !out_size) return MZ_ERR_ARG;
    *out_buf = NULL;
    *out_size = 0;
    if (!z || !e) return MZ_ERR_ARG;
    if (e->uncomp_size > (uint64_t)SIZE_MAX) return MZ_ERR_UNSUPPORTED;

    size_t need = (size_t)e->uncomp_size;
    uint8_t *p = malloc(need ? need : 1);
    if (!p) return MZ_ERR_MEM;

    size_t n = 0;
    int rc = mz_read(z, e, p, need, &n);
    if (rc != MZ_OK) {
        free(p);
        return rc;
    }
    *out_buf = p;
    *out_size = n;
    return MZ_OK;
}

int mz_read_prefix(mz_archive *z, const mz_entry *e, void *buf, size_t buf_size,
                   size_t *out_size)
{
    if (out_size) *out_size = 0;
    if (!z || !z->fp || !e) return MZ_ERR_ARG;
    if (!buf && buf_size > 0) return MZ_ERR_ARG;
    if (buf_size == 0) return MZ_OK;
    if (e->flags & 1u) return MZ_ERR_UNSUPPORTED;
    if (e->method != 0 && e->method != 8) return MZ_ERR_UNSUPPORTED;
    if (e->uncomp_size > (uint64_t)SIZE_MAX) return MZ_ERR_UNSUPPORTED;

    /* 与 mz_read 同一套局部头/数据范围校验（结构非法一律拒绝）。 */
    uint8_t lh[MZ_LOCAL_MIN];
    int rc = read_at(z->fp, e->local_offset, lh, sizeof lh);
    if (rc != MZ_OK) return rc;
    if (rd32(lh) != MZ_SIG_LOCAL) return MZ_ERR_FORMAT;
    if (rd16(lh + 8) != e->method) return MZ_ERR_FORMAT;

    uint64_t data_off = 0;
    uint64_t data_end = 0;
    if (!u64_add(e->local_offset, MZ_LOCAL_MIN, &data_off) ||
        !u64_add(data_off, rd16(lh + 26), &data_off) ||
        !u64_add(data_off, rd16(lh + 28), &data_off) ||
        !u64_add(data_off, e->comp_size, &data_end) ||
        data_end > z->file_size) {
        return MZ_ERR_TRUNCATED;
    }

    size_t want = buf_size;
    if (e->uncomp_size < (uint64_t)want) want = (size_t)e->uncomp_size;

    if (e->method == 0) {
        if (e->comp_size != e->uncomp_size) return MZ_ERR_FORMAT;
        if (want > 0) {
            rc = read_at(z->fp, data_off, buf, want);
            if (rc != MZ_OK) return rc;
        }
    } else {
        size_t made = 0;
        rc = inflate_entry(z->fp, data_off, e->comp_size, (uint8_t *)buf, want, &made, 1);
        if (rc != MZ_OK) return rc;
        if (made != want) return MZ_ERR_FORMAT;   /* 声明尺寸内提前结束 = 坏流 */
    }

    if (out_size) *out_size = want;
    return MZ_OK;
}

const char *mz_strerror(int err)
{
    switch (err) {
    case MZ_OK:              return "ok";
    case MZ_ERR_IO:          return "I/O error";
    case MZ_ERR_FORMAT:      return "malformed archive";
    case MZ_ERR_TRUNCATED:   return "archive truncated or offset out of bounds";
    case MZ_ERR_UNSUPPORTED: return "unsupported zip feature";
    case MZ_ERR_BUF:         return "output buffer too small";
    case MZ_ERR_MEM:         return "out of memory";
    case MZ_ERR_ARG:         return "invalid argument";
    case MZ_ERR_CRC:         return "CRC mismatch";
    default:                 return "unknown error";
    }
}