#include <openvdb/openvdb.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace fs = std::filesystem;

namespace {

constexpr int kImageSize = 64;
constexpr std::array<float, 3> kColor = {0.72f, 0.78f, 0.86f};

float sample_linear(const openvdb::FloatGrid::ConstAccessor &accessor,
                    const std::array<float, 3> &position)
{
  const std::array<int, 3> base = {
      static_cast<int>(std::floor(position[0])),
      static_cast<int>(std::floor(position[1])),
      static_cast<int>(std::floor(position[2])),
  };
  const std::array<float, 3> fraction = {
      position[0] - static_cast<float>(base[0]),
      position[1] - static_cast<float>(base[1]),
      position[2] - static_cast<float>(base[2]),
  };
  float value = 0.0f;
  for (int x = 0; x < 2; ++x) {
    for (int y = 0; y < 2; ++y) {
      for (int z = 0; z < 2; ++z) {
        const float weight = (x ? fraction[0] : 1.0f - fraction[0]) *
                             (y ? fraction[1] : 1.0f - fraction[1]) *
                             (z ? fraction[2] : 1.0f - fraction[2]);
        value += accessor.getValue(openvdb::Coord(base[0] + x, base[1] + y, base[2] + z)) *
                 weight;
      }
    }
  }
  return value;
}

uint8_t pack_unorm(float value)
{
  return static_cast<uint8_t>(std::lround(std::clamp(value, 0.0f, 1.0f) * 255.0f));
}

std::vector<uint8_t> render_axis(const openvdb::FloatGrid &density,
                                 const openvdb::CoordBBox &bounds,
                                 int view_axis)
{
  const auto accessor = density.getConstAccessor();
  const openvdb::Coord minimum = bounds.min();
  const openvdb::Coord maximum = bounds.max();
  const std::array<int, 3> plane_a = view_axis == 0 ? std::array<int, 3>{1, 2, 0} :
                                      view_axis == 1 ? std::array<int, 3>{0, 2, 1} :
                                                       std::array<int, 3>{0, 1, 2};
  const int ray_axis = plane_a[2];
  const int ray_min = minimum[ray_axis];
  const int ray_max = maximum[ray_axis];
  const int ray_count = std::max(1, ray_max - ray_min + 1);
  const int stride = std::max(1, (ray_count + 255) / 256);
  const float voxel_size = std::max(0.01f, static_cast<float>(density.voxelSize()[ray_axis]));
  const float phase = 1.0f / 12.5663706f;
  const float source_scale = 0.5f + 8.0f * phase;
  std::vector<uint8_t> pixels(kImageSize * kImageSize * 4, 0);

  for (int y = 0; y < kImageSize; ++y) {
    for (int x = 0; x < kImageSize; ++x) {
      const std::array<int, 2> plane_min = {minimum[plane_a[0]], minimum[plane_a[1]]};
      const std::array<int, 2> plane_max = {maximum[plane_a[0]], maximum[plane_a[1]]};
      const std::array<float, 2> extent = {
          static_cast<float>(plane_max[0] - plane_min[0] + 1),
          static_cast<float>(plane_max[1] - plane_min[1] + 1),
      };
      const std::array<float, 2> plane_position = {
          static_cast<float>(plane_min[0]) +
              ((static_cast<float>(x) + 0.5f) / static_cast<float>(kImageSize)) * extent[0] - 0.5f,
          static_cast<float>(plane_min[1]) +
              ((static_cast<float>(y) + 0.5f) / static_cast<float>(kImageSize)) * extent[1] - 0.5f,
      };

      float transmittance = 1.0f;
      std::array<float, 3> radiance = {0.0f, 0.0f, 0.0f};
      for (int ray = ray_min; ray <= ray_max; ray += stride) {
        std::array<float, 3> position{};
        position[plane_a[0]] = plane_position[0];
        position[plane_a[1]] = plane_position[1];
        position[ray_axis] = static_cast<float>(ray) + 0.5f;
        const float sampled_density = std::max(0.0f, sample_linear(accessor, position));
        const float alpha = 1.0f - std::exp(-sampled_density * voxel_size * static_cast<float>(stride));
        for (int channel = 0; channel < 3; ++channel) {
          radiance[channel] += transmittance * alpha * kColor[channel] * source_scale;
        }
        transmittance *= 1.0f - alpha;
        if (transmittance < 0.005f) {
          break;
        }
      }

      const size_t offset = static_cast<size_t>(y * kImageSize + x) * 4;
      pixels[offset] = pack_unorm(radiance[0]);
      pixels[offset + 1] = pack_unorm(radiance[1]);
      pixels[offset + 2] = pack_unorm(radiance[2]);
      pixels[offset + 3] = pack_unorm(1.0f - transmittance);
    }
  }
  return pixels;
}

void write_bytes(const fs::path &path, const std::vector<uint8_t> &bytes)
{
  std::ofstream output(path, std::ios::binary | std::ios::trunc);
  if (!output) {
    throw std::runtime_error("failed to create golden image: " + path.string());
  }
  output.write(reinterpret_cast<const char *>(bytes.data()), static_cast<std::streamsize>(bytes.size()));
  if (!output) {
    throw std::runtime_error("failed to write golden image: " + path.string());
  }
}

}  // namespace

int main(int argc, char **argv)
{
  if (argc != 3) {
    std::cerr << "usage: vdb_volume_golden INPUT.vdb OUTPUT_PREFIX\n";
    return 2;
  }

  try {
    openvdb::initialize();
    const fs::path input = fs::absolute(argv[1]);
    const fs::path output_prefix = fs::absolute(argv[2]);
    if (input.extension() != ".vdb") {
      throw std::runtime_error("input extension must be .vdb");
    }
    fs::create_directories(output_prefix.parent_path());

    openvdb::io::File file(input.string());
    file.open(false);
    const openvdb::GridBase::Ptr base = file.readGrid("density");
    file.close();
    const openvdb::FloatGrid::Ptr density = openvdb::gridPtrCast<openvdb::FloatGrid>(base);
    if (!density || density->getGridClass() != openvdb::GRID_FOG_VOLUME) {
      throw std::runtime_error("density must be an OpenVDB FloatGrid fog volume");
    }
    const openvdb::CoordBBox bounds = density->evalActiveVoxelBoundingBox();
    if (bounds.empty()) {
      throw std::runtime_error("density grid has no active voxels");
    }

    const std::array<const char *, 3> names = {"x", "y", "z"};
    for (int axis = 0; axis < 3; ++axis) {
      write_bytes(output_prefix.string() + "-" + names[axis] + ".rgba",
                  render_axis(*density, bounds, axis));
    }
    std::cout << "{\"schemaVersion\":1,\"openVDBVersion\":\""
              << openvdb::getLibraryVersionString() << "\",\"grid\":\"density\",\"width\":"
              << kImageSize << ",\"height\":" << kImageSize << ",\"activeVoxelCount\":"
              << density->activeVoxelCount() << ",\"indexBounds\":{\"min\":["
              << bounds.min().x() << ',' << bounds.min().y() << ',' << bounds.min().z()
              << "],\"max\":[" << bounds.max().x() << ',' << bounds.max().y() << ','
              << bounds.max().z() << "]}}\n";
    openvdb::uninitialize();
    return 0;
  }
  catch (const std::exception &error) {
    std::cerr << "VDB_VOLUME_GOLDEN_FAILED: " << error.what() << '\n';
    return 1;
  }
}