#include "data/store/ChunkedVolumeStore.hpp"

#include <QByteArray>
#include <nlohmann/json.hpp>

#include <cstring>
#include <filesystem>
#include <fstream>
#include <stdexcept>

#include "core/algo/GprVolumeBuilder.hpp"  // geopro::core::BuiltI16

namespace geopro::data {

namespace {
using nlohmann::json;
namespace fs = std::filesystem;

constexpr const char* kMetaFile = "meta.json";
constexpr const char* kDataFile = "data.bin";

int ceilDiv(int n, int brick) { return (n + brick - 1) / brick; }

// 块尺寸（边缘块 < brick）：第 b 块沿该轴的体素数。
int extent(int n, int b, int brick) {
  const int got = n - b * brick;
  return got < brick ? got : brick;
}

// 从体中拷出一块的 int16（块内 i 最快、k 最慢，与体一致）。
std::vector<std::int16_t> sliceBrick(const geopro::core::ScalarVolumeI16& vol,
                                     int bx, int by, int bz, int brick,
                                     int bw, int bh, int bd) {
  std::vector<std::int16_t> out(static_cast<std::size_t>(bw) * bh * bd);
  const int i0 = bx * brick, j0 = by * brick, k0 = bz * brick;
  std::size_t w = 0;
  for (int kk = 0; kk < bd; ++kk) {
    for (int jj = 0; jj < bh; ++jj) {
      for (int ii = 0; ii < bw; ++ii) {
        out[w++] = vol.at(i0 + ii, j0 + jj, k0 + kk);
      }
    }
  }
  return out;
}
}  // namespace

void ChunkedVolumeStore::write(const std::string& dir,
                               const geopro::core::BuiltI16& b, int brick) {
  if (brick <= 0) throw std::invalid_argument("ChunkedVolumeStore: brick must be > 0");

  const auto& vol = b.vol;
  const int nx = vol.nx(), ny = vol.ny(), nz = vol.nz();

  fs::create_directories(fs::path(dir));

  const int bX = ceilDiv(nx, brick);
  const int bY = ceilDiv(ny, brick);
  const int bZ = ceilDiv(nz, brick);

  std::ofstream data((fs::path(dir) / kDataFile).string(),
                     std::ios::binary | std::ios::trunc);
  if (!data) throw std::runtime_error("ChunkedVolumeStore: cannot open data.bin for write");

  json bricks = json::array();
  std::int64_t offset = 0;

  // 固定遍历顺序：bz 最慢、bx 最快（与 brickIndex 一致）。
  for (int bz = 0; bz < bZ; ++bz) {
    for (int by = 0; by < bY; ++by) {
      for (int bx = 0; bx < bX; ++bx) {
        const int bw = extent(nx, bx, brick);
        const int bh = extent(ny, by, brick);
        const int bd = extent(nz, bz, brick);
        auto raw = sliceBrick(vol, bx, by, bz, brick, bw, bh, bd);

        const int rawBytes =
            static_cast<int>(raw.size() * sizeof(std::int16_t));
        const QByteArray compressed = qCompress(
            reinterpret_cast<const uchar*>(raw.data()), rawBytes);
        const std::int64_t clen = compressed.size();

        data.write(compressed.constData(), clen);
        if (!data) throw std::runtime_error("ChunkedVolumeStore: data.bin write failed");

        bricks.push_back(json{{"offset", offset},
                              {"compressedLen", clen},
                              {"bw", bw},
                              {"bh", bh},
                              {"bd", bd}});
        offset += clen;
      }
    }
  }
  data.close();

  json meta;
  meta["nx"] = nx;
  meta["ny"] = ny;
  meta["nz"] = nz;
  meta["brick"] = brick;
  meta["origin"] = {b.origin[0], b.origin[1], b.origin[2]};
  meta["spacing"] = {b.spacing[0], b.spacing[1], b.spacing[2]};
  meta["quant"] = {{"scale", b.quant.scale}, {"offset", b.quant.offset}};
  meta["vminPhys"] = b.vminPhys;
  meta["vmaxPhys"] = b.vmaxPhys;
  meta["bricks"] = std::move(bricks);

  std::ofstream metaOut((fs::path(dir) / kMetaFile).string(),
                        std::ios::trunc);
  if (!metaOut) throw std::runtime_error("ChunkedVolumeStore: cannot open meta.json for write");
  metaOut << meta.dump(2);
  if (!metaOut) throw std::runtime_error("ChunkedVolumeStore: meta.json write failed");
}

StoreMeta ChunkedVolumeStore::readMeta(const std::string& dir) {
  std::ifstream in((fs::path(dir) / kMetaFile).string());
  if (!in) throw std::runtime_error("ChunkedVolumeStore: cannot open meta.json");
  json meta;
  in >> meta;

  StoreMeta m;
  m.nx = meta.at("nx").get<int>();
  m.ny = meta.at("ny").get<int>();
  m.nz = meta.at("nz").get<int>();
  m.brick = meta.at("brick").get<int>();
  const auto& o = meta.at("origin");
  m.origin = {o[0].get<double>(), o[1].get<double>(), o[2].get<double>()};
  const auto& s = meta.at("spacing");
  m.spacing = {s[0].get<double>(), s[1].get<double>(), s[2].get<double>()};
  m.quant.scale = meta.at("quant").at("scale").get<double>();
  m.quant.offset = meta.at("quant").at("offset").get<double>();
  m.vminPhys = meta.at("vminPhys").get<double>();
  m.vmaxPhys = meta.at("vmaxPhys").get<double>();
  return m;
}

ChunkedVolumeStore::ChunkedVolumeStore(const std::string& dir) : dir_(dir) {
  std::ifstream in((fs::path(dir) / kMetaFile).string());
  if (!in) throw std::runtime_error("ChunkedVolumeStore: cannot open meta.json");
  json meta;
  in >> meta;

  meta_.nx = meta.at("nx").get<int>();
  meta_.ny = meta.at("ny").get<int>();
  meta_.nz = meta.at("nz").get<int>();
  meta_.brick = meta.at("brick").get<int>();
  const auto& o = meta.at("origin");
  meta_.origin = {o[0].get<double>(), o[1].get<double>(), o[2].get<double>()};
  const auto& sp = meta.at("spacing");
  meta_.spacing = {sp[0].get<double>(), sp[1].get<double>(), sp[2].get<double>()};
  meta_.quant.scale = meta.at("quant").at("scale").get<double>();
  meta_.quant.offset = meta.at("quant").at("offset").get<double>();
  meta_.vminPhys = meta.at("vminPhys").get<double>();
  meta_.vmaxPhys = meta.at("vmaxPhys").get<double>();

  bricksX_ = ceilDiv(meta_.nx, meta_.brick);
  bricksY_ = ceilDiv(meta_.ny, meta_.brick);
  bricksZ_ = ceilDiv(meta_.nz, meta_.brick);

  const auto& arr = meta.at("bricks");
  bricks_.reserve(arr.size());
  for (const auto& e : arr) {
    BrickEntry be;
    be.offset = e.at("offset").get<std::int64_t>();
    be.compressedLen = e.at("compressedLen").get<std::int64_t>();
    be.bw = e.at("bw").get<int>();
    be.bh = e.at("bh").get<int>();
    be.bd = e.at("bd").get<int>();
    bricks_.push_back(be);
  }
}

std::vector<std::int16_t> ChunkedVolumeStore::readBrick(int bx, int by,
                                                        int bz) const {
  if (bx < 0 || by < 0 || bz < 0 || bx >= bricksX_ || by >= bricksY_ ||
      bz >= bricksZ_) {
    throw std::out_of_range("ChunkedVolumeStore::readBrick: brick index out of range");
  }
  const BrickEntry& be = bricks_.at(static_cast<std::size_t>(brickIndex(bx, by, bz)));

  std::ifstream data((fs::path(dir_) / kDataFile).string(), std::ios::binary);
  if (!data) throw std::runtime_error("ChunkedVolumeStore: cannot open data.bin");
  data.seekg(static_cast<std::streamoff>(be.offset), std::ios::beg);

  QByteArray compressed;
  compressed.resize(static_cast<int>(be.compressedLen));
  data.read(compressed.data(), be.compressedLen);
  if (!data) throw std::runtime_error("ChunkedVolumeStore: data.bin read failed");

  const QByteArray raw = qUncompress(compressed);
  const std::size_t count =
      static_cast<std::size_t>(be.bw) * be.bh * be.bd;
  if (static_cast<std::size_t>(raw.size()) != count * sizeof(std::int16_t)) {
    throw std::runtime_error("ChunkedVolumeStore: decompressed size mismatch");
  }

  std::vector<std::int16_t> out(count);
  std::memcpy(out.data(), raw.constData(), raw.size());
  return out;
}

}  // namespace geopro::data