#include "procgen/Biome.h"
#include "procgen/PointGraph.h"
#include "procgen/ShapeGrammar.h"

#include <memory>
#include <string>
#include "zeroerr/assert.h"
#include "zeroerr/unittest.h"

using namespace eve::procgen;

namespace {

bool pointSetsEqual(const PointSet& first, const PointSet& second) {
    if (first.getCount() != second.getCount()) return false;
    if (first.attributes().columnCount() != second.attributes().columnCount()) return false;
    for (int index = 0; index < first.getCount(); ++index) {
        const ProcgenPoint& a = first.points()[size_t(index)];
        const ProcgenPoint& b = second.points()[size_t(index)];
        if (a.x != b.x || a.y != b.y || a.z != b.z || a.normalX != b.normalX || a.normalY != b.normalY ||
            a.normalZ != b.normalZ || a.yaw != b.yaw || a.scaleX != b.scaleX || a.scaleY != b.scaleY ||
            a.scaleZ != b.scaleZ || a.density != b.density || a.seed != b.seed)
            return false;
        for (size_t column = 0; column < first.attributes().columnCount(); ++column) {
            const std::string name(first.attributes().columnName(column));
            if (first.attributes().typeOf(name) != second.attributes().typeOf(name) ||
                first.getAttributeType(index, name) != second.getAttributeType(index, name))
                return false;
            switch (*first.attributes().typeOf(name)) {
                case ProcgenAttributeType::Float:
                    if (first.attributes().getFloat(index, name) != second.attributes().getFloat(index, name))
                        return false;
                    break;
                case ProcgenAttributeType::Int:
                    if (first.attributes().getInt(index, name) != second.attributes().getInt(index, name)) return false;
                    break;
                case ProcgenAttributeType::Bool:
                    if (first.attributes().getBool(index, name) != second.attributes().getBool(index, name))
                        return false;
                    break;
                case ProcgenAttributeType::String:
                    if (first.getStringAttribute(index, name, {}) != second.getStringAttribute(index, name, {}))
                        return false;
                    break;
                case ProcgenAttributeType::Vector: {
                    const auto firstValue  = first.attributes().getVector(index, name);
                    const auto secondValue = second.attributes().getVector(index, name);
                    if (firstValue.has_value() != secondValue.has_value()) return false;
                    if (firstValue && (firstValue->x != secondValue->x || firstValue->y != secondValue->y ||
                                       firstValue->z != secondValue->z))
                        return false;
                    break;
                }
            }
        }
    }
    return true;
}

}  // namespace

TEST_CASE("procgen.pointGraphStress.resumesDeepGraphAcrossBudgetWindows") {
    constexpr int nodeCount = 64;
    PointSet input;
    input.add(0.f, 0.f, 0.f);
    PointGraph graph;
    CHECK(graph.addNode("source", "input"));
    CHECK(graph.setNodePoints("source", &input));
    std::string previous = "source";
    for (int index = 0; index < nodeCount; ++index) {
        const std::string id = "move-" + std::to_string(index);
        CHECK(graph.addNode(id, "transform"));
        CHECK(graph.connect(previous, id));
        CHECK(graph.setNodeFloat(id, "x", 0.25f));
        previous = id;
    }
    CHECK(graph.validate());
    graph.setExecutionNodeBudget(17);

    std::unique_ptr<PointSet> output;
    int budgetWindows = 0;
    while (!output && budgetWindows < 20) {
        output.reset(graph.execute(previous));
        ++budgetWindows;
        if (!output) CHECK(graph.wasCancelled());
    }
    REQUIRE(bool(output));
    CHECK_EQ(output->getX(0), 16.f);
    CHECK_EQ(budgetWindows, 4);
    CHECK(graph.getCacheHitCount() > 0);
}

TEST_CASE("procgen.pointGraphStress.computePolicyPreservesDeterministicCpuPath") {
    PointSet input;
    input.add(2.f, 3.f, 4.f);
    PointGraph graph;
    CHECK(graph.addNode("source", "input"));
    CHECK(graph.addNode("move", "transform"));
    CHECK(graph.connect("source", "move"));
    CHECK(graph.setNodePoints("source", &input));
    CHECK(graph.setNodeFloat("move", "x", 5.f));
    CHECK(!graph.setComputePolicy("invalid"));
    CHECK(graph.setComputePolicy("cpu"));
    graph.setComputeMinimumPoints(0);

    std::unique_ptr<PointSet> output(graph.execute("move"));
    REQUIRE(bool(output));
    CHECK_EQ(output->getX(0), 7.f);
    CHECK_EQ(graph.getComputePolicy(), "cpu");
    CHECK_EQ(graph.getComputeMinimumPoints(), 0);
    CHECK_EQ(graph.getMetricBackend(graph.getMetricCount() - 1), "cpu");
    CHECK(graph.getComputeFallbackReason().empty());
    CHECK_EQ(graph.getComputeUploadCount(), uint64_t(0));
    CHECK_EQ(graph.getComputeDispatchCount(), uint64_t(0));
    CHECK_EQ(graph.getComputeReadbackCount(), uint64_t(0));
    CHECK_EQ(graph.getComputeBufferReuseCount(), uint64_t(0));
    CHECK_EQ(graph.getComputePeakBufferBytes(), uint64_t(0));
    CHECK_EQ(graph.getLastFusedTransformCount(), 0);
    CHECK(graph.debugReport().find("backend=cpu") != std::string::npos);
}

TEST_CASE("procgen.pointGraphStress.isBitStableAcrossIndependentBiomeInstances") {
    SpatialData domain = SpatialData::box(-32.f, 0.f, -32.f, 32.f, 0.f, 32.f);
    SpatialData center = SpatialData::sphere(0.f, 0.f, 0.f, 12.f);
    BiomeRules rules;
    CHECK(rules.addLayer("forest", &domain, 1, 0.85f));
    CHECK(rules.addLayer("clearing", &center, 5, 0.35f));
    CHECK(rules.addAsset("forest", "oak", 3.f, 0.8f, 1.3f, true));
    CHECK(rules.addAsset("forest", "pine", 1.f, 0.9f, 1.4f, true));
    CHECK(rules.addAsset("clearing", "rock", 1.f, 0.7f, 1.1f, true));

    PointGraph asset;
    CHECK(asset.addNode("biome", "biome.generate"));
    CHECK(asset.setNodeFloat("biome", "spacing", 1.f));
    CHECK(asset.setNodeFloat("biome", "jitter", 0.75f));
    CHECK(asset.setNodeInt("biome", "seed", 0x5a17));
    std::unique_ptr<PointGraph> first(asset.instantiate());
    std::unique_ptr<PointGraph> second(asset.instantiate());
    REQUIRE(bool(first));
    REQUIRE(bool(second));
    CHECK(first->setNodeSpatial("biome", &domain));
    CHECK(first->setNodeBiomeRules("biome", &rules));
    CHECK(second->setNodeSpatial("biome", &domain));
    CHECK(second->setNodeBiomeRules("biome", &rules));

    std::unique_ptr<PointSet> firstOutput(first->execute("biome"));
    std::unique_ptr<PointSet> secondOutput(second->execute("biome"));
    REQUIRE(bool(firstOutput));
    REQUIRE(bool(secondOutput));
    CHECK(firstOutput->getCount() > 3000);
    CHECK(firstOutput->getCount() < 4225);
    CHECK(pointSetsEqual(*firstOutput, *secondOutput));
}

TEST_CASE("procgen.pointGraphStress.recoversAfterExternalRuleAndDataErrors") {
    PointSet path;
    path.add(0.f, 0.f, 0.f);
    path.add(10.f, 0.f, 0.f);
    ShapeGrammar grammar;
    CHECK(grammar.addModule("A", "wall", 2.f));
    PointGraph graph;
    CHECK(graph.addNode("path", "input"));
    CHECK(graph.addNode("facade", "grammar.generate"));
    CHECK(graph.connect("path", "facade"));
    CHECK(graph.setNodeString("facade", "grammar", "Missing+"));
    CHECK(!graph.execute("facade"));
    CHECK(graph.getError().find("no points") != std::string::npos);

    CHECK(graph.setNodePoints("path", &path));
    CHECK(graph.setNodeShapeGrammar("facade", &grammar));
    CHECK(!graph.execute("facade"));
    CHECK(graph.getError().find("unknown module") != std::string::npos);
    CHECK(graph.setNodeString("facade", "grammar", "A+"));
    std::unique_ptr<PointSet> recovered(graph.execute("facade"));
    REQUIRE(bool(recovered));
    CHECK_EQ(recovered->getCount(), 5);
    CHECK_EQ(recovered->getStringAttribute(0, "asset", ""), std::string("wall"));
}

TEST_CASE("procgen.pointGraphStress.rejectsOversizedOutputsBeforeSpatialAllocation") {
    SpatialData domain = SpatialData::box(0.f, 0.f, 0.f, 100.f, 0.f, 100.f);
    PointGraph graph;
    CHECK(graph.addNode("sample", "spatial.sample"));
    CHECK(graph.setNodeSpatial("sample", &domain));
    CHECK(graph.setNodeFloat("sample", "spacing", 1.f));
    graph.setMaxNodeOutputPoints(1000);
    CHECK_EQ(graph.getMaxNodeOutputPoints(), 1000);
    CHECK(!graph.execute("sample"));
    CHECK(graph.getError().find("point budget") != std::string::npos);
    CHECK(!graph.getNodeOutput("sample"));

    graph.setMaxNodeOutputPoints(11000);
    std::unique_ptr<PointSet> output(graph.execute("sample"));
    REQUIRE(bool(output));
    CHECK_EQ(output->getCount(), 10201);

    graph.setMaxNodeOutputPoints(1000);
    CHECK(!graph.execute("sample"));
    CHECK(graph.getError().find("point budget") != std::string::npos);

    Heightmap heightmap(11, 11);
    heightmap.setHeight(10, 10, 100.f);
    SpatialData terrain = SpatialData::heightfield(heightmap, 0.f, 0.f, 1.f, 1.f);
    PointGraph terrainGraph;
    CHECK(terrainGraph.addNode("sample", "spatial.sample"));
    CHECK(terrainGraph.setNodeSpatial("sample", &terrain));
    terrainGraph.setMaxNodeOutputPoints(121);
    output.reset(terrainGraph.execute("sample"));
    REQUIRE(bool(output));
    CHECK_EQ(output->getCount(), 121);
}

TEST_CASE("procgen.pointGraphStress.preflightsInputMergeAndCopyAllocations") {
    PointSet first;
    PointSet second;
    for (int index = 0; index < 60; ++index) {
        first.add(float(index), 0.f, 0.f);
        second.add(0.f, 0.f, float(index));
    }
    PointGraph graph;
    CHECK(graph.addNode("first", "input"));
    CHECK(graph.addNode("second", "input"));
    CHECK(graph.addNode("merge", "merge"));
    CHECK(graph.addNode("copy", "copy.points"));
    graph.setMaxNodeOutputPoints(100);
    CHECK(graph.setNodePoints("first", &first));
    CHECK(graph.setNodePoints("second", &second));
    CHECK(graph.connect("first", "merge", 0));
    CHECK(graph.connect("second", "merge", 1));
    CHECK(graph.connect("first", "copy", 0));
    CHECK(graph.connect("second", "copy", 1));
    CHECK(graph.setNodeInt("copy", "maxPoints", 10000));
    CHECK(!graph.execute("merge"));
    CHECK(graph.getError().find("point budget") != std::string::npos);
    CHECK(!graph.execute("copy"));
    CHECK(graph.getError().find("point budget") != std::string::npos);

    PointSet oversized;
    for (int index = 0; index < 101; ++index) oversized.add(float(index), 0.f, 0.f);
    CHECK(!graph.setNodePoints("first", &oversized));
    CHECK(graph.getError().find("input exceeds") != std::string::npos);
}