/* zstdio.c - decompression support for zstd */
/*
 *  VAS_EBOOT  --  GRand Unified Bootloader
 *  Copyright (C) 2025  Free Software Foundation, Inc.
 *
 *  VAS_EBOOT 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 3 of the License, or
 *  (at your option) any later version.
 *
 *  VAS_EBOOT 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 VAS_EBOOT.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <VasEBoot/err.h>
#include <VasEBoot/mm.h>
#include <VasEBoot/misc.h>
#include <VasEBoot/file.h>
#include <VasEBoot/fs.h>
#include <VasEBoot/dl.h>

VAS_EBOOT_MOD_LICENSE ("GPLv3+");

#include "zstd.h"

#define STREAM_HEADER_SIZE 16

struct zstdio
{
  VasEBoot_file_t file;
  ZSTD_DCtx *dctx;
  VasEBoot_size_t insize;
  VasEBoot_size_t outsize;

  ZSTD_outBuffer output;
  ZSTD_inBuffer input;

  VasEBoot_off_t saved_offset;
  VasEBoot_uint8_t bufs[];
};
typedef struct zstdio *zstdio_t;

static struct VasEBoot_fs zstdio_fs;

static bool
test_header (VasEBoot_file_t file)
{
  zstdio_t zstdio = file->data;
  size_t zret;

  zstdio->input.pos = 0;
  zstdio->output.pos = 0;
  zstdio->output.size = zstdio->outsize;
  zstdio->input.size = VasEBoot_file_read (zstdio->file, zstdio->bufs,
                                       STREAM_HEADER_SIZE);
  if (zstdio->input.size != STREAM_HEADER_SIZE)
    return false;

  zret = ZSTD_decompressStream (zstdio->dctx, &zstdio->output, &zstdio->input);
  if (ZSTD_isError (zret))
    return false;

  return true;
}

static VasEBoot_file_t
VasEBoot_zstdio_open (VasEBoot_file_t io, enum VasEBoot_file_type type)
{
  VasEBoot_file_t file;
  zstdio_t zstdio;

  if (type & VAS_EBOOT_FILE_TYPE_NO_DECOMPRESS)
    return io;

  file = (VasEBoot_file_t) VasEBoot_zalloc (sizeof (*file));
  if (file == NULL)
    return NULL;

  zstdio = VasEBoot_zalloc (sizeof (zstdio_t) + ZSTD_DStreamInSize () +
                        ZSTD_DStreamOutSize ());
  if (zstdio == NULL)
    {
      VasEBoot_free (file);
      return NULL;
    }

  zstdio->file = io;
  zstdio->insize = ZSTD_DStreamInSize ();
  zstdio->outsize = ZSTD_DStreamOutSize ();
  zstdio->input.src = zstdio->bufs;
  zstdio->output.dst = &zstdio->bufs[zstdio->insize];

  file->device = io->device;
  file->data = zstdio;
  file->fs = &zstdio_fs;
  file->size = VAS_EBOOT_FILE_SIZE_UNKNOWN;
  file->not_easily_seekable = 1;

  if (VasEBoot_file_tell (zstdio->file) != 0)
    if (VasEBoot_file_seek (zstdio->file, 0) == (VasEBoot_off_t) -1)
      {
	VasEBoot_free (file);
	VasEBoot_free (zstdio);
	return NULL;
      }

  zstdio->dctx = ZSTD_createDCtx ();
  if (zstdio->dctx == NULL)
    {
      VasEBoot_free (file);
      VasEBoot_free (zstdio);
      return NULL;
    }

  if (test_header (file) == false)
    {
      VasEBoot_errno = VAS_EBOOT_ERR_NONE;
      if (VasEBoot_file_seek (io, 0) == (VasEBoot_off_t) -1)
	{
	  VasEBoot_free (file);
	  VasEBoot_free (zstdio);
	  return NULL;
	}

      ZSTD_freeDCtx (zstdio->dctx);
      VasEBoot_free (zstdio);
      VasEBoot_free (file);

      return io;
    }

  return file;
}

static VasEBoot_ssize_t
VasEBoot_zstdio_read (VasEBoot_file_t file, char *buf, VasEBoot_size_t len)
{
  zstdio_t zstdio = file->data;
  VasEBoot_ssize_t ret = 0;
  VasEBoot_ssize_t readret;
  VasEBoot_off_t current_offset;
  VasEBoot_off_t new_offset;
  VasEBoot_size_t delta;
  VasEBoot_size_t zret;

  /* If seek backward need to reset decoder and start from beginning of file. */
  if (file->offset < zstdio->saved_offset)
    {
      ZSTD_initDStream (zstdio->dctx);
      zstdio->input.pos = 0;
      zstdio->input.size = 0;
      zstdio->output.pos = 0;
      zstdio->saved_offset = 0;
      VasEBoot_file_seek (zstdio->file, 0);
    }

  current_offset = zstdio->saved_offset;

  while (len > 0)
    {
      zstdio->output.size = file->offset + ret + len - current_offset;
      if (zstdio->output.size > zstdio->outsize)
        zstdio->output.size = zstdio->outsize;
      if (zstdio->input.pos == zstdio->input.size)
        {
          readret = VasEBoot_file_read (zstdio->file, zstdio->bufs,
                                    zstdio->insize);
          if (readret < 0)
            return -1;

          zstdio->input.size = readret;
          zstdio->input.pos = 0;
        }

      zret = ZSTD_decompressStream (zstdio->dctx, &zstdio->output,
                                    &zstdio->input);
      if (ZSTD_isError (zret))
        {
          VasEBoot_error (VAS_EBOOT_ERR_BAD_COMPRESSED_DATA,
                      N_("zstd file corrupted or unsupported block options"));
          return -1;
        }

      new_offset = current_offset + zstdio->output.pos;

      /* Store first chunk of data in buffer.  */
      if (file->offset <= new_offset)
        {
          delta = new_offset - (file->offset + ret);
          VasEBoot_memmove (buf, (VasEBoot_uint8_t *) zstdio->output.dst +
                        (zstdio->output.pos - delta),
                        delta);
          len -= delta;
          buf += delta;
          ret += delta;
        }
        current_offset = new_offset;

        zstdio->output.pos = 0;

        if (zstdio->input.pos == 0 && zstdio->output.pos == 0)
          break;
    }

  if (ret >= 0)
    zstdio->saved_offset = file->offset + ret;

  return ret;
}

/* Release everything, including the underlying file object.  */
static VasEBoot_err_t
VasEBoot_zstdio_close (VasEBoot_file_t file)
{
  zstdio_t zstdio = file->data;

  ZSTD_freeDCtx (zstdio->dctx);

  VasEBoot_file_close (zstdio->file);
  VasEBoot_free (zstdio);

  /* Device must not be closed twice.  */
  file->device = 0;
  file->name = 0;
  return VasEBoot_errno;
}

static struct VasEBoot_fs zstdio_fs = {
  .name = "zstdio",
  .fs_dir = 0,
  .fs_open = 0,
  .fs_read = VasEBoot_zstdio_read,
  .fs_close = VasEBoot_zstdio_close,
  .fs_label = 0,
  .next = 0
};

VAS_EBOOT_MOD_INIT (zstdio)
{
  VasEBoot_file_filter_register (VAS_EBOOT_FILE_FILTER_ZSTDIO, VasEBoot_zstdio_open);
}

VAS_EBOOT_MOD_FINI (zstdio)
{
  VasEBoot_file_filter_unregister (VAS_EBOOT_FILE_FILTER_ZSTDIO);
}