#include "VISIBuilder.hpp"
#include "logvisor/logvisor.hpp"

#ifndef _WIN32
#include <unistd.h>
#include <signal.h>
#endif

#define VISI_MAX_LEVEL 10
#define VISI_MIN_LENGTH 8.0

static logvisor::Module Log("VISIBuilder");

VISIBuilder::PVSRenderCache::PVSRenderCache(VISIRenderer& renderer) : m_renderer(renderer) { m_cache.reserve(1000); }

static std::unique_ptr<VISIRenderer::RGBA8[]> RGBABuf(new VISIRenderer::RGBA8[256 * 256 * 6]);

const VISIBuilder::Leaf& VISIBuilder::PVSRenderCache::GetLeaf(const zeus::CVector3f& vec) {
  auto search = m_cache.find(vec);
  if (search != m_cache.cend()) {
    // Log.report(logvisor::Info, FMT_STRING("Cache hit"));
    return *search->second;
  }

  // Log.report(logvisor::Info, FMT_STRING("Rendering"));
  bool needsTransparent = false;
  m_renderer.RenderPVSOpaque(RGBABuf.get(), vec, needsTransparent);
  std::unique_ptr<Leaf> leafOut = std::make_unique<Leaf>();
  for (unsigned i = 0; i < 768 * 512; ++i) {
    const VISIRenderer::RGBA8& pixel = RGBABuf[i];
    uint32_t id = (pixel.b << 16) | (pixel.g << 8) | pixel.r;
    if (id != 0)
      leafOut->setBit(id - 1);
  }

  auto setBitLambda = [&](int idx) { leafOut->setBit(idx); };
  auto setLightLambda = [&](int idx, EPVSVisSetState state) {
    if (state != EPVSVisSetState::EndOfTree)
      leafOut->setLightEnum(m_lightMetaBit + idx * 2, state);
  };
  if (needsTransparent)
    m_renderer.RenderPVSTransparent(setBitLambda, vec);
  m_renderer.RenderPVSEntitiesAndLights(setBitLambda, setLightLambda, vec);

  return *m_cache.emplace(std::make_pair(vec, std::move(leafOut))).first->second;
}

void VISIBuilder::Progress::report(int divisions) {
  m_prog += 1.f / divisions;
  // printf(" %g%%        \r", m_prog * 100.f);
  // fflush(stdout);
  if (m_updatePercent)
    m_updatePercent(m_prog);
}

void VISIBuilder::Node::buildChildren(int level, int divisions, const zeus::CAABox& curAabb, PVSRenderCache& rc,
                                      Progress& prog, const std::function<bool()>& terminate) {
  if (terminate())
    return;

  // Recurse in while building node structure
  if (level < VISI_MAX_LEVEL) {
    // Heuristic split
    int splits[3];
    splits[0] = (curAabb.max.x() - curAabb.min.x() >= VISI_MIN_LENGTH) ? 2 : 1;
    splits[1] = (curAabb.max.y() - curAabb.min.y() >= VISI_MIN_LENGTH) ? 2 : 1;
    splits[2] = (curAabb.max.z() - curAabb.min.z() >= VISI_MIN_LENGTH) ? 2 : 1;

    if (splits[0] == 2)
      flags |= 0x1;
    if (splits[1] == 2)
      flags |= 0x2;
    if (splits[2] == 2)
      flags |= 0x4;

    int thisdiv = splits[0] * splits[1] * splits[2] * divisions;

    if (flags) {
      childNodes.resize(8);

      // Inward subdivide
      zeus::CAABox Z[2];
      if (flags & 0x4)
        curAabb.splitZ(Z[0], Z[1]);
      else
        Z[0] = curAabb;
      for (int i = 0; i < splits[2]; ++i) {
        zeus::CAABox Y[2];
        if (flags & 0x2)
          Z[i].splitY(Y[0], Y[1]);
        else
          Y[0] = Z[i];
        for (int j = 0; j < splits[1]; ++j) {
          zeus::CAABox X[2];
          if (flags & 0x1)
            Y[j].splitX(X[0], X[1]);
          else
            X[0] = Y[j];
          for (int k = 0; k < splits[0]; ++k) {
            childNodes[i * 4 + j * 2 + k].buildChildren(level + 1, thisdiv, X[k], rc, prog, terminate);
          }
        }
      }

      // Outward unsubdivide for like-leaves
      for (int i = 0; i < 3; ++i) {
        if (flags & 0x4 && childNodes[0] == childNodes[4] && (!(flags & 0x1) || childNodes[1] == childNodes[5]) &&
            (!(flags & 0x2) || childNodes[2] == childNodes[6]) && (!(flags & 0x3) || childNodes[3] == childNodes[7])) {
          flags &= ~0x4;
          // Log.report(logvisor::Info, FMT_STRING("Unsub Z"));
          continue;
        }
        if (flags & 0x2 && childNodes[0] == childNodes[2] && (!(flags & 0x1) || childNodes[1] == childNodes[3]) &&
            (!(flags & 0x4) || childNodes[4] == childNodes[6]) && (!(flags & 0x5) || childNodes[5] == childNodes[7])) {
          flags &= ~0x2;
          // Log.report(logvisor::Info, FMT_STRING("Unsub Y"));
          continue;
        }
        if (flags & 0x1 && childNodes[0] == childNodes[1] && (!(flags & 0x2) || childNodes[2] == childNodes[3]) &&
            (!(flags & 0x4) || childNodes[4] == childNodes[5]) && (!(flags & 0x6) || childNodes[6] == childNodes[7])) {
          flags &= ~0x1;
          // Log.report(logvisor::Info, FMT_STRING("Unsub X"));
          continue;
        }
        break;
      }

      if (!flags) {
        // This is now a leaf node
        for (int i = 0; i < 8; ++i)
          leaf |= childNodes[i].leaf;
        // Log.report(logvisor::Info, FMT_STRING("Leaf Promote"));
        return;
      }
    }
  }

  if (!flags) {
    // This is a child node
    zeus::CVector3f center = curAabb.center();
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), curAabb.min.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), curAabb.min.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), curAabb.min.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), center.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), center.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), center.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), curAabb.max.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), curAabb.max.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), curAabb.max.y(), curAabb.min.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), curAabb.min.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), curAabb.min.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), curAabb.min.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), center.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), center.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), center.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), curAabb.max.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), curAabb.max.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), curAabb.max.y(), center.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), curAabb.min.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), curAabb.min.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), curAabb.min.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), center.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), center.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), center.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.min.x(), curAabb.max.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(center.x(), curAabb.max.y(), curAabb.max.z()));
    leaf |= rc.GetLeaf(zeus::CVector3f(curAabb.max.x(), curAabb.max.y(), curAabb.max.z()));

    prog.report(divisions);
  }
}

static const int NumChildTable[] = {0, 2, 2, 4, 2, 4, 4, 8};

void VISIBuilder::Node::calculateSizesAndOffs(size_t& cur, size_t leafSz) {
  cur += 1;
  flags |= 0x18;

  if (flags & 0x7) {
    int splits[3];
    splits[0] = (flags & 0x1) ? 2 : 1;
    splits[1] = (flags & 0x2) ? 2 : 1;
    splits[2] = (flags & 0x4) ? 2 : 1;

    // Inward accumulate
    const size_t startCur = cur;
    size_t maxDelta = 0;
    for (int i = 0; i < splits[2]; ++i)
      for (int j = 0; j < splits[1]; ++j)
        for (int k = 0; k < splits[0]; ++k) {
          const size_t nodeSel = i * 4 + j * 2 + k;
          const size_t delta = cur - startCur;
          if (delta > maxDelta)
            maxDelta = delta;
          childRelOffs[nodeSel] = delta;
          childNodes[nodeSel].calculateSizesAndOffs(cur, leafSz);
        }

    const int numChildren = NumChildTable[flags & 0x7];
    if (maxDelta > 0xffff) {
      cur += (numChildren - 1) * 3;
      flags |= 0x40;
    } else if (maxDelta > 0xff) {
      cur += (numChildren - 1) * 2;
    } else {
      cur += numChildren - 1;
      flags |= 0x20;
    }
  } else {
    if (!leaf)
      flags &= ~0x8;
    else
      cur += leafSz;
  }
}

void VISIBuilder::Node::writeNodes(athena::io::MemoryWriter& w, size_t leafBytes) const {
  w.writeUByte(flags);

  if (flags & 0x7) {
    int splits[3];
    splits[0] = (flags & 0x1) ? 2 : 1;
    splits[1] = (flags & 0x2) ? 2 : 1;
    splits[2] = (flags & 0x4) ? 2 : 1;

    // Write offsets
    for (int i = 0; i < splits[2]; ++i)
      for (int j = 0; j < splits[1]; ++j)
        for (int k = 0; k < splits[0]; ++k) {
          const size_t nodeSel = i * 4 + j * 2 + k;
          if (nodeSel == 0)
            continue;
          const size_t offset = childRelOffs[nodeSel];
          if (flags & 0x40) {
            w.writeUByte((offset >> 16) & 0xff);
            w.writeUByte((offset >> 8) & 0xff);
            w.writeUByte(offset & 0xff);
          } else if (flags & 0x20) {
            w.writeUByte(offset & 0xff);
          } else {
            w.writeUint16Big(offset);
          }
        }

    // Inward iterate
    for (int i = 0; i < splits[2]; ++i)
      for (int j = 0; j < splits[1]; ++j)
        for (int k = 0; k < splits[0]; ++k) {
          const size_t nodeSel = i * 4 + j * 2 + k;
          childNodes[nodeSel].writeNodes(w, leafBytes);
        }
  } else if (leaf) {
    leaf.write(w, leafBytes);
  }
}

std::vector<uint8_t> VISIBuilder::build(const zeus::CAABox& fullAabb, size_t modelCount,
                                        const std::vector<VISIRenderer::Entity>& entities,
                                        const std::vector<VISIRenderer::Light>& lights, size_t layer2LightCount,
                                        FPercent updatePercent, ProcessType parentPid) {
  // Log.report(logvisor::Info, FMT_STRING("Started!"));

  size_t featureCount = modelCount + entities.size();
  renderCache.m_lightMetaBit = featureCount;

  Progress prog(updatePercent);
#ifndef _WIN32
  auto terminate = [this, parentPid]() {
    return renderCache.m_renderer.m_terminate || (parentPid ? kill(parentPid, 0) : false);
  };
#else
  auto terminate = [this, parentPid]() {
    DWORD exitCode = 0;
    if (!GetExitCodeProcess(parentPid, &exitCode))
      return renderCache.m_renderer.m_terminate;
    return renderCache.m_renderer.m_terminate || (parentPid ? exitCode != STILL_ACTIVE : false);
  };
#endif
  rootNode.buildChildren(0, 1, fullAabb, renderCache, prog, terminate);
  if (terminate())
    return {};

  // Lights cache their CPVSVisSet result enum as 2 bits
  size_t leafBitsCount = featureCount + lights.size() * 2;
  size_t leafBytesCount = ROUND_UP_8(leafBitsCount) / 8;

  // Calculate octree size and store relative offsets
  size_t octreeSz = 0;
  rootNode.calculateSizesAndOffs(octreeSz, leafBytesCount);

  size_t visiSz = 34 + entities.size() * 4 + lights.size() * leafBytesCount + 36 + octreeSz;
  size_t roundedVisiSz = ROUND_UP_32(visiSz);

  std::vector<uint8_t> dataOut(roundedVisiSz, 0);
  athena::io::MemoryWriter w(dataOut.data(), roundedVisiSz);
  w.writeUint32Big('VISI');
  w.writeUint32Big(2);
  w.writeBool(true);
  w.writeBool(true);
  w.writeUint32Big(featureCount);
  w.writeUint32Big(lights.size());
  w.writeUint32Big(layer2LightCount);
  w.writeUint32Big(entities.size());
  w.writeUint32Big(leafBytesCount);
  w.writeUint32Big(lights.size());

  for (const VISIRenderer::Entity& e : entities) {
    w.writeUint32Big(e.entityId);
  }

  for (const VISIRenderer::Light& l : lights) {
    const VISIBuilder::Leaf& leaf = renderCache.GetLeaf(l.point);
    leaf.write(w, leafBytesCount);
  }

  w.writeVec3fBig(fullAabb.min);
  w.writeVec3fBig(fullAabb.max);
  w.writeUint32Big(featureCount + lights.size());
  w.writeUint32Big(lights.size());
  w.writeUint32Big(octreeSz);
  rootNode.writeNodes(w, leafBytesCount);

  w.seekAlign32();

  // Log.report(logvisor::Info, FMT_STRING("Finished!"));
  return dataOut;
}