#include <nd-wfc/wfc.hpp>
#include <nd-wfc/wfc_builder.hpp>
#include <nd-wfc/worlds.hpp>
#include <iostream>

// Tile types for the dungeon
enum class Tile : int {
    Empty = 0,
    Wall = 1,
    Floor = 2
};

constexpr size_t DungeonWidth = 16;
constexpr size_t DungeonHeight = 16;

using World = WFC::Array2D<Tile, DungeonWidth, DungeonHeight>;

void printDungeon(const World& world) {
    for (size_t y = 0; y < DungeonHeight; ++y) {
        for (size_t x = 0; x < DungeonWidth; ++x) {
            switch (world.at(x, y)) {
                case Tile::Floor:    std::cout << '.'; break;
                case Tile::Wall:     std::cout << '#'; break;
                case Tile::Empty:    std::cout << ' '; break;
            }
        }
        std::cout << '\n';
    }
}

int main() {
    std::cout << "2D Dungeon WFC Demo\n";
    std::cout << "Dungeon size: " << DungeonWidth << "x" << DungeonHeight << "\n\n";

    using DungeonBuilder = WFC::Builder<World>
        ::DefineIDs<Tile::Floor>
        ::Constrain<decltype([](const World& world, size_t index, WFC::WorldValue<Tile> val, auto& constrainer) constexpr {

            auto [x, y] = world.getCoord(index);

            // floor cannot be adjacent to empty space
            constrainer.Exclude(Tile::Empty, world.getCoordOffset(x, y, -1, 0)); // Left
            constrainer.Exclude(Tile::Empty, world.getCoordOffset(x, y, 1, 0));  // Right
            constrainer.Exclude(Tile::Empty, world.getCoordOffset(x, y, 0, -1)); // Up
            constrainer.Exclude(Tile::Empty, world.getCoordOffset(x, y, 0, 1));  // Down
        })>
        ::DefineIDs<Tile::Wall, Tile::Empty>
        ::SetInitialState<decltype([](World& world, auto& constrainer, auto& rng) constexpr {
            // make it impossible for the edge to be floor
            for (size_t x = 0; x < world.width(); ++x) {
                constrainer.Exclude(world.getId({x, 0}), Tile::Floor);
                constrainer.Exclude(world.getId({x, world.height() - 1}), Tile::Floor);
            }
            for (size_t y = 0; y < world.height(); ++y) {
                constrainer.Exclude(world.getId({0, y}), Tile::Wall);
                constrainer.Exclude(world.getId({world.width() - 1, y}), Tile::Wall);
            }
        })>
        ::Build;

    World world{};
    DungeonBuilder::Solve(world, 0xDEADBEEF);
    printDungeon(world);

    return 0;
}