/* grib2.c : décodeur GRIB2 (voir grib2.h). Références : manuel des codes OMM n° 306,
 * FM 92 GRIB édition 2, sections 0 à 8 et modèles 3.0, 4.0, 4.8, 5.0, 5.3, 5.42, 6.0, 7.x. */
#include "grib2.h"

#include <libaec.h>
#include <math.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "util.h"

/* --------------------------------------------------------------- lecture d'octets */
/* Les entiers GRIB2 sont gros-boutistes ; les entiers signés sont en signe + magnitude
 * (bit de poids fort = signe), pas en complément à deux. */
static uint32_t u8(const uint8_t *p) { return p[0]; }
static uint32_t u16(const uint8_t *p) { return (uint32_t)p[0] << 8 | p[1]; }
static uint32_t u32(const uint8_t *p) { return (uint32_t)p[0] << 24 | (uint32_t)p[1] << 16 | (uint32_t)p[2] << 8 | p[3]; }
static uint64_t u64(const uint8_t *p) { return (uint64_t)u32(p) << 32 | u32(p + 4); }
static int32_t s16(const uint8_t *p) { int32_t v = (int32_t)(u16(p) & 0x7fff); return (p[0] & 0x80) ? -v : v; }
static int32_t s32(const uint8_t *p) { int32_t v = (int32_t)(u32(p) & 0x7fffffffu); return (p[0] & 0x80) ? -v : v; }
static float f32(const uint8_t *p) { uint32_t b = u32(p); float f; memcpy(&f, &b, 4); return f; }

/* entier signé (signe + magnitude) sur n octets, pour les descripteurs du 5.3 */
static int64_t sn(const uint8_t *p, int n)
{
    uint64_t v = 0;
    for (int k = 0; k < n; k++)
        v = v << 8 | p[k];
    uint64_t sign = (uint64_t)1 << (8 * n - 1);
    return (v & sign) ? -(int64_t)(v & (sign - 1)) : (int64_t)v;
}

/* --------------------------------------------------------------- lecteur de bits */
typedef struct {
    const uint8_t *p;
    size_t nbits, pos;
} Bits;

/* lit n bits (0..32) ; retourne -1 si le tampon est dépassé */
static int bits_get(Bits *b, unsigned n, uint32_t *out)
{
    if (n > 32 || b->pos + n > b->nbits)
        return -1;
    uint64_t v = 0;
    size_t pos = b->pos;
    unsigned left = n;
    while (left) {
        unsigned off = (unsigned)(pos & 7), take = 8 - off;
        if (take > left)
            take = left;
        unsigned c = (b->p[pos >> 3] >> (8 - off - take)) & ((1u << take) - 1);
        v = v << take | c;
        pos += take;
        left -= take;
    }
    b->pos = pos;
    *out = (uint32_t)v;
    return 0;
}

static void bits_align(Bits *b) { b->pos = (b->pos + 7) & ~(size_t)7; }

static int fail(char *err, size_t errlen, const char *fmt, ...) __attribute__((format(printf, 3, 4)));
static int fail(char *err, size_t errlen, const char *fmt, ...)
{
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(err, errlen, fmt, ap);
    va_end(ap);
    return -1;
}

/* durée selon l'unité de la table 4.4, en secondes */
static int64_t unit_seconds(int unit)
{
    switch (unit) {
    case 0: return 60;
    case 1: return 3600;
    case 2: return 86400;
    case 10: return 3 * 3600;
    case 11: return 6 * 3600;
    case 12: return 12 * 3600;
    case 13: return 1;
    default: return -1;
    }
}

/* --------------------------------------------------------------- dépaquetage */
typedef struct {
    float R;
    int E, D, bits, drt;
    size_t ndata;               /* valeurs codées (sans les points masqués par le bitmap) */
    /* 5.42 */
    int ccsds_flags, block_size, rsi;
    /* 5.3 */
    int group_split, missing_mgmt, ng, ref_width, bits_width, len_incr, bits_len, order, extra;
    uint32_t ref_len, last_len;
} Pack;

static int unpack_simple(const Pack *k, const uint8_t *d, size_t dlen, uint32_t *x, char *err, size_t el)
{
    if (k->bits == 0) {
        memset(x, 0, k->ndata * sizeof *x);
        return 0;
    }
    if (k->bits > 32)
        return fail(err, el, "5.0 : %d bits par valeur", k->bits);
    Bits b = {d, dlen * 8, 0};
    for (size_t n = 0; n < k->ndata; n++)
        if (bits_get(&b, (unsigned)k->bits, &x[n]))
            return fail(err, el, "5.0 : données tronquées");
    return 0;
}

static int unpack_ccsds(const Pack *k, const uint8_t *d, size_t dlen, uint32_t *x, char *err, size_t el)
{
    if (k->bits == 0) {
        memset(x, 0, k->ndata * sizeof *x);
        return 0;
    }
    if (k->bits > 32 || k->block_size <= 0 || k->rsi <= 0)
        return fail(err, el, "5.42 : paramètres invalides");
    /* échantillons en sortie : 1, 2 ou 4 octets, ordre natif (on retire les drapeaux MSB et 3BYTE) */
    size_t nb = (size_t)(k->bits + 7) / 8;
    if (nb == 3)
        nb = 4;
    if (k->ndata > SIZE_MAX / nb)
        return fail(err, el, "5.42 : taille");
    uint8_t *out = xmalloc(k->ndata, nb);
    struct aec_stream s;
    memset(&s, 0, sizeof s);
    s.next_in = d;
    s.avail_in = dlen;
    s.next_out = out;
    s.avail_out = k->ndata * nb;
    s.bits_per_sample = (unsigned)k->bits;
    s.block_size = (unsigned)k->block_size;
    s.rsi = (unsigned)k->rsi;
    s.flags = (unsigned)k->ccsds_flags & ~(unsigned)(AEC_DATA_3BYTE | AEC_DATA_MSB);
    int rc = aec_buffer_decode(&s);
    if (rc != AEC_OK || s.total_out != k->ndata * nb) {
        free(out);
        return fail(err, el, "5.42 : décompression CCSDS en échec (%d)", rc);
    }
    for (size_t n = 0; n < k->ndata; n++) {
        if (nb == 1)
            x[n] = out[n];
        else if (nb == 2) {
            uint16_t v;
            memcpy(&v, out + 2 * n, 2);
            x[n] = v;
        } else
            memcpy(&x[n], out + 4 * n, 4);
    }
    free(out);
    return 0;
}

/* 5.3 : valeurs en groupes (référence, largeur, longueur par groupe), puis différenciation
 * spatiale d'ordre 1 ou 2 à défaire sur la suite des valeurs non manquantes.
 * x reçoit les entiers (après différenciation), miss[n] = 1 si la valeur est manquante. */
static int unpack_complex(const Pack *k, const uint8_t *d, size_t dlen, int64_t *x, uint8_t *miss, char *err, size_t el)
{
    if (k->order < 1 || k->order > 2 || k->extra < 1 || k->extra > 4)
        return fail(err, el, "5.3 : différenciation non gérée");
    if (k->ng <= 0 || (size_t)k->ng > k->ndata || k->bits > 32 || k->bits_width > 32 || k->bits_len > 32)
        return fail(err, el, "5.3 : groupes invalides");
    size_t need = (size_t)k->extra * (size_t)(k->order + 1);
    if (dlen < need)
        return fail(err, el, "5.3 : descripteurs tronqués");
    int64_t ival1 = sn(d, k->extra), ival2 = k->order == 2 ? sn(d + k->extra, k->extra) : 0;
    int64_t minsd = sn(d + (size_t)k->extra * (size_t)k->order, k->extra);
    Bits b = {d + need, (dlen - need) * 8, 0};
    size_t ng = (size_t)k->ng;
    uint32_t *gref = xmalloc(ng, 4), *gw = xmalloc(ng, 4), *gl = xmalloc(ng, 4);
    int rc = -1;
    for (size_t g = 0; g < ng; g++)
        if (bits_get(&b, (unsigned)k->bits, &gref[g]))
            goto out;
    bits_align(&b);
    for (size_t g = 0; g < ng; g++) {
        if (bits_get(&b, (unsigned)k->bits_width, &gw[g]))
            goto out;
        gw[g] += (uint32_t)k->ref_width;
        if (gw[g] > 32)
            goto out;
    }
    bits_align(&b);
    size_t total = 0;
    for (size_t g = 0; g < ng; g++) {
        uint32_t v;
        if (bits_get(&b, (unsigned)k->bits_len, &v))
            goto out;
        gl[g] = g + 1 == ng ? k->last_len : k->ref_len + v * (uint32_t)k->len_incr;
        total += gl[g];
        if (total > k->ndata)
            goto out;
    }
    if (total != k->ndata)
        goto out;
    bits_align(&b);
    uint32_t allbits = k->bits ? (uint32_t)((1ull << k->bits) - 1) : 0;
    size_t n = 0;
    for (size_t g = 0; g < ng; g++) {
        uint32_t w = gw[g], wall = w ? (uint32_t)((1ull << w) - 1) : 0;
        for (uint32_t m = 0; m < gl[g]; m++, n++) {
            uint32_t v = 0;
            miss[n] = 0;
            if (w) {
                if (bits_get(&b, w, &v))
                    goto out;
                if (k->missing_mgmt >= 1 && v == wall)
                    miss[n] = 1;
                else if (k->missing_mgmt == 2 && v == wall - 1)
                    miss[n] = 1;
                x[n] = (int64_t)v + gref[g];
            } else {
                if (k->missing_mgmt >= 1 && gref[g] == allbits)
                    miss[n] = 1;
                else if (k->missing_mgmt == 2 && gref[g] == allbits - 1)
                    miss[n] = 1;
                x[n] = gref[g];
            }
        }
    }
    /* différenciation spatiale, sur les seules valeurs présentes */
    int64_t p1 = 0, p2 = 0;
    size_t seen = 0;
    for (n = 0; n < k->ndata; n++) {
        if (miss[n])
            continue;
        int64_t v;
        if (seen == 0)
            v = ival1;
        else if (seen == 1 && k->order == 2)
            v = ival2;
        else if (k->order == 1)
            v = x[n] + minsd + p1;
        else
            v = x[n] + minsd + 2 * p1 - p2;
        p2 = p1;
        p1 = v;
        x[n] = v;
        seen++;
    }
    rc = 0;
out:
    free(gref);
    free(gw);
    free(gl);
    if (rc)
        fail(err, el, "5.3 : données tronquées ou incohérentes");
    return rc;
}

/* --------------------------------------------------------------- message */
void g2_free(G2Field *f)
{
    free(f->values);
    f->values = NULL;
}

int g2_next(const uint8_t *buf, size_t len, size_t *pos, G2Field *f, g2_filter want, void *ctx,
            char *err, size_t el)
{
    memset(f, 0, sizeof *f);
    f->stat_proc = -1;
    /* recherche de "GRIB" (des octets parasites entre messages sont tolérés) */
    size_t p = *pos;
    while (p + 16 <= len && memcmp(buf + p, "GRIB", 4) != 0)
        p++;
    if (p + 16 > len) {
        *pos = len;
        return 0;
    }
    const uint8_t *m = buf + p;
    if (m[7] != 2)
        return *pos = p + 4, fail(err, el, "édition %u non gérée", m[7]);
    uint64_t mlen = u64(m + 8);
    if (mlen < 16 + 4 || mlen > len - p)
        return *pos = len, fail(err, el, "longueur de message invalide");
    *pos = p + (size_t)mlen;
    if (memcmp(m + mlen - 4, "7777", 4) != 0)
        return fail(err, el, "fin de message absente");
    f->discipline = (int)m[6];

    const uint8_t *s5 = NULL, *s6 = NULL, *s7 = NULL;
    size_t l5 = 0, l6 = 0, l7 = 0;
    int have3 = 0, have4 = 0;
    size_t q = 16, end = (size_t)mlen - 4;
    while (q < end) {
        if (end - q < 5)
            return fail(err, el, "section tronquée");
        size_t sl = u32(m + q);
        int sn_ = (int)m[q + 4];
        if (sl < 5 || sl > end - q)
            return fail(err, el, "section %d de longueur invalide", sn_);
        const uint8_t *s = m + q;
        switch (sn_) {
        case 1:
            if (sl < 21)
                return fail(err, el, "section 1 courte");
            /* octets 13-19 : année, mois, jour, heure, minute, seconde (numérotation à partir de 1) */
            f->ref_time = epoch_utc((int)u16(s + 12), (int)s[14], (int)s[15], (int)s[16], (int)s[17], (int)s[18]);
            break;
        case 2:
            break;
        case 3: {
            if (sl < 72 || u16(s + 12) != 0)
                return fail(err, el, "grille non gérée (modèle %u)", sl >= 14 ? u16(s + 12) : 0);
            f->ni = (int)u32(s + 30);
            f->nj = (int)u32(s + 34);
            uint32_t basic = u32(s + 38), sub = u32(s + 42);
            double unit = (basic == 0 || basic == 0xffffffffu || sub == 0xffffffffu) ? 1e-6 : (double)basic / sub;
            f->la1 = s32(s + 46) * unit;
            f->lo1 = s32(s + 50) * unit;
            f->la2 = s32(s + 55) * unit;
            f->lo2 = s32(s + 59) * unit;
            f->di = u32(s + 63) * unit;
            f->dj = u32(s + 67) * unit;
            f->scan = (int)s[71];
            if (f->ni <= 0 || f->nj <= 0 || (uint64_t)f->ni * (uint64_t)f->nj > MAX_POINTS)
                return fail(err, el, "grille trop grande");
            if (f->scan & 0x20)
                return fail(err, el, "balayage par colonnes non géré");
            have3 = 1;
            break;
        }
        case 4: {
            if (sl < 34)
                return fail(err, el, "section 4 courte");
            f->pdt = (int)u16(s + 7);
            if (f->pdt != 0 && f->pdt != 8)
                return fail(err, el, "produit 4.%d non géré", f->pdt);
            f->category = (int)s[9];
            f->number = (int)s[10];
            int unit = (int)s[17];
            int64_t ft = (int64_t)u32(s + 18), us = unit_seconds(unit);
            f->surface = (int)s[22];
            int sc = (int8_t)s[23];
            uint32_t sv = u32(s + 24);
            f->surface_value = sv == 0xffffffffu ? NAN : sv * pow(10.0, -((s[23] & 0x80) ? -(sc & 0x7f) : sc));
            if (f->pdt == 0) {
                if (us < 0)
                    return fail(err, el, "unité de temps %d inconnue", unit);
                f->valid_time = f->ref_time + ft * us;
            } else {
                if (sl < 58)
                    return fail(err, el, "section 4.8 courte");
                /* octets 35-41 : fin de l'intervalle statistique, écrite en clair */
                f->valid_time = epoch_utc((int)u16(s + 34), (int)s[36], (int)s[37], (int)s[38], (int)s[39], (int)s[40]);
                f->stat_proc = (int)s[46];
            }
            have4 = 1;
            break;
        }
        case 5:
            s5 = s;
            l5 = sl;
            break;
        case 6:
            s6 = s;
            l6 = sl;
            break;
        case 7:
            s7 = s;
            l7 = sl;
            break;
        default:
            return fail(err, el, "section %d inattendue", sn_);
        }
        q += sl;
    }
    if (!have3 || !have4 || !s5 || !s6 || !s7)
        return fail(err, el, "message incomplet");
    if (want && !want(f->discipline, f->category, f->number, ctx))
        return 1;

    /* ---- section 5 */
    Pack k;
    memset(&k, 0, sizeof k);
    if (l5 < 21)
        return fail(err, el, "section 5 courte");
    k.ndata = u32(s5 + 5);
    k.drt = (int)u16(s5 + 9);
    k.R = f32(s5 + 11);
    k.E = s16(s5 + 15);
    k.D = s16(s5 + 17);
    k.bits = (int)u8(s5 + 19);
    f->drt = k.drt;
    f->bits = k.bits;
    size_t npts = (size_t)f->ni * (size_t)f->nj;
    if (k.ndata > npts)
        return fail(err, el, "plus de valeurs que de points");
    if (k.drt == 42) {
        if (l5 < 25)
            return fail(err, el, "section 5.42 courte");
        k.ccsds_flags = (int)s5[21];
        k.block_size = (int)s5[22];
        k.rsi = (int)u16(s5 + 23);
    } else if (k.drt == 3) {
        if (l5 < 49)
            return fail(err, el, "section 5.3 courte");
        k.group_split = (int)s5[21];
        k.missing_mgmt = (int)s5[22];
        k.ng = (int)u32(s5 + 31);
        k.ref_width = (int)s5[35];
        k.bits_width = (int)s5[36];
        k.ref_len = u32(s5 + 37);
        k.len_incr = (int)s5[41];
        k.last_len = u32(s5 + 42);
        k.bits_len = (int)s5[46];
        k.order = (int)s5[47];
        k.extra = (int)s5[48];
        if (k.group_split != 1)
            return fail(err, el, "5.3 : découpage de groupes %d non géré", k.group_split);
    } else if (k.drt != 0) {
        return fail(err, el, "emballage 5.%d non géré", k.drt);
    }

    /* ---- section 6 : bitmap */
    const uint8_t *bitmap = NULL;
    if (l6 < 6)
        return fail(err, el, "section 6 courte");
    if (s6[5] == 0) {
        if (l6 - 6 < (npts + 7) / 8)
            return fail(err, el, "bitmap tronqué");
        bitmap = s6 + 6;
    } else if (s6[5] != 255) {
        return fail(err, el, "bitmap %u non géré", s6[5]);
    }
    size_t present = npts;
    if (bitmap) {
        present = 0;
        for (size_t n = 0; n < npts; n++)
            present += (bitmap[n >> 3] >> (7 - (n & 7))) & 1;
    }
    if (present != k.ndata)
        return fail(err, el, "bitmap et nombre de valeurs incohérents");

    /* ---- section 7 : données */
    const uint8_t *d = s7 + 5;
    size_t dlen = l7 - 5;
    double scale_e = ldexp(1.0, k.E), scale_d = pow(10.0, -k.D);
    double *vals = xmalloc(k.ndata ? k.ndata : 1, sizeof(double));
    int rc;
    if (k.drt == 3) {
        int64_t *x = xmalloc(k.ndata ? k.ndata : 1, sizeof(int64_t));
        uint8_t *miss = xmalloc(k.ndata ? k.ndata : 1, 1);
        rc = unpack_complex(&k, d, dlen, x, miss, err, el);
        for (size_t n = 0; rc == 0 && n < k.ndata; n++)
            vals[n] = miss[n] ? NAN : ((double)k.R + (double)x[n] * scale_e) * scale_d;
        free(x);
        free(miss);
    } else {
        uint32_t *x = xmalloc(k.ndata ? k.ndata : 1, sizeof(uint32_t));
        rc = k.drt == 42 ? unpack_ccsds(&k, d, dlen, x, err, el) : unpack_simple(&k, d, dlen, x, err, el);
        for (size_t n = 0; rc == 0 && n < k.ndata; n++)
            vals[n] = ((double)k.R + (double)x[n] * scale_e) * scale_d;
        free(x);
    }
    if (rc) {
        free(vals);
        return -1;
    }
    f->values = xmalloc(npts, sizeof(double));
    if (bitmap) {
        size_t c = 0;
        for (size_t n = 0; n < npts; n++)
            f->values[n] = ((bitmap[n >> 3] >> (7 - (n & 7))) & 1) ? vals[c++] : NAN;
    } else {
        memcpy(f->values, vals, npts * sizeof(double));
    }
    free(vals);
    return 1;
}