/*
 *  $Id: zonfile.c 29809 2026-04-15 11:54:46Z yeti-dn $
 *  Copyright (C) 2023-2026 David Necas (Yeti).
 *
 *  This program is free software; you can redistribute it and/or modify it under the terms of the GNU General Public
 *  License as published by the Free Software Foundation; either version 2 of the License, or (at your option) any
 *  later version.
 *
 *  This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied
 *  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
 *  details.
 *
 *  You should have received a copy of the GNU General Public License along with this program; if not, write to the
 *  Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/**
 * [FILE-MAGIC-USERGUIDE]
 * Keyence ZON data
 * .zon
 * Read
 **/

/**
 * [FILE-MAGIC-FREEDESKTOP]
 * <mime-type type="application/x-zon-profilometry">
 *   <comment>Zon profilometry data</comment>
 *   <glob pattern="*.zon"/>
 *   <glob pattern="*.ZON"/>
 * </mime-type>
 **/

#include "config.h"
#include <glib/gi18n-lib.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <glib/gstdio.h>
#include <gwy.h>

#include "err.h"
#include "gwyzip.h"
#ifdef HAVE_ZSTD
#include "zstd.h"
#endif

#define MAGIC_PK "PK\x03\x04"
#define MAGIC_PK_SIZE (sizeof(MAGIC_PK)-1)

#define MAGIC_ZSTD "\x28\xb5\x2f\xfd"
#define MAGIC_ZSTD_SIZE (sizeof(MAGIC_ZSTD)-1)

#define BLOODY_UTF8_BOM "\xef\xbb\xbf"
#define EXTENSION ".zon"

enum {
    ZON_PK_HEADER_SIZE = 8,
    ZON_BMP_HEADER_SIZE = 54,
    ZON_HEADER_SIZE = ZON_PK_HEADER_SIZE + ZON_BMP_HEADER_SIZE,
};

typedef struct {
    gchar *uuid;
    gchar *freeme;
    guint xres;
    guint yres;
    guint itemsize;
    guint rowstride;
    guchar *data;
} ZONData;

typedef struct {
    GHashTable *hash;
    GArray *datafiles;
    GString *path;
    GString *str;
    gboolean ignore;
#ifdef HAVE_ZSTD
    ZSTD_DCtx *zstd_context;
#endif
} ZONFile;

static gboolean module_register(void);
static gint     detect_file    (const GwyFileDetectInfo *fileinfo,
                                gboolean only_name);
static GwyFile* load_file      (const gchar *filename,
                                GwyRunModeFlags mode,
                                GError **error);
static gboolean parse_one_file (GwyZipFile zipfile,
                                ZONFile *zonfile,
                                GError **error);
static void     free_zonfile   (ZONFile *zonfile);
#ifdef HAVE_ZSTD
static guchar*  decompress_zstd(ZSTD_DCtx *context,
                                guchar *content,
                                gsize compsize,
                                gsize *decompsize,
                                const gchar *filename,
                                        GError **error);
#endif

static GwyModuleInfo module_info = {
    GWY_MODULE_ABI_VERSION,
    &module_register,
    N_("Reads Keyence ZON files."),
    "Yeti <yeti@gwyddion.net>",
    "2.2",
    "David Nečas (Yeti)",
    "2023",
};

GWY_MODULE_QUERY(module_info)

static gboolean
module_register(void)
{
    gwy_file_func_register("zonfile",
                           N_("Keyence ZON data (.zon)"),
                           detect_file, load_file, NULL, NULL);

    return TRUE;
}

/* The weird ZIP + BMP header. Since it is nonstandard, consider it sufficient. */
static gboolean
check_magic_header(const guchar *buffer, gsize size)
{
    if (size < ZON_HEADER_SIZE)
        return FALSE;
    if (memcmp(buffer, "KPK", 3) != 0)
        return FALSE;
    if (memcmp(buffer + 8, "BM", 2) != 0)
        return FALSE;
    return TRUE;
}

static gint
detect_file(const GwyFileDetectInfo *fileinfo, gboolean only_name)
{
    if (only_name)
        return g_str_has_suffix(fileinfo->name_lowercase, EXTENSION) ? 15 : 0;

    if (!check_magic_header(fileinfo->head, fileinfo->buffer_len))
        return 0;

    return 95;
}

static GwyZipFile
open_contained_zip_file(const gchar *filename, const guchar **pbuffer, gsize *psize, GError **error)
{
    GwyZipFile zipfile = NULL;
    GError *err = NULL;
    const guchar *p, *zipstart, *buffer = NULL;
    G_GNUC_UNUSED gsize headersize;
    gsize bmpsize, zipsize, size = 0;

    *pbuffer = NULL;
    *psize = 0;
    if (!gwy_map_file_contents(filename, &buffer, &size, &err)) {
        err_GET_FILE_CONTENTS(error, &err);
        return NULL;
    }
    if (!check_magic_header(buffer, size)) {
        err_FILE_TYPE(error, "ZON");
        goto end;
    }

    p = buffer + 4;
    headersize = gwy_get_guint32_le(&p);

    /* Check if we find the ZIP header after the bitmap. */
    p = buffer + ZON_PK_HEADER_SIZE + 2;  /* 2 for BM **/
    bmpsize = gwy_get_guint32_le(&p);

    gwy_debug("header size %lu, bmp size %lu", (gulong)headersize, (gulong)bmpsize);

    if (ZON_PK_HEADER_SIZE + bmpsize + MAGIC_PK_SIZE > size) {
        err_FILE_TYPE(error, "ZON");
        goto end;
    }

    zipsize = size - ZON_PK_HEADER_SIZE - bmpsize;
    zipstart = buffer + ZON_PK_HEADER_SIZE + bmpsize;
    if (memcmp(zipstart, MAGIC_PK, MAGIC_PK_SIZE) != 0) {
        gwy_debug("header size is wrong, trying to search forward for a PK ZIP signature");
        const guchar *t = gwy_memmem(zipstart, MIN(zipsize, 1024), MAGIC_PK, MAGIC_PK_SIZE);
        if (t) {
            gwy_debug("found a PK ZIP signature %ld bytes after", (glong)(t - zipstart));
            zipsize -= t - zipstart;
            zipstart = t;
        }
        else {
            gwy_debug("did not find a PK ZIP signature");
        }
    }

    gwy_debug("opening ZIP");
    zipfile = gwyzip_make_temporary_archive(zipstart, zipsize, "gwyddion-zonfile-XXXXXX.zip", error);

end:
    if (zipfile) {
        *pbuffer = buffer;
        *psize = size;
    }
    else
        gwy_unmap_file_contents(buffer);

    return zipfile;
}

static GwyFile*
load_file(const gchar *filename,
          G_GNUC_UNUSED GwyRunModeFlags mode,
          GError **error)
{
    GwyFile *file = NULL;
    ZONFile zonfile;
    GwyZipFile zipfile = NULL;
    ZONData *zd;
    gchar *xml_filename;
    const gchar *s;
    const guchar *buffer = NULL;
    gsize size = 0;
    GwyField *field;
    gdouble *d;
    gdouble dxy, q;
    guint i, j, id, mask_i = G_MAXUINT;

    gwy_clear1(zonfile);
    if (!(zipfile = open_contained_zip_file(filename, &buffer, &size, error)))
        goto fail;

    gwy_debug("OK");
    if (!gwyzip_first_file(zipfile, error))
        goto fail;

    zonfile.hash = g_hash_table_new_full(g_str_hash, g_str_equal, g_free, g_free);
    zonfile.datafiles = g_array_new(FALSE, FALSE, sizeof(ZONData));
    zonfile.path = g_string_new(NULL);
    zonfile.str = g_string_new(NULL);

    do {
        if (!gwyzip_get_current_filename(zipfile, &xml_filename, error))
            goto fail;
        gwy_debug("found file: <%s>", xml_filename);
        if (!parse_one_file(zipfile, &zonfile, error))
            goto fail;
        g_free(xml_filename);
    } while (gwyzip_next_file(zipfile, NULL));

    if (!(s = g_hash_table_lookup(zonfile.hash, "/Calibration/XYCalibration/MeterPerPixel"))) {
        err_MISSING_FIELD(error, "/Calibration/XYCalibration/MeterPerPixel");
        goto fail;
    }
    dxy = g_ascii_strtod(s, NULL);
    sanitise_real_size(&dxy, "pixel size");

    if (!(s = g_hash_table_lookup(zonfile.hash, "/Calibration/ZCalibration/MeterPerUnit"))) {
        err_MISSING_FIELD(error, "/Calibration/ZCalibration/MeterPerUnit");
        goto fail;
    }
    q = g_ascii_strtod(s, NULL);

    file = gwy_file_new_in_construction();
    for (i = id = 0; i < zonfile.datafiles->len; i++) {
        zd = &g_array_index(zonfile.datafiles, ZONData, i);
        if (zd->itemsize == 1)
            mask_i = i;
        if (zd->itemsize != 4)
            continue;

        field = gwy_field_new(zd->xres, zd->yres, zd->xres*dxy, zd->yres*dxy, FALSE);
        gwy_unit_set_from_string(gwy_field_get_unit_xy(field), "m");
        gwy_unit_set_from_string(gwy_field_get_unit_z(field), "m");
        d = gwy_field_get_data(field);
        for (j = 0; j < zd->yres; j++) {
            gwy_convert_raw_data(zd->data + j*zd->rowstride, zd->xres, 1,
                                 GWY_RAW_DATA_SINT32, GWY_BYTE_ORDER_LITTLE_ENDIAN, d + j*zd->xres, q, 0.0);
        }
        gwy_file_pass_image(file, id, field);
        gwy_log_add_import(file, GWY_FILE_IMAGE, id, NULL, filename);
        gwy_image_title_fall_back(file, id);
        id++;
    }
    if (mask_i != G_MAXUINT) {
        zd = &g_array_index(zonfile.datafiles, ZONData, mask_i);
        GwyNield *mask = gwy_nield_new(zd->xres, zd->yres);
        gint *m = gwy_nield_get_data(mask);
        for (i = 0; i < zd->yres; i++) {
            for (j = 0; j < zd->xres; j++)
                m[i*zd->xres + j] = !!zd->data[i*zd->rowstride + j];
        }
        for (i = 0; i < id; i++) {
            if (i == id-1)
                gwy_file_set_image_mask(file, i, mask);
            else
                gwy_file_pass_image_mask(file, i, gwy_nield_copy(mask));
        }
        g_object_unref(mask);
    }

    if (!gwy_dict_get_n_items(GWY_DICT(file))) {
        g_clear_object(&file);
        err_NO_DATA(error);
        goto fail;
    }

fail:
    if (zipfile)
        gwyzip_close(zipfile);
    gwy_unmap_file_contents(buffer);
    free_zonfile(&zonfile);

    return file;
}

static void
start_element(G_GNUC_UNUSED GMarkupParseContext *context,
              const gchar *element_name,
              G_GNUC_UNUSED const gchar **attribute_names,
              G_GNUC_UNUSED const gchar **attribute_values,
              gpointer user_data,
              G_GNUC_UNUSED GError **error)
{
    ZONFile *zonfile = (ZONFile*)user_data;

    if (zonfile->ignore)
        return;

    g_string_append_c(zonfile->path, '/');
    g_string_append(zonfile->path, element_name);
    //gwy_debug("%s", zonfile->path->str);
    if (gwy_stramong(zonfile->path->str,
                     "/DataMap", "/ColorPaletteData/Palette", "/DeviceDumpInformations",
                     NULL)) {
        zonfile->ignore = TRUE;
        gwy_debug("found %s, ignoring", zonfile->path->str);
    }
}

static void
end_element(G_GNUC_UNUSED GMarkupParseContext *context,
            const gchar *element_name,
            gpointer user_data,
            G_GNUC_UNUSED GError **error)
{
    ZONFile *zonfile = (ZONFile*)user_data;
    guint n, len = zonfile->path->len;
    gchar *path = zonfile->path->str;

    if (zonfile->ignore)
        return;

    n = strlen(element_name);
    g_return_if_fail(g_str_has_suffix(path, element_name));
    g_return_if_fail(len > n);
    g_return_if_fail(path[len-1 - n] == '/');
    //gwy_debug("%s", path);

    g_string_set_size(zonfile->path, len-1 - n);
}

static void
text(G_GNUC_UNUSED GMarkupParseContext *context,
     const gchar *value,
     G_GNUC_UNUSED gsize value_len,
     gpointer user_data,
     G_GNUC_UNUSED GError **error)
{
    ZONFile *zonfile = (ZONFile*)user_data;
    gchar *path = zonfile->path->str;
    GString *str = zonfile->str;

    if (zonfile->ignore || !strlen(value))
        return;

    g_string_assign(str, value);
    g_strstrip(str->str);
    if (!*(str->str))
        return;

    gwy_debug("%s = <%s>", path, str->str);
    g_hash_table_replace(zonfile->hash, g_strdup(path), g_strdup(str->str));
}

static gboolean
gather_data_file(GwyZipFile zipfile, ZONFile *zonfile,
                 guchar *buffer, gsize size,
                 GError **error)
{
    ZONData zd;
    const guchar *p = buffer;

    if (!gwyzip_get_current_filename(zipfile, &zd.uuid, error))
        return FALSE;
    gwy_debug("non-XML file %s", zd.uuid);

    if (size <= 16) {
        err_TRUNCATED_PART(error, zd.uuid);
        goto fail;
    }

    zd.xres = gwy_get_guint32_le(&p);
    zd.yres = gwy_get_guint32_le(&p);
    zd.itemsize = gwy_get_guint32_le(&p);
    zd.rowstride = gwy_get_guint32_le(&p);
    gwy_debug("xres %u, yres %u, item size %u, rowstride %u", zd.xres, zd.yres, zd.itemsize, zd.rowstride);
    if (err_DIMENSION(error, zd.xres) || err_DIMENSION(error, zd.yres))
        goto fail;
    if (zd.rowstride/zd.itemsize < zd.xres) {
        g_set_error(error, GWY_MODULE_FILE_ERROR, GWY_MODULE_FILE_ERROR_DATA,
                    _("Rows stride %u does not match length and bytes per item (%u×%u)."),
                    zd.rowstride, zd.itemsize, zd.xres);
        goto fail;
    }
    size -= 16;
    if (size % zd.rowstride || size/zd.rowstride != zd.yres) {
        g_set_error(error, GWY_MODULE_FILE_ERROR, GWY_MODULE_FILE_ERROR_DATA,
                    _("Data size %u does not match data dimensions (%u×%u)."),
                    (guint)size, zd.xres, zd.yres);
        goto fail;
    }

    zd.data = buffer + 16;
    zd.freeme = buffer;
    g_array_append_val(zonfile->datafiles, zd);
    return TRUE;

fail:
    g_free(buffer);
    return FALSE;
}

static gboolean
parse_one_file(GwyZipFile zipfile, ZONFile *zonfile, GError **error)
{
    GMarkupParser parser = {
        &start_element, &end_element, &text, NULL, NULL,
    };
    GMarkupParseContext *context = NULL;
    guchar *content = NULL, *s;
    gboolean ok = FALSE;
    GError *err = NULL;
    gsize size;

    if (!(content = gwyzip_get_file_content(zipfile, &size, error)))
        return FALSE;

    /* Sometimes the files inside the ZIP archive are also compressed individually using zstd. Because no matter how
     * messy things are, someone always figures out a way to make them worse. Give an informative error message if
     * this is not supported. */
    if (size >= 4 && memcmp(content, MAGIC_ZSTD, MAGIC_ZSTD_SIZE) == 0) {
#ifdef HAVE_ZSTD
        if (!zonfile->zstd_context)
            zonfile->zstd_context = ZSTD_createDCtx();

        gchar *filename = NULL;
        if (!gwyzip_get_current_filename(zipfile, &filename, NULL))
            filename = g_strdup("???");

        /* It always eats the content passed to it, even on failure. */
        gsize decompsize;
        if (!(content = decompress_zstd(zonfile->zstd_context, content, size, &decompsize, filename, error))) {
            g_free(filename);
            return FALSE;
        }
        g_free(filename);
        size = decompsize;
#else
        g_set_error(error, GWY_MODULE_FILE_ERROR, GWY_MODULE_FILE_ERROR_SPECIFIC,
                    /* TRANSLATORS: %s is replaced by a software library name like zlib. */
                    _("Cannot decompress compressed data.  Gwyddion was built without %s support."), "zstd");
        g_free(content);
        return FALSE;
#endif
    }

    if (g_str_has_prefix(content, "<"))
        s = gwy_strkill(content, "\r");
    else if (g_str_has_prefix(content, BLOODY_UTF8_BOM "<"))
        s = gwy_strkill(content + strlen(BLOODY_UTF8_BOM), "\r");
    else if (size >= MAGIC_PK_SIZE && memcmp(content, MAGIC_PK, MAGIC_PK_SIZE) == 0) {
        gwy_debug("ignoring nested compressed file");
        g_free(content);
        /* Just ignore nested ZIP files. They probably contain other ZIP files inside, which in turn contain stuff
         * already present in the main archive. Or more nested ZIP files. */
        return TRUE;
    }
    else
        return gather_data_file(zipfile, zonfile, content, size, error);

    g_string_truncate(zonfile->path, 0);
    zonfile->ignore = FALSE;
    context = g_markup_parse_context_new(&parser, 0, zonfile, NULL);
    if (!g_markup_parse_context_parse(context, s, -1, &err) || !g_markup_parse_context_end_parse(context, &err))
        err_XML(error, &err);
    else
        ok = TRUE;

    if (context)
        g_markup_parse_context_free(context);
    g_free(content);

    return ok;
}

#ifdef HAVE_ZSTD
static guchar*
decompress_zstd(ZSTD_DCtx *context,
                guchar *content, gsize compsize, gsize *decompsize,
                const gchar *filename, GError **error)
{
    // We must use the streaming API because they compress the files in a way which does not store the size.
    // 4096 bytes should be enough for the various XML bits and pieces.
    ZSTD_inBuffer input = { content, compsize, 0 };
    ZSTD_outBuffer output = { NULL, 4096, 0 };

    gwy_debug("trying to zstd-decompress file %s", filename);
    output.dst = g_new(guchar, output.size);

    while (TRUE) {
        gsize status = ZSTD_decompressStream(context, &output, &input);
        gwy_debug("input size=%zu, pos=%zu; output size=%zu, pos=%zu",
                   input.size, input.pos, output.size, output.pos);
        if (ZSTD_isError(status)) {
            g_set_error(error, GWY_MODULE_FILE_ERROR, GWY_MODULE_FILE_ERROR_SPECIFIC,
                        _("Libzstd error while reading the zst file: %s."), filename);
            g_free(content);
            g_free(output.dst);
            ZSTD_DCtx_reset(context, ZSTD_reset_session_only);
            return NULL;
        }
        if (input.pos == input.size)
            break;

        output.size = 2*output.size;
        output.dst = g_renew(guchar, output.dst, output.size);
    }

    g_free(content);
    ZSTD_DCtx_reset(context, ZSTD_reset_session_only);

    guchar *retval = output.dst;
    if (output.pos == output.size)
        retval = g_renew(guchar, retval, output.size+1);
    retval[output.pos] = '\0';
    *decompsize = output.pos;
    return retval;
}
#endif

static void
free_zonfile(ZONFile *zonfile)
{
    guint i;

    if (zonfile->hash)
        g_hash_table_destroy(zonfile->hash);
    if (zonfile->path)
        g_string_free(zonfile->path, TRUE);
    if (zonfile->str)
        g_string_free(zonfile->str, TRUE);
    if (zonfile->datafiles) {
        for (i = 0; i < zonfile->datafiles->len; i++) {
            ZONData *zondata = &g_array_index(zonfile->datafiles, ZONData, i);
            g_free(zondata->uuid);
            g_free(zondata->freeme);
        }
        g_array_free(zonfile->datafiles, TRUE);
    }
#ifdef HAVE_ZSTD
    ZSTD_freeDCtx(zonfile->zstd_context);
#endif
}

/* vim: set cin columns=120 tw=118 et ts=4 sw=4 cino=>1s,e0,n0,f0,{0,}0,^0,\:1s,=0,g1s,h0,t0,+1s,c3,(0,u0 : */
