#include <openvdb/openvdb.h>

#include <nanovdb/io/IO.h>
#include <nanovdb/tools/CreateNanoGrid.h>

#include <algorithm>
#include <atomic>
#include <array>
#include <chrono>
#include <csignal>
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

#include <unistd.h>

namespace fs = std::filesystem;

constexpr uint64_t MAX_SOURCE_BYTES = 512ULL * 1024ULL * 1024ULL;
constexpr uint64_t MAX_OUTPUT_BYTES = 1024ULL * 1024ULL * 1024ULL;
constexpr uint64_t MAX_ACTIVE_VOXELS = 64ULL * 1000ULL * 1000ULL;
constexpr size_t MAX_GRIDS = 64;

struct Options {
  fs::path input;
  fs::path output;
  fs::path report;
  std::vector<std::string> grids;
  std::string quantization = "LOSSLESS";
  fs::path cancel_file;
  uint64_t timeout_ms = 0;
};

static std::atomic<bool> interrupted(false);

static void request_interrupt(int)
{
  interrupted.store(true, std::memory_order_relaxed);
}

struct GridReport {
  std::string name;
  std::string source_type;
  std::string value_type;
  std::string grid_class;
  uint64_t active_voxels = 0;
  uint64_t segment_offset = 0;
  uint64_t segment_length = 0;
  uint64_t grid_offset = 0;
  uint64_t grid_length = 0;
  openvdb::CoordBBox index_bounds;
  openvdb::BBoxd world_bounds;
  openvdb::Vec3d voxel_size;
  std::array<double, 16> index_to_world{};
  struct ScalarSample {
    openvdb::Coord coord;
    float value;
    bool active;
  };
  struct VectorSample {
    openvdb::Coord coord;
    std::array<float, 3> value;
    bool active;
  };
  std::vector<ScalarSample> scalar_samples;
  std::vector<VectorSample> vector_samples;
};

static std::string json_string(const std::string &value)
{
  std::ostringstream stream;
  stream << '"';
  for (const unsigned char character : value) {
    switch (character) {
      case '"': stream << "\\\""; break;
      case '\\': stream << "\\\\"; break;
      case '\b': stream << "\\b"; break;
      case '\f': stream << "\\f"; break;
      case '\n': stream << "\\n"; break;
      case '\r': stream << "\\r"; break;
      case '\t': stream << "\\t"; break;
      default:
        if (character < 0x20) {
          stream << "\\u" << std::hex << std::setw(4) << std::setfill('0') << int(character) << std::dec;
        }
        else {
          stream << character;
        }
    }
  }
  stream << '"';
  return stream.str();
}

static Options parse_options(int argc, char **argv)
{
  Options options;
  for (int index = 1; index < argc; ++index) {
    const std::string argument = argv[index];
    auto value = [&](const char *name) -> std::string {
      if (++index >= argc) {
        throw std::runtime_error(std::string("missing value for ") + name);
      }
      return argv[index];
    };
    if (argument == "--input") options.input = value("--input");
    else if (argument == "--output") options.output = value("--output");
    else if (argument == "--report") options.report = value("--report");
    else if (argument == "--grid") options.grids.push_back(value("--grid"));
    else if (argument == "--quantization") options.quantization = value("--quantization");
    else if (argument == "--cancel-file") options.cancel_file = value("--cancel-file");
    else if (argument == "--timeout-ms") {
      const std::string raw = value("--timeout-ms");
      size_t consumed = 0;
      options.timeout_ms = std::stoull(raw, &consumed);
      if (consumed != raw.size() || options.timeout_ms < 1 || options.timeout_ms > 60ULL * 60ULL * 1000ULL) {
        throw std::runtime_error("timeout must be between 1ms and 1h");
      }
    }
    else if (argument == "--help") {
      std::cout << "usage: vdb_to_nanovdb --input FILE.vdb --output FILE.nvdb --report REPORT.json "
                   "[--grid NAME] [--quantization LOSSLESS|FP16] [--cancel-file FILE] [--timeout-ms MS]\n";
      std::exit(0);
    }
    else {
      throw std::runtime_error("unknown argument: " + argument);
    }
  }
  if (options.input.empty() || options.output.empty() || options.report.empty()) {
    throw std::runtime_error("--input, --output and --report are required");
  }
  if (options.input.extension() != ".vdb" || options.output.extension() != ".nvdb" || options.report.extension() != ".json") {
    throw std::runtime_error("input/output/report extensions must be .vdb/.nvdb/.json");
  }
  if (options.quantization != "LOSSLESS" && options.quantization != "FP16") {
    throw std::runtime_error("quantization must be LOSSLESS or FP16");
  }
  if (options.grids.size() > MAX_GRIDS || std::set<std::string>(options.grids.begin(), options.grids.end()).size() != options.grids.size()) {
    throw std::runtime_error("selected grid list is duplicated or exceeds the budget");
  }
  return options;
}

static void check_interrupted(const Options &options,
                              const std::chrono::steady_clock::time_point started,
                              const char *stage)
{
  if (interrupted.load(std::memory_order_relaxed) ||
      (!options.cancel_file.empty() && fs::exists(options.cancel_file)))
  {
    throw std::runtime_error(std::string("conversion cancelled during ") + stage);
  }
  if (options.timeout_ms > 0) {
    const uint64_t elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                                 std::chrono::steady_clock::now() - started)
                                 .count();
    if (elapsed >= options.timeout_ms) {
      throw std::runtime_error(std::string("conversion timed out during ") + stage);
    }
  }
}

static std::string grid_class_name(openvdb::GridClass grid_class)
{
  switch (grid_class) {
    case openvdb::GRID_LEVEL_SET: return "LEVEL_SET";
    case openvdb::GRID_FOG_VOLUME: return "FOG_VOLUME";
    case openvdb::GRID_STAGGERED: return "STAGGERED";
    default: return "UNKNOWN";
  }
}

static openvdb::BBoxd world_bounds(const openvdb::GridBase &grid, const openvdb::CoordBBox &bbox)
{
  auto corner_point = [&](int corner) {
    return openvdb::Vec3d(
        corner & 1 ? bbox.max().x() + 1.0 : bbox.min().x(),
        corner & 2 ? bbox.max().y() + 1.0 : bbox.min().y(),
        corner & 4 ? bbox.max().z() + 1.0 : bbox.min().z());
  };
  const openvdb::Vec3d first = grid.transform().indexToWorld(corner_point(0));
  openvdb::BBoxd result(first, first);
  for (int corner = 1; corner < 8; ++corner) {
    result.expand(grid.transform().indexToWorld(corner_point(corner)));
  }
  return result;
}

static std::array<double, 16> index_to_world(const openvdb::GridBase &grid)
{
  const openvdb::Vec3d origin = grid.transform().indexToWorld(openvdb::Vec3d(0.0));
  const openvdb::Vec3d x = grid.transform().indexToWorld(openvdb::Vec3d(1.0, 0.0, 0.0)) - origin;
  const openvdb::Vec3d y = grid.transform().indexToWorld(openvdb::Vec3d(0.0, 1.0, 0.0)) - origin;
  const openvdb::Vec3d z = grid.transform().indexToWorld(openvdb::Vec3d(0.0, 0.0, 1.0)) - origin;
  return {x.x(), y.x(), z.x(), origin.x(),
          x.y(), y.y(), z.y(), origin.y(),
          x.z(), y.z(), z.z(), origin.z(),
          0.0, 0.0, 0.0, 1.0};
}

static nanovdb::GridHandle<nanovdb::HostBuffer> convert_grid(
    const openvdb::GridBase::Ptr &grid, const std::string &quantization)
{
  if (grid->isType<openvdb::FloatGrid>()) {
    auto typed = openvdb::GridBase::grid<openvdb::FloatGrid>(grid);
    if (quantization == "FP16") {
      nanovdb::tools::CreateNanoGrid<openvdb::FloatGrid> converter(*typed);
      converter.setStats(nanovdb::tools::StatsMode::All);
      converter.setChecksum(nanovdb::CheckMode::Full);
      return converter.getHandle<nanovdb::Fp16>();
    }
  }
  else if (!grid->isType<openvdb::Vec3SGrid>()) {
    throw std::runtime_error("unsupported OpenVDB grid type for " + grid->getName() + ": " + grid->valueType());
  }
  if (quantization != "LOSSLESS") {
    throw std::runtime_error("FP16 is supported only for FloatGrid: " + grid->getName());
  }
  return nanovdb::tools::openToNanoVDB(
      grid, nanovdb::tools::StatsMode::All, nanovdb::CheckMode::Full, 0);
}

static void write_vec3(std::ostream &output, const openvdb::Vec3d &value)
{
  output << '[' << value.x() << ',' << value.y() << ',' << value.z() << ']';
}

static void write_coord(std::ostream &output, const openvdb::Coord &value)
{
  output << '[' << value.x() << ',' << value.y() << ',' << value.z() << ']';
}

static void write_report(const Options &options,
                         const fs::path &report_path,
                         const std::vector<GridReport> &grids)
{
  std::ofstream output(report_path, std::ios::out | std::ios::trunc);
  if (!output) throw std::runtime_error("failed to create conversion report");
  char nano_version[16];
  nanovdb::toStr(nano_version, nanovdb::Version());
  output << std::setprecision(17)
         << "{\n  \"schemaVersion\":1,\n"
         << "  \"input\":" << json_string(fs::absolute(options.input).string()) << ",\n"
         << "  \"output\":" << json_string(fs::absolute(options.output).string()) << ",\n"
         << "  \"quantization\":" << json_string(options.quantization) << ",\n"
         << "  \"openVDBVersion\":" << json_string(openvdb::getLibraryVersionString()) << ",\n"
         << "  \"nanoVDBVersion\":" << json_string(nano_version) << ",\n";
  using FloatRootData = nanovdb::RootData<nanovdb::NanoUpper<float>>;
  using FloatUpperData = nanovdb::InternalData<nanovdb::NanoLower<float>, 5>;
  using FloatLowerData = nanovdb::InternalData<nanovdb::NanoLeaf<float>, 4>;
  using FloatLeafData = nanovdb::LeafData<float, nanovdb::Coord, nanovdb::Mask, 3>;
  using Vec3RootData = nanovdb::RootData<nanovdb::NanoUpper<nanovdb::Vec3f>>;
  using Vec3UpperData = nanovdb::InternalData<nanovdb::NanoLower<nanovdb::Vec3f>, 5>;
  using Vec3LowerData = nanovdb::InternalData<nanovdb::NanoLeaf<nanovdb::Vec3f>, 4>;
  using Vec3LeafData = nanovdb::LeafData<nanovdb::Vec3f, nanovdb::Coord, nanovdb::Mask, 3>;
  output << "  \"float32TreeLayout\":{"
         << "\"gridDataBytes\":" << sizeof(nanovdb::GridData)
         << ",\"treeDataBytes\":" << sizeof(nanovdb::TreeData)
         << ",\"treeRootOffsetOffset\":" << offsetof(nanovdb::TreeData, mNodeOffset[3])
         << ",\"rootDataBytes\":" << sizeof(FloatRootData)
         << ",\"rootTableSizeOffset\":" << offsetof(FloatRootData, mTableSize)
         << ",\"rootTileBytes\":" << sizeof(FloatRootData::Tile)
         << ",\"rootTileKeyOffset\":" << offsetof(FloatRootData::Tile, key)
         << ",\"rootTileChildOffset\":" << offsetof(FloatRootData::Tile, child)
         << ",\"rootTileStateOffset\":" << offsetof(FloatRootData::Tile, state)
         << ",\"rootTileValueOffset\":" << offsetof(FloatRootData::Tile, value)
         << ",\"upperNodeBytes\":" << sizeof(FloatUpperData)
         << ",\"upperValueMaskOffset\":" << offsetof(FloatUpperData, mValueMask)
         << ",\"upperChildMaskOffset\":" << offsetof(FloatUpperData, mChildMask)
         << ",\"upperTableOffset\":" << offsetof(FloatUpperData, mTable)
         << ",\"lowerNodeBytes\":" << sizeof(FloatLowerData)
         << ",\"lowerValueMaskOffset\":" << offsetof(FloatLowerData, mValueMask)
         << ",\"lowerChildMaskOffset\":" << offsetof(FloatLowerData, mChildMask)
         << ",\"lowerTableOffset\":" << offsetof(FloatLowerData, mTable)
         << ",\"leafNodeBytes\":" << sizeof(FloatLeafData)
         << ",\"leafValueMaskOffset\":" << offsetof(FloatLeafData, mValueMask)
         << ",\"leafValuesOffset\":" << offsetof(FloatLeafData, mValues)
         << "},\n"
         << "  \"vec3fTreeLayout\":{"
         << "\"gridDataBytes\":" << sizeof(nanovdb::GridData)
         << ",\"treeDataBytes\":" << sizeof(nanovdb::TreeData)
         << ",\"treeRootOffsetOffset\":" << offsetof(nanovdb::TreeData, mNodeOffset[3])
         << ",\"rootDataBytes\":" << sizeof(Vec3RootData)
         << ",\"rootTableSizeOffset\":" << offsetof(Vec3RootData, mTableSize)
         << ",\"rootTileBytes\":" << sizeof(Vec3RootData::Tile)
         << ",\"rootTileKeyOffset\":" << offsetof(Vec3RootData::Tile, key)
         << ",\"rootTileChildOffset\":" << offsetof(Vec3RootData::Tile, child)
         << ",\"rootTileStateOffset\":" << offsetof(Vec3RootData::Tile, state)
         << ",\"rootTileValueOffset\":" << offsetof(Vec3RootData::Tile, value)
         << ",\"upperNodeBytes\":" << sizeof(Vec3UpperData)
         << ",\"upperValueMaskOffset\":" << offsetof(Vec3UpperData, mValueMask)
         << ",\"upperChildMaskOffset\":" << offsetof(Vec3UpperData, mChildMask)
         << ",\"upperTableOffset\":" << offsetof(Vec3UpperData, mTable)
         << ",\"lowerNodeBytes\":" << sizeof(Vec3LowerData)
         << ",\"lowerValueMaskOffset\":" << offsetof(Vec3LowerData, mValueMask)
         << ",\"lowerChildMaskOffset\":" << offsetof(Vec3LowerData, mChildMask)
         << ",\"lowerTableOffset\":" << offsetof(Vec3LowerData, mTable)
         << ",\"leafNodeBytes\":" << sizeof(Vec3LeafData)
         << ",\"leafValueMaskOffset\":" << offsetof(Vec3LeafData, mValueMask)
         << ",\"leafValuesOffset\":" << offsetof(Vec3LeafData, mValues)
         << "},\n"
         << "  \"grids\":[\n";
  for (size_t index = 0; index < grids.size(); ++index) {
    const GridReport &grid = grids[index];
    output << "    {\"name\":" << json_string(grid.name)
           << ",\"sourceType\":" << json_string(grid.source_type)
           << ",\"valueType\":" << json_string(grid.value_type)
           << ",\"gridClass\":" << json_string(grid.grid_class)
           << ",\"activeVoxelCount\":" << grid.active_voxels
           << ",\"segmentByteOffset\":" << grid.segment_offset
           << ",\"segmentByteLength\":" << grid.segment_length
           << ",\"byteOffset\":" << grid.grid_offset
           << ",\"byteLength\":" << grid.grid_length
           << ",\"indexBounds\":{\"min\":";
    write_coord(output, grid.index_bounds.min());
    output << ",\"max\":";
    write_coord(output, grid.index_bounds.max());
    output << "},\"worldBounds\":{\"min\":";
    write_vec3(output, grid.world_bounds.min());
    output << ",\"max\":";
    write_vec3(output, grid.world_bounds.max());
    output << "},\"voxelSize\":";
    write_vec3(output, grid.voxel_size);
    output << ",\"indexToWorld\":[";
    for (size_t matrix_index = 0; matrix_index < grid.index_to_world.size(); ++matrix_index) {
      if (matrix_index) output << ',';
      output << grid.index_to_world[matrix_index];
    }
    output << ']';
    if (!grid.scalar_samples.empty()) {
      output << ",\"scalarSamples\":[";
      for (size_t sample_index = 0; sample_index < grid.scalar_samples.size(); ++sample_index) {
        const auto &sample = grid.scalar_samples[sample_index];
        if (sample_index) output << ',';
        output << "{\"coord\":";
        write_coord(output, sample.coord);
        output << ",\"value\":" << sample.value << ",\"active\":" << (sample.active ? "true" : "false") << '}';
      }
      output << ']';
    }
    if (!grid.vector_samples.empty()) {
      output << ",\"vectorSamples\":[";
      for (size_t sample_index = 0; sample_index < grid.vector_samples.size(); ++sample_index) {
        const auto &sample = grid.vector_samples[sample_index];
        if (sample_index) output << ',';
        output << "{\"coord\":";
        write_coord(output, sample.coord);
        output << ",\"value\":[" << sample.value[0] << ',' << sample.value[1] << ',' << sample.value[2]
               << "],\"active\":" << (sample.active ? "true" : "false") << '}';
      }
      output << ']';
    }
    output << '}' << (index + 1 == grids.size() ? "\n" : ",\n");
  }
  output << "  ]\n}\n";
  if (!output) throw std::runtime_error("failed to write conversion report");
}

int main(int argc, char **argv)
{
  fs::path staged_output;
  fs::path staged_report;
  try {
    const Options options = parse_options(argc, argv);
    const auto started = std::chrono::steady_clock::now();
    std::signal(SIGINT, request_interrupt);
    std::signal(SIGTERM, request_interrupt);
    check_interrupted(options, started, "startup");
    if (!fs::is_regular_file(options.input)) throw std::runtime_error("input VDB does not exist");
    const uint64_t source_size = fs::file_size(options.input);
    if (source_size == 0 || source_size > MAX_SOURCE_BYTES) throw std::runtime_error("input VDB exceeds the source byte budget");
    fs::create_directories(fs::absolute(options.output).parent_path());
    fs::create_directories(fs::absolute(options.report).parent_path());
    const std::string stage_suffix = "." + std::to_string(static_cast<uint64_t>(getpid())) + ".stage";
    staged_output = options.output.string() + stage_suffix;
    staged_report = options.report.string() + stage_suffix;
    fs::remove(staged_output);
    fs::remove(staged_report);

    openvdb::initialize();
    openvdb::io::File input(options.input.string());
    input.open(false);
    check_interrupted(options, started, "OpenVDB inventory");
    openvdb::GridPtrVecPtr source_grids = input.getGrids();
    if (!source_grids || source_grids->empty() || source_grids->size() > MAX_GRIDS) throw std::runtime_error("VDB grid count exceeds the budget");

    const std::set<std::string> selected(options.grids.begin(), options.grids.end());
    std::set<std::string> found;
    uint64_t total_active_voxels = 0;
    std::ofstream output(staged_output, std::ios::binary | std::ios::trunc);
    if (!output) throw std::runtime_error("failed to create NanoVDB output");
    std::vector<GridReport> report;
    for (const openvdb::GridBase::Ptr &grid : *source_grids) {
      check_interrupted(options, started, "grid inventory");
      if (!selected.empty() && !selected.count(grid->getName())) continue;
      if (!found.insert(grid->getName()).second) throw std::runtime_error("duplicate source grid name: " + grid->getName());
      total_active_voxels += grid->activeVoxelCount();
      if (total_active_voxels > MAX_ACTIVE_VOXELS) throw std::runtime_error("active voxel budget exceeded");

      auto handle = convert_grid(grid, options.quantization);
      check_interrupted(options, started, "NanoVDB conversion");
      const uint64_t segment_offset = static_cast<uint64_t>(output.tellp());
      nanovdb::io::writeGrid(output, handle, nanovdb::io::Codec::NONE);
      const uint64_t segment_end = static_cast<uint64_t>(output.tellp());
      const uint64_t name_size = grid->getName().size() + 1;
      const uint64_t grid_offset = segment_offset + sizeof(nanovdb::io::FileHeader) + sizeof(nanovdb::io::FileMetaData) + name_size;
      if (grid_offset + handle.gridSize() != segment_end) throw std::runtime_error("unexpected NanoVDB segment layout");

      GridReport item;
      item.name = grid->getName();
      item.source_type = grid->valueType();
      item.value_type = options.quantization == "FP16" ? "FLOAT16" : grid->isType<openvdb::FloatGrid>() ? "FLOAT32" : "VEC3F32";
      item.grid_class = grid_class_name(grid->getGridClass());
      item.active_voxels = grid->activeVoxelCount();
      item.segment_offset = segment_offset;
      item.segment_length = segment_end - segment_offset;
      item.grid_offset = grid_offset;
      item.grid_length = handle.gridSize();
      item.index_bounds = grid->evalActiveVoxelBoundingBox();
      item.world_bounds = world_bounds(*grid, item.index_bounds);
      item.voxel_size = grid->voxelSize();
      item.index_to_world = index_to_world(*grid);
      if (options.quantization == "LOSSLESS" && grid->isType<openvdb::FloatGrid>()) {
        const nanovdb::NanoGrid<float> *nano_grid = handle.grid<float>();
        if (!nano_grid) throw std::runtime_error("NanoVDB Float32 grid payload is unavailable");
        const std::array<openvdb::Coord, 5> sample_coords = {
            item.index_bounds.min(), openvdb::Coord(0, 0, 0), item.index_bounds.max(),
            openvdb::Coord(item.index_bounds.min().x() - 1, 0, 0),
            openvdb::Coord(item.index_bounds.max().x() + 1, 0, 0)};
        for (const openvdb::Coord &coord : sample_coords) {
          float value = 0.0f;
          const bool active = nano_grid->tree().probeValue(nanovdb::Coord(coord.x(), coord.y(), coord.z()), value);
          item.scalar_samples.push_back({coord, value, active});
        }
      }
      else if (options.quantization == "LOSSLESS" && grid->isType<openvdb::Vec3SGrid>()) {
        const nanovdb::NanoGrid<nanovdb::Vec3f> *nano_grid = handle.grid<nanovdb::Vec3f>();
        if (!nano_grid) throw std::runtime_error("NanoVDB Vec3f grid payload is unavailable");
        const std::array<openvdb::Coord, 5> sample_coords = {
            item.index_bounds.min(), openvdb::Coord(0, 0, 0), item.index_bounds.max(),
            openvdb::Coord(item.index_bounds.min().x() - 1, 0, 0),
            openvdb::Coord(item.index_bounds.max().x() + 1, 0, 0)};
        for (const openvdb::Coord &coord : sample_coords) {
          nanovdb::Vec3f value(0.0f);
          const bool active = nano_grid->tree().probeValue(nanovdb::Coord(coord.x(), coord.y(), coord.z()), value);
          item.vector_samples.push_back({coord, {value[0], value[1], value[2]}, active});
        }
      }
      report.push_back(std::move(item));

      if (segment_end > MAX_OUTPUT_BYTES) throw std::runtime_error("NanoVDB output exceeds the byte budget");
    }
    input.close();
    output.close();
    check_interrupted(options, started, "artifact commit");
    if (!selected.empty() && found != selected) throw std::runtime_error("one or more selected grids were not found");
    if (report.empty()) throw std::runtime_error("no supported grids were selected");
    write_report(options, staged_report, report);
    check_interrupted(options, started, "report commit");
    fs::rename(staged_output, options.output);
    fs::rename(staged_report, options.report);
    openvdb::uninitialize();
    std::cout << "vdb-to-nanovdb-ok input=" << options.input.string()
              << " output=" << options.output.string()
              << " grids=" << report.size()
              << " bytes=" << fs::file_size(options.output) << '\n';
    return 0;
  }
  catch (const std::exception &error) {
    if (!staged_output.empty()) fs::remove(staged_output);
    if (!staged_report.empty()) fs::remove(staged_report);
    std::cerr << "VDB_CONVERSION_FAILED: " << error.what() << '\n';
    return 1;
  }
}