#include <algorithm>
#include <chrono>
#include <libzfs.h>
#include <map>
#include <optional>
#include <set>
#include <stdio.h>
#include <string.h>
#include <vector>

struct Partition {
  std::string pool;
  std::string key;
  std::string ds;
};

std::vector<Partition> partitions = {
    {"archive", "pu", "archive/public"},
    {"archive", "pu/Steam", "archive/public/SteamLibrary"},
    {"archive", "pu/_downl", "archive/public/_downloading"},
    {"archive", "pu/books", "archive/public/books"},
    {"archive", "pu/files", "archive/public/files"},
    {"archive", "pu/instal", "archive/public/install"},
    {"archive", "pu/libre", "archive/public/libreelec"},
    {"archive", "pu/mov", "archive/public/movies"},
    {"archive", "pu/mus", "archive/public/music"},
    {"archive", "pu/photo", "archive/public/photos"},
    {"archive", "replica", "archive/replica"},
    {"archive", "replica", "archive/replica/Minecraft"},
    {"archive", "replica", "archive/replica/games_300"},
    {"archive", "replica", "archive/replica/games"},
    {"archive", "replica", "archive/replica/games/SteamLibrary"},
    {"archive", "replica", "archive/replica/games/baltica"},
    {"archive", "replica", "archive/replica/games/drive_d"},
    {"archive", "replica", "archive/replica/home"},
    {"archive", "replica", "archive/replica/home/ivan-cache"},
    {"archive", "replica", "archive/replica/home/katya"},
    {"archive", "replica", "archive/replica/home/katya/cache"},
    {"archive", "replica", "archive/replica/marianna"},
    {"archive", "replica", "archive/replica/marianna/ROOT"},
    {"archive", "replica", "archive/replica/marianna/win10"},
    {"archive", "replica", "archive/replica/root"},
    {"archive", "replica", "archive/replica/temp"},
    {"archive", "replica", "archive/replica/win10"},
    {"archive", "replica", "archive/replica/win10b"},
    {"archive", "replica", "archive/replica/win81"},
    {"archive", "replica", "archive/replica/var-cache"},
    {"archive", "replica", "archive/replica/var-log"},
    {"archive", "replica", "archive/replica/kernel"},
    {"archive", "sy/baltica", "archive/system/baltica"},
    {"archive", "sy/distfs", "archive/system/distfiles"},
    {"archive", "sy/gosync", "archive/system/gosync"},
    {"archive", "sy/photosy", "archive/system/photos-sync"},
    {"archive", "sy/printer", "archive/system/printer"},
    {"archive", "sy/q/w10c", "archive/system/qemu/win10c"},

    {"games", "steam", "games/SteamLibrary"},
    {"games", "drive_d", "games/drive_d"},
    {"games", "baltica", "games/baltica"},

    {"ssd", "Minecrt", "ssd/Minecraft"},
    {"ssd", "gams300", "ssd/games_300"},
    {"ssd", "home", "ssd/home"},
    {"ssd", "h/iv-c", "ssd/home/ivan-cache"},
    {"ssd", "h/kat", "ssd/home/katya"},
    {"ssd", "h/kat/c", "ssd/home/katya/cache"},
    {"ssd", "kernel", "ssd/kernel"},
    {"ssd", "portge", "ssd/portage"},
    {"ssd", "root",   "ssd/root"},
    {"ssd", "v/log", "ssd/var-log"},
    {"ssd", "v/cache", "ssd/var-cache"},
    {"ssd", "temp", "ssd/temp"},
    {"ssd", "temp", "ssd/temp/ROOT"},
    {"ssd", "temp", "ssd/temp/acl"},
    {"ssd", "temp", "ssd/temp/master"},
    {"ssd", "temp", "ssd/temp/master/img"},
    {"ssd", "w10b", "ssd/win10b"},
    {"ssd", "w81", "ssd/win81"},
};

int main() {
  const char* myzfs = "myzfs3";

  typedef std::chrono::high_resolution_clock Clock;

  libzfs_handle_t *h = libzfs_init();
  char property[ZFS_MAXPROPLEN];
  zprop_source_t source;

  std::map<std::string,std::set<std::string>> charts;
  for (const auto& it : partitions) {
    charts[it.pool].insert(it.key);
  }
  for (auto chart : charts) {
    printf("CHART %s.%s-ds '' 'Fill rate of %s pool datasets' 'KiB' %s ds-rate line 1000 1\n",
           myzfs, chart.first.c_str(), chart.first.c_str(), chart.first.c_str());
    for (auto dim : chart.second) {
      printf("DIMENSION %s '' absolute 1 1024\n", dim.c_str());
    }
  }
  for (auto chart : charts) {
    printf("CHART %s.%s-snap '' 'Fill rate of %s pool snapshots' 'KiB' %s snap-rate line 1000 1\n",
           myzfs, chart.first.c_str(), chart.first.c_str(), chart.first.c_str());
    for (auto dim : chart.second) {
      printf("DIMENSION %s '' absolute 1 1024\n", dim.c_str());
    }
  }
  for (auto chart : charts) {
    printf("CHART %s.%s-usage '' 'Usage of %s pool' 'KiB' %s usage stacked 1000 1\n",
           myzfs, chart.first.c_str(), chart.first.c_str(), chart.first.c_str());
    for (auto dim : chart.second) {
      printf("DIMENSION %s '' absolute 1 1024\n", dim.c_str());
      printf("DIMENSION %sSNAP '' absolute 1 1024\n", dim.c_str());
    }
  }
  for (auto chart : charts) {
    printf("CHART %s.%s-allocated '' 'Percent of allocation for %s pool' '%%' %s allocated line 1000 1\n",
           myzfs, chart.first.c_str(), chart.first.c_str(), chart.first.c_str());
      printf("DIMENSION percent '' absolute 1 10000000000\n");
  }
  
  auto prev = Clock::now();
  char buf[1024];
  buf[0] = 0;
  auto start = Clock::now();
  int iter = 0;
  std::map<std::string, std::map<std::string, int64_t>> prev_ds_usage;
  std::map<std::string, std::map<std::string, int64_t>> prev_snap_usage;

  while (true) {
    iter++;
    auto now = Clock::now();
    uint64_t ms = std::chrono::duration_cast<std::chrono::microseconds>(now - prev).count();
    prev = now;

    if (iter != 1) {
      snprintf(buf, 1024, " %ld", ms);
    }

    std::map<std::string, std::map<std::string, int64_t>> ds_usage;
    std::map<std::string, std::map<std::string, int64_t>> snap_usage;
    for (const auto& it : partitions) {
      auto fs = zfs_open(h, it.ds.c_str(), ZFS_TYPE_DATASET);
      if (fs == NULL) continue;

      uint64_t ds = -1, snap = -1;
      zfs_prop_get_numeric(fs, ZFS_PROP_USEDDS, &ds, &source, property, sizeof(property));
      if (ds != -1) ds_usage[it.pool][it.key] += ds;

      zfs_prop_get_numeric(fs, ZFS_PROP_USEDSNAP, &snap, &source, property, sizeof(property));
      if (snap != -1) snap_usage[it.pool][it.key] += snap;
      zfs_close(fs);
    }

    for (const auto& it : ds_usage) {
      auto prev = prev_ds_usage.find(it.first);
      if (prev == prev_ds_usage.end()) {
        // skip chart
        continue;
      }
      printf("BEGIN %s.%s-ds%s\n", myzfs, it.first.c_str(), buf);
      for (const auto& it2 : it.second) {

        auto prev_ds = prev->second.find(it2.first);
        if (prev_ds == prev->second.end()) {
          // skip dimension
          continue;
        }
        printf("SET %s = %ld\n", it2.first.c_str(), it2.second - prev_ds->second);
      }
      printf("END\n");
    }
    for (const auto& it : snap_usage) {
      auto prev = prev_snap_usage.find(it.first);
      if (prev == prev_snap_usage.end()) {
        // skip chart
        continue;
      }
      printf("BEGIN %s.%s-snap%s\n", myzfs, it.first.c_str(), buf);
      for (const auto& it2 : it.second) {

        auto prev_snap = prev->second.find(it2.first);
        if (prev_snap == prev->second.end()) {
          // skip dimension
          continue;
        }
        printf("SET %s = %ld\n", it2.first.c_str(), it2.second - prev_snap->second);
      }
      printf("END\n");
    }
    for (const auto& it : ds_usage) {
      printf("BEGIN %s.%s-usage%s\n", myzfs, it.first.c_str(), buf);
      for (const auto& it2 : it.second) {
        printf("SET %s = %ld\n", it2.first.c_str(), it2.second);
      }
      for (const auto& it2 : snap_usage[it.first]) {
        printf("SET %sSNAP = %ld\n", it2.first.c_str(), it2.second);
      }
      printf("END\n");
    }

    for (const auto& chart: charts) {
      const std::string& pool_name = chart.first;
      printf("BEGIN %s.%s-allocated%s\n", myzfs, pool_name.c_str(), buf);
      uint64_t size, allocated;
      auto pool = zpool_open(h, pool_name.c_str());
      if (pool == nullptr) continue;
      size = zpool_get_prop_int(pool, ZPOOL_PROP_SIZE, nullptr);
      allocated = zpool_get_prop_int(pool, ZPOOL_PROP_ALLOCATED, nullptr);
      //printf("%ld %ld\n", size, allocated);
      uint64_t usage = (uint64_t)(((double)allocated) * 1e12 / size);
      printf("SET percent = %ld\n", usage);
      printf("END\n");
      zpool_close(pool);
    }

    fflush(stdout);
    prev_ds_usage = ds_usage;
    prev_snap_usage = snap_usage;

    int64_t elapsed = std::chrono::duration_cast<std::chrono::microseconds>(Clock::now() - start).count();
    int64_t target = iter * 1000000ll;
    int64_t to_sleep = std::max<int64_t>(500000, target - elapsed);
    usleep(to_sleep);
  }
}