/****************************************************************************
 * apps/examples/nxflatxip/nxflatxip_main.c
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.  The
 * ASF licenses this file to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance with the
 * License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 ****************************************************************************/

/****************************************************************************
 * End to end demonstration of what a writable execute-in-place file system
 * is for:
 *
 *   1. write an NXFLAT module into xipfs at run time, the way a download
 *      would,
 *   2. confirm the file system can hand out a direct flash pointer for it,
 *   3. run two instances of it concurrently,
 *   4. confirm both ran correctly, that their data did not interfere, and
 *      that the shared text was pinned in place while they ran.
 *
 * The same module in a ROMFS image would give the same execute-in-place
 * behaviour; what it could not do is arrive after the firmware was built.
 ****************************************************************************/

/****************************************************************************
 * Included Files
 ****************************************************************************/

#include <nuttx/config.h>

#include <sys/ioctl.h>
#include <sys/mount.h>
#include <sys/stat.h>
#include <sys/wait.h>

#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <syslog.h>
#include <time.h>
#include <unistd.h>

#include <nuttx/binfmt/binfmt.h>
#include <nuttx/fs/ioctl.h>
#include <nuttx/fs/xipfs.h>
#include <nuttx/symtab.h>

#include "xipmod_bin.h"

/****************************************************************************
 * Pre-processor Definitions
 ****************************************************************************/

#define MOUNTPT    CONFIG_EXAMPLES_NXFLATXIP_MOUNTPT
#define MTDDEV     CONFIG_EXAMPLES_NXFLATXIP_MTD
#define MODPATH    MOUNTPT "/xipmod"
#define BENCH_PATH MOUNTPT "/big"
#define BENCH_SIZE (256 * 1024)
#define DF_FILES   55
#define DF_MAPCOLS 50

/****************************************************************************
 * External Symbols
 ****************************************************************************/

/* Generated by module/Makefile from the import list mknxflat left in the
 * module's thunk file.  The NXFLAT loader has no other way to resolve the
 * functions the module calls.
 */

extern const struct symtab_s g_nxflat_exports[];
extern const int g_nxflat_nexports;

/****************************************************************************
 * Private Data
 ****************************************************************************/

static char g_map[512];

/****************************************************************************
 * Public Functions
 ****************************************************************************/

/****************************************************************************
 * Name: xipmod_report
 *
 * Description:
 *   Called by the module, running in place from flash, to report what it
 *   found.  It passes numbers rather than formatting its own output; see
 *   module/xipmod.c for why.
 *
 *   This is also the direction of the interface that needs r10 to survive:
 *   the module reaches its own data through that register, and this
 *   function runs with the module's data base still in it.  That is what
 *   the --fixed-r10 the ARM toolchain adds under CONFIG_PIC is for.
 *
 ****************************************************************************/

void xipmod_report(int seed, FAR const void *text, FAR const void *stack,
                   long sum, int errors)
{
  syslog(LOG_INFO, "  instance seed %d: text %p, stack %p, sum %ld -- %s\n",
         seed, text, stack, sum,
         errors == 0 ? "data private and intact" : "DATA CORRUPTED");
}

/****************************************************************************
 * Private Functions
 ****************************************************************************/

/****************************************************************************
 * Name: stage_module
 *
 * Description:
 *   Write the module into the file system exactly as a download would:
 *   declare the size up front so xipfs reserves an exactly sized
 *   contiguous extent, then write it once.
 *
 ****************************************************************************/

static int stage_module(void)
{
  ssize_t nwritten;
  int fd;

  unlink(MODPATH);

  fd = open(MODPATH, O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0755);
  if (fd < 0)
    {
      syslog(LOG_INFO, "ERROR: open %s failed: %d\n", MODPATH, errno);
      return -errno;
    }

  if (ftruncate(fd, g_xipmod_len) < 0)
    {
      syslog(LOG_INFO, "ERROR: ftruncate failed: %d\n", errno);
      close(fd);
      return -errno;
    }

  nwritten = write(fd, g_xipmod, g_xipmod_len);
  close(fd);

  if (nwritten != (ssize_t)g_xipmod_len)
    {
      syslog(LOG_INFO, "ERROR: short write: %zd of %u\n", nwritten,
             g_xipmod_len);
      return -EIO;
    }

  syslog(LOG_INFO, "staged %u bytes to %s\n", g_xipmod_len, MODPATH);
  return 0;
}

/****************************************************************************
 * Name: extent_info_path
 *
 * Description:
 *   Where a file physically lies: its extent, its flash address, and how
 *   many live mappings are pinning it in place.
 *
 ****************************************************************************/

static int extent_info_path(FAR const char *path,
                            FAR struct xipfs_extent_info_s *info)
{
  int fd;
  int ret;

  fd = open(path, O_RDONLY | O_CLOEXEC);
  if (fd < 0)
    {
      return -errno;
    }

  ret = ioctl(fd, XIPFSIOC_EXTENTINFO, (unsigned long)(uintptr_t)info);
  close(fd);

  return ret < 0 ? -errno : 0;
}

/****************************************************************************
 * Name: bench_read
 *
 * Description:
 *   Time a bulk read straight out of the memory mapped flash, big enough to
 *   miss the XIP cache.  Run once on a freshly booted system, where the
 *   bootrom's fast read mode is still in effect, and again after a flash
 *   write has forced the driver to restore XIP itself.  The difference is
 *   the cost of whichever restore path the driver used.
 *
 ****************************************************************************/

static long bench_read(void)
{
  struct timespec t0;
  struct timespec t1;
  FAR uint8_t *buf;
  ssize_t n;
  long total = 0;
  int fd;

  buf = malloc(4096);
  if (buf == NULL)
    {
      return -1;
    }

  fd = open(BENCH_PATH, O_RDONLY | O_CLOEXEC);
  if (fd < 0)
    {
      free(buf);
      return -1;
    }

  clock_gettime(CLOCK_MONOTONIC, &t0);

  while ((n = read(fd, buf, 4096)) > 0)
    {
      total += n;
    }

  clock_gettime(CLOCK_MONOTONIC, &t1);

  close(fd);
  free(buf);

  if (total <= 0)
    {
      return -1;
    }

  return (t1.tv_sec - t0.tv_sec) * 1000 +
         (t1.tv_nsec - t0.tv_nsec) / 1000000;
}

static int bench_main(void)
{
  long before;
  long after;
  int fd;

  syslog(LOG_INFO, "\nXIP read benchmark (%d KB)\n", BENCH_SIZE / 1024);

  before = bench_read();
  if (before < 0)
    {
      syslog(LOG_INFO, "  create %s first: "
             "dd if=/dev/zero of=%s bs=512 count=512\n",
             BENCH_PATH, BENCH_PATH);
      return EXIT_FAILURE;
    }

  syslog(LOG_INFO, "  before any flash write : %ld ms\n", before);

  /* Any write forces the driver to tear XIP down and put it back */

  fd = open(MOUNTPT "/benchw",
            O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
  if (fd >= 0)
    {
      write(fd, "x", 1);
      close(fd);
    }

  after = bench_read();
  syslog(LOG_INFO, "  after  a flash write   : %ld ms\n", after);

  if (before > 0)
    {
      syslog(LOG_INFO, "  ratio                  : %ld.%02ldx\n",
             after / before, ((after * 100) / before) % 100);
    }

  unlink(MOUNTPT "/benchw");
  return EXIT_SUCCESS;
}

static int remount(void)
{
  int ret;

  ret = umount(MOUNTPT);
  if (ret < 0 && errno != ENOENT && errno != EINVAL)
    {
      return -errno;
    }

  ret = mount(MTDDEV, MOUNTPT, "xipfs", 0, NULL);
  return ret < 0 ? -errno : 0;
}

static void fill_pattern(FAR uint8_t *buf, size_t len, uint8_t seed)
{
  size_t i;

  for (i = 0; i < len; i++)
    {
      buf[i] = (uint8_t)(seed + (i * 7));
    }
}

static int create_file(FAR const char *path, size_t len, uint8_t seed)
{
  FAR uint8_t *buf;
  ssize_t nw;
  int fd;
  int ret = 0;

  buf = malloc(len);
  if (buf == NULL)
    {
      return -ENOMEM;
    }

  fill_pattern(buf, len, seed);

  fd = open(path, O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
  if (fd < 0)
    {
      free(buf);
      return -errno;
    }

  if (ftruncate(fd, len) < 0)
    {
      ret = -errno;
    }
  else
    {
      nw = write(fd, buf, len);
      if (nw != (ssize_t)len)
        {
          ret = nw < 0 ? -errno : -EIO;
        }
    }

  close(fd);
  free(buf);
  return ret;
}

static bool verify_file(FAR const char *path, size_t len, uint8_t seed)
{
  FAR uint8_t *got;
  FAR uint8_t *want;
  bool ok = false;
  ssize_t n;
  int fd;

  got  = malloc(len);
  want = malloc(len);
  if (got == NULL || want == NULL)
    {
      free(got);
      free(want);
      return false;
    }

  fill_pattern(want, len, seed);

  fd = open(path, O_RDONLY | O_CLOEXEC);
  if (fd >= 0)
    {
      n = read(fd, got, len);
      ok = (n == (ssize_t)len) && (memcmp(got, want, len) == 0);
      close(fd);
    }

  free(got);
  free(want);
  return ok;
}

/****************************************************************************
 * Name: run_defrag
 *
 * Description:
 *   Compact the volume.  The command acts on the volume, so it goes through
 *   a descriptor for the mountpoint directory: a descriptor for a file
 *   inside it would hold that file open, and an open extent cannot be
 *   relocated.
 *
 ****************************************************************************/

static int run_defrag(uint32_t max_ms,
                      FAR struct xipfs_defrag_result_s *result)
{
  struct xipfs_defrag_arg_s arg;
  int fd;
  int ret;

  memset(&arg, 0, sizeof(arg));
  arg.max_ms = max_ms;

  fd = open(MOUNTPT, O_RDONLY | O_DIRECTORY | O_CLOEXEC);
  if (fd < 0)
    {
      return -errno;
    }

  ret = ioctl(fd, XIPFSIOC_DEFRAG, (unsigned long)(uintptr_t)&arg);
  close(fd);

  if (ret < 0)
    {
      return -errno;
    }

  *result = arg.result;
  return 0;
}

/****************************************************************************
 * Defragmentation showcase
 *
 * Fragments the volume until an allocation genuinely cannot be satisfied --
 * plenty of free space, but none of it contiguous -- then compacts and
 * retries the same allocation.  On a real board everything below runs
 * against real NOR: each move is a page-at-a-time copy, a metadata commit,
 * and an erase of the vacated blocks.
 ****************************************************************************/

static int df_mapfill(FAR uint32_t *dstart, FAR uint32_t *dblocks)
{
  struct xipfs_extent_info_s info;
  FAR struct dirent *de;
  FAR DIR *dirp;
  char path[80];
  int files = 0;
  uint32_t i;

  memset(g_map, '.', sizeof(g_map));
  *dstart = 0;
  *dblocks = 0;

  dirp = opendir(MOUNTPT);
  if (dirp == NULL)
    {
      return -1;
    }

  while ((de = readdir(dirp)) != NULL)
    {
      snprintf(path, sizeof(path), "%s/%s", MOUNTPT, de->d_name);
      if (extent_info_path(path, &info) < 0)
        {
          continue;
        }

      *dstart  = info.data_start;
      *dblocks = info.data_nblocks;
      files++;

      for (i = 0; i < info.nblocks; i++)
        {
          uint32_t idx = info.start_block - info.data_start + i;
          if (idx < sizeof(g_map))
            {
              g_map[idx] = '#';
            }
        }
    }

  closedir(dirp);
  return files;
}

static void df_showmap(FAR const char *title, uint32_t dblocks,
                       FAR uint32_t *freeout, FAR uint32_t *largestout)
{
  uint32_t largest = 0;
  uint32_t run = 0;
  uint32_t free = 0;
  uint32_t i;
  char row[DF_MAPCOLS + 1];

  syslog(LOG_INFO, "\n  %s\n", title);

  for (i = 0; i < dblocks; i++)
    {
      if (g_map[i] == '.')
        {
          free++;
          if (++run > largest)
            {
              largest = run;
            }
        }
      else
        {
          run = 0;
        }
    }

  for (i = 0; i < dblocks; i += DF_MAPCOLS)
    {
      uint32_t n = dblocks - i;
      if (n > DF_MAPCOLS)
        {
          n = DF_MAPCOLS;
        }

      memcpy(row, &g_map[i], n);
      row[n] = '\0';
      syslog(LOG_INFO, "    %3" PRIu32 " |%s|\n", i, row);
    }

  syslog(LOG_INFO, "    free %" PRIu32
         " blocks, largest contiguous run %" PRIu32 "\n", free, largest);

  if (freeout != NULL)
    {
      *freeout = free;
    }

  if (largestout != NULL)
    {
      *largestout = largest;
    }
}

static int df_try_alloc(uint32_t nblocks, uint32_t erasesize)
{
  int fd;
  int ret;

  unlink(MOUNTPT "/bigalloc");

  fd = open(MOUNTPT "/bigalloc",
            O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
  if (fd < 0)
    {
      return -errno;
    }

  ret = ftruncate(fd, (off_t)nblocks * erasesize);
  close(fd);

  if (ret < 0)
    {
      ret = -errno;
      unlink(MOUNTPT "/bigalloc");
    }

  return ret;
}

/****************************************************************************
 * Name: empty_dir
 *
 * Description:
 *   Delete everything below 'dir', descending into subdirectories.
 *
 *   One entry is taken per pass rather than a listing being collected up
 *   front, because this recurses and the task stack is the default 2 KB: an
 *   array of names big enough to be useful costs more than a kilobyte a
 *   frame, and two levels of that overflows and returns through a smashed
 *   frame.  A pass costs an opendir instead, which no one is timing.
 *
 *   The directory is closed before anything is removed, since unlinking
 *   while iterating skips entries.
 *
 ****************************************************************************/

static void empty_dir(FAR const char *dir)
{
  char path[80];
  bool more = true;

  while (more)
    {
      FAR struct dirent *de;
      FAR DIR *dirp;
      bool isdir;

      dirp = opendir(dir);
      if (dirp == NULL)
        {
          return;
        }

      de = readdir(dirp);
      if (de == NULL)
        {
          closedir(dirp);
          return;
        }

      isdir = (de->d_type == DTYPE_DIRECTORY);
      snprintf(path, sizeof(path), "%s/%s", dir, de->d_name);
      closedir(dirp);

      if (isdir)
        {
          empty_dir(path);
          more = rmdir(path) == 0;
        }
      else
        {
          more = unlink(path) == 0;
        }
    }
}

static int defrag_main(void)
{
  struct xipfs_defrag_result_s res;
  struct xipfs_extent_info_s info;
  struct timespec t0;
  struct timespec t1;
  uint32_t dstart;
  uint32_t dblocks;
  uint32_t erasesize = 4096;
  uint32_t want;
  uint32_t freeblk = 0;
  uint32_t largest = 0;
  char path[80];
  long ms;
  int i;
  int ret;

  memset(&res, 0, sizeof(res));

  syslog(LOG_INFO, "\n=== xipfs defragmentation ===\n");

  /* Clean slate, so the map below accounts for every used block */

  empty_dir(MOUNTPT);

  syslog(LOG_INFO, "\n[1] writing %d files...\n", DF_FILES);

  for (i = 0; i < DF_FILES; i++)
    {
      snprintf(path, sizeof(path), MOUNTPT "/m%02d", i);
      if (create_file(path, 2000, (uint8_t)(i + 1)) < 0)
        {
          syslog(LOG_INFO, "    stopped at %d (volume full)\n", i);
          break;
        }
    }

  if (extent_info_path(MOUNTPT "/m00", &info) == 0)
    {
      erasesize = info.erasesize;
    }

  df_mapfill(&dstart, &dblocks);
  df_showmap("map, all files present:", dblocks, NULL, NULL);

  syslog(LOG_INFO, "\n[2] deleting every other file to fragment it...\n");

  for (i = 0; i < DF_FILES; i += 2)
    {
      snprintf(path, sizeof(path), MOUNTPT "/m%02d", i);
      unlink(path);
    }

  df_mapfill(&dstart, &dblocks);
  df_showmap("map, fragmented (# used, . free):", dblocks,
             &freeblk, &largest);

  /* Bigger than any single hole, but comfortably inside the total free
   * space -- so it can only be satisfied once the holes are coalesced.
   */

  want = largest + (freeblk - largest) / 2;

  syslog(LOG_INFO, "\n[3] requesting a %" PRIu32
         "-block (%" PRIu32 " KB) contiguous file\n", want,
         want * erasesize / 1024);

  ret = df_try_alloc(want, erasesize);
  syslog(LOG_INFO, "    result: %s\n",
         ret < 0 ? "FAILED -ENOSPC (free space is not contiguous)" : "ok");

  syslog(LOG_INFO, "\n[4] compacting...\n");

  clock_gettime(CLOCK_MONOTONIC, &t0);
  ret = run_defrag(0, &res);
  clock_gettime(CLOCK_MONOTONIC, &t1);

  ms = (t1.tv_sec - t0.tv_sec) * 1000 + (t1.tv_nsec - t0.tv_nsec) / 1000000;

  if (ret < 0)
    {
      syslog(LOG_INFO, "    defrag failed: %d\n", ret);
      return EXIT_FAILURE;
    }

  syslog(LOG_INFO, "    elapsed          : %ld ms\n", ms);
  syslog(LOG_INFO, "    extents relocated: %" PRIu32 "\n",
         res.extents_moved);
  syslog(LOG_INFO, "    blocks reclaimed : %" PRIu32 "\n",
         res.blocks_reclaimed);
  syslog(LOG_INFO, "    largest free run : %zu bytes (%zu blocks)\n",
         res.largest_free_run, res.largest_free_run / erasesize);
  syslog(LOG_INFO, "    stop reason      : %d (%s)\n", res.reason,
         res.reason == XIPFS_DEFRAG_DONE         ? "done" :
         res.reason == XIPFS_DEFRAG_BLOCKED_OPEN ? "blocked by an open "
                                                   "file" :
         res.reason == XIPFS_DEFRAG_BLOCKED_PINS ? "blocked by a live "
                                                   "XIP mapping" :
                                                   "other");

  df_mapfill(&dstart, &dblocks);
  df_showmap("map after defrag:", dblocks, NULL, NULL);

  want = res.largest_free_run / erasesize;

  syslog(LOG_INFO, "\n[5] retrying, sized to the reported largest run"
         " (%" PRIu32 " blocks)\n", want);
  ret = df_try_alloc(want, erasesize);
  syslog(LOG_INFO, "    result: %s\n",
         ret < 0 ? "still failed" : "OK -- allocation now fits");
  unlink(MOUNTPT "/bigalloc");

  syslog(LOG_INFO, "\n[6] verifying every relocated file byte-for-byte\n");

  ret = 0;
  for (i = 1; i < DF_FILES; i += 2)
    {
      snprintf(path, sizeof(path), MOUNTPT "/m%02d", i);
      if (!verify_file(path, 2000, (uint8_t)(i + 1)))
        {
          syslog(LOG_INFO, "    MISMATCH in %s\n", path);
          ret = -1;
        }
    }

  syslog(LOG_INFO, "    %s\n", ret == 0 ? "all survivors intact" : "FAILED");

  syslog(LOG_INFO, "\n[7] remounting to prove it is durable\n");
  if (remount() < 0)
    {
      syslog(LOG_INFO, "    remount FAILED\n");
      return EXIT_FAILURE;
    }

  ret = 0;
  for (i = 1; i < DF_FILES; i += 2)
    {
      snprintf(path, sizeof(path), MOUNTPT "/m%02d", i);
      if (!verify_file(path, 2000, (uint8_t)(i + 1)))
        {
          ret = -1;
        }
    }

  syslog(LOG_INFO, "    %s\n",
         ret == 0 ? "all files intact after remount" : "FAILED");

  syslog(LOG_INFO, "\n=== defrag complete ===\n");
  return EXIT_SUCCESS;
}

/****************************************************************************
 * Public Functions
 ****************************************************************************/

int main(int argc, FAR char *argv[])
{
  struct xipfs_extent_info_s info;
  FAR char *args[3];
  char seed[2][8];
  pid_t pid[2];
  int status;
  int ret;
  int i;

  if (argc > 1 && strcmp(argv[1], "bench") == 0)
    {
      return bench_main();
    }

  if (argc > 1 && strcmp(argv[1], "defrag") == 0)
    {
      return defrag_main();
    }

  syslog(LOG_INFO,
         "\n=== NXFLAT module executed in place from xipfs ===\n\n");

  ret = stage_module();
  if (ret < 0)
    {
      return EXIT_FAILURE;
    }

  ret = extent_info_path(MODPATH, &info);
  if (ret < 0)
    {
      syslog(LOG_INFO, "ERROR: cannot query extent: %d\n", ret);
      return EXIT_FAILURE;
    }

  syslog(LOG_INFO, "extent: block %" PRIu32 " x%" PRIu32 ", size %" PRIu32
         ", flash addr 0x%08" PRIxPTR "\n",
         info.start_block, info.nblocks, info.size, info.xipaddr);

  if (info.xipaddr == 0)
    {
      syslog(LOG_INFO,
             "ERROR: the file system cannot expose a flash pointer, so the "
             "module would have to be copied to RAM.  Does the MTD driver "
             "answer BIOC_XIPBASE?\n");
      return EXIT_FAILURE;
    }

  /* Launch two instances.  They overlap deliberately: each fills its array,
   * sleeps, then checks it.  Shared data would show up as a mismatch.
   */

  syslog(LOG_INFO, "\nrunning two concurrent instances...\n\n");

  for (i = 0; i < 2; i++)
    {
      snprintf(seed[i], sizeof(seed[i]), "%d", i + 1);

      args[0] = (FAR char *)"xipmod";
      args[1] = seed[i];
      args[2] = NULL;

      /* exec() rather than posix_spawn() because the symbol table the
       * module links against is the demo's, not one the OS holds.
       */

      pid[i] = exec(MODPATH, args, NULL, g_nxflat_exports,
                    g_nxflat_nexports);
      if (pid[i] < 0)
        {
          syslog(LOG_INFO, "ERROR: exec of instance %d failed: %d\n", i + 1,
                 pid[i]);
          return EXIT_FAILURE;
        }
    }

  /* While both are alive the shared text must be pinned once per instance,
   * which is what stops the defragmenter relocating code that is executing.
   */

  usleep(100000);

  if (extent_info_path(MODPATH, &info) == 0)
    {
      syslog(LOG_INFO,
             "\n[while running] pin count on the shared text = %" PRIu32
             "\n\n", info.pincount);
    }

  for (i = 0; i < 2; i++)
    {
      if (waitpid(pid[i], &status, 0) >= 0 &&
          (!WIFEXITED(status) || WEXITSTATUS(status) != EXIT_SUCCESS))
        {
          syslog(LOG_INFO, "instance %d did not exit cleanly\n", i + 1);
          ret = -1;
        }
    }

  /* The pins outlive the instances.  binfmt maps the text in the context of
   * whoever called exec(), so the mappings belong to this task's address
   * space, not to the tasks that ran, and they are released when this task
   * exits.  Run "xipfs -n" after this demo returns to see the count back at
   * zero and the extent movable again.
   */

  if (extent_info_path(MODPATH, &info) == 0)
    {
      syslog(LOG_INFO,
             "[after exit]    pin count on the shared text = %" PRIu32
             " (released when this task exits)\n", info.pincount);
    }

  syslog(LOG_INFO, "\n=== %s ===\n", ret == 0 ? "done" : "FAILED");
  return ret == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}
