// SPDX-License-Identifier: GPL-2.0-or-later
// mkscff: write a block image as SCFF, the SchOS Compressed File Format.
//
//   mkscff IN.img OUT.scff [--chunk KB] [--level N] [--threads N]
//
// The layout is include/fs/scff.h. Chunks are compressed on every core, since
// zstd's high levels are slow and the image is hundreds of megabytes, and each
// one is decompressed again and compared before it is written: an image the
// kernel cannot read is found here rather than at boot.
//
// Built against the zstd release the kernel's decoder comes from (ports/zstd),
// not the host's libzstd, so both sides are one version and nothing needs to
// be installed for it.

#include <errno.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "fs/scff.h"
#include "zstd.h"

static unsigned char *img;
static uint64_t img_bytes;
static uint32_t chunk_bytes = 128 * 1024;
static int level = 19;
static uint32_t nchunks;

static unsigned char **out;     // chunk i as stored, or 0 for an all-zero one
static size_t *outlen;
static uint32_t next_chunk;
static pthread_mutex_t next_lock = PTHREAD_MUTEX_INITIALIZER;

static void
die(const char *what)
{
  fprintf(stderr, "mkscff: %s\n", what);
  exit(1);
}

static int
all_zero(const unsigned char *p, size_t n)
{
  for (size_t i = 0; i < n; i++)
    if (p[i])
      return 0;
  return 1;
}

static void *
worker(void *arg)
{
  (void)arg;
  ZSTD_CCtx *cc = ZSTD_createCCtx();
  ZSTD_DCtx *dc = ZSTD_createDCtx();
  size_t cap = ZSTD_compressBound(chunk_bytes);
  unsigned char *buf = malloc(cap), *back = malloc(chunk_bytes);

  if (!cc || !dc || !buf || !back)
    die("out of memory");
  ZSTD_CCtx_setParameter(cc, ZSTD_c_compressionLevel, level);
  ZSTD_CCtx_setParameter(cc, ZSTD_c_checksumFlag, 1);

  for (;;) {
    pthread_mutex_lock(&next_lock);
    uint32_t i = next_chunk++;
    pthread_mutex_unlock(&next_lock);
    if (i >= nchunks)
      break;

    const unsigned char *src = img + (uint64_t)i * chunk_bytes;
    size_t n = img_bytes - (uint64_t)i * chunk_bytes;
    if (n > chunk_bytes)
      n = chunk_bytes;

    if (all_zero(src, n)) {
      outlen[i] = 0;
      continue;
    }
    size_t c = ZSTD_compress2(cc, buf, cap, src, n);
    if (ZSTD_isError(c))
      die(ZSTD_getErrorName(c));
    if (c >= n) {                     // did not compress: stored as it is
      out[i] = malloc(n);
      memcpy(out[i], src, n);
      outlen[i] = n;
      continue;
    }
    size_t d = ZSTD_decompressDCtx(dc, back, chunk_bytes, buf, c);
    if (ZSTD_isError(d) || d != n || memcmp(back, src, n) != 0)
      die("a chunk did not decompress to itself");
    out[i] = malloc(c);
    memcpy(out[i], buf, c);
    outlen[i] = c;
  }
  ZSTD_freeCCtx(cc);
  ZSTD_freeDCtx(dc);
  free(buf);
  free(back);
  return 0;
}

int
main(int argc, char **argv)
{
  const char *in = 0, *outpath = 0;
  long threads = sysconf(_SC_NPROCESSORS_ONLN);

  for (int i = 1; i < argc; i++) {
    if (!strcmp(argv[i], "--chunk") && i + 1 < argc)
      chunk_bytes = (uint32_t)atoi(argv[++i]) * 1024;
    else if (!strcmp(argv[i], "--level") && i + 1 < argc)
      level = atoi(argv[++i]);
    else if (!strcmp(argv[i], "--threads") && i + 1 < argc)
      threads = atol(argv[++i]);
    else if (!in)
      in = argv[i];
    else if (!outpath)
      outpath = argv[i];
    else
      die("usage: mkscff IN.img OUT.scff [--chunk KB] [--level N] [--threads N]");
  }
  if (!in || !outpath)
    die("usage: mkscff IN.img OUT.scff [--chunk KB] [--level N] [--threads N]");
  if (chunk_bytes < 4096 || chunk_bytes > SCFF_CHUNK_MAX ||
      (chunk_bytes & (chunk_bytes - 1)))
    die("--chunk must be a power of two from 4 to 1024 KB");
  if (threads < 1)
    threads = 1;

  FILE *f = fopen(in, "rb");
  if (!f)
    die(strerror(errno));
  fseek(f, 0, SEEK_END);
  img_bytes = (uint64_t)ftell(f);
  fseek(f, 0, SEEK_SET);
  if (img_bytes == 0 || img_bytes % 4096)
    die("the image is not a whole number of 4 KB blocks");
  if (!(img = malloc(img_bytes)) || fread(img, 1, img_bytes, f) != img_bytes)
    die("cannot read the image");
  fclose(f);

  nchunks = (uint32_t)((img_bytes + chunk_bytes - 1) / chunk_bytes);
  out = calloc(nchunks, sizeof(*out));
  outlen = calloc(nchunks, sizeof(*outlen));
  if (!out || !outlen)
    die("out of memory");

  pthread_t *t = calloc((size_t)threads, sizeof(*t));
  for (long i = 0; i < threads; i++)
    pthread_create(&t[i], 0, worker, 0);
  for (long i = 0; i < threads; i++)
    pthread_join(t[i], 0);

  struct scff_header h;
  memset(&h, 0, sizeof(h));
  memcpy(h.magic, SCFF_MAGIC, 4);
  h.version = SCFF_VERSION;
  h.codec = SCFF_CODEC_ZSTD;
  h.chunk_bytes = chunk_bytes;
  h.nchunks = nchunks;
  h.image_bytes = img_bytes;
  h.index_off = sizeof(h);

  uint64_t *index = malloc(((size_t)nchunks + 1) * sizeof(uint64_t));
  uint64_t pos = sizeof(h) + ((uint64_t)nchunks + 1) * sizeof(uint64_t);
  for (uint32_t i = 0; i < nchunks; i++) {
    index[i] = pos;
    pos += outlen[i];
  }
  index[nchunks] = pos;
  h.stored_bytes = pos;

  if (!(f = fopen(outpath, "wb")))
    die(strerror(errno));
  fwrite(&h, sizeof(h), 1, f);
  fwrite(index, sizeof(uint64_t), (size_t)nchunks + 1, f);
  for (uint32_t i = 0; i < nchunks; i++)
    if (outlen[i])
      fwrite(out[i], 1, outlen[i], f);
  if (fclose(f) != 0)
    die("cannot write the output");

  printf("mkscff: %llu KB -> %llu KB (%.1f%%), %u chunks of %u KB, zstd -%d\n",
         (unsigned long long)(img_bytes >> 10), (unsigned long long)(pos >> 10),
         100.0 * (double)pos / (double)img_bytes, nchunks, chunk_bytes >> 10, level);
  return 0;
}