#include <iostream>
#include <blt/std/allocator.h>
#include <blt/gfx/window.h>
#include <blt/gfx/state.h>
#include <blt/gfx/stb/stb_image.h>
#include <blt/gfx/stb/stb_perlin.h>
#include <blt/std/hashmap.h>
#include <image.h>
#include "blt/gfx/renderer/resource_manager.h"
#include "blt/gfx/renderer/batch_2d_renderer.h"
#include "blt/std/assert.h"
#include <variant>
#include <random>

blt::gfx::matrix_state_manager global_matrices;
blt::gfx::resource_manager resources;
blt::gfx::batch_renderer_2d renderer_2d(resources);

struct node;

blt::area_allocator<node, 32000> node_allocator;
blt::area_allocator<image, 32000> img_allocator;

using data_t = std::array<float, 5>;

node* createNode()
{
    auto* n = node_allocator.allocate(1);
    node_allocator.construct(n);
    return n;
}

image* createImage()
{
    auto* i = img_allocator.allocate(1);
    img_allocator.construct(i);
    return i;
}

void destroyNode(node* n)
{
    node_allocator.destroy(n);
    node_allocator.deallocate(n, 1);
}

void destroyImage(image* img)
{
    img_allocator.destroy(img);
    img_allocator.deallocate(img, 1);
}

static constexpr int OPERATOR_COUNT = 13;

class tree;

struct node
{
        friend tree;
    private:
        node* left;
        node* right;
        function_t type;
        data_t data;
        image* img = nullptr;
        
        [[nodiscard]] bool has_both() const
        {
            return left->img != nullptr && right->img != nullptr;
        }
        
        [[nodiscard]] bool has_one() const
        {
            return left->img != nullptr || right->img != nullptr;
        }
        
        [[nodiscard]] image* get_one() const
        {
            if (left->img != nullptr)
                return left->img;
            if (right->img != nullptr)
                return right->img;
            BLT_THROW(std::runtime_error("Unable to get when one is missing! You have an error in your GA!"));
        }
    
    public:
        
        static node* construct_random()
        {
            node* n = createNode();
            
            std::random_device dev;
            std::mt19937_64 engine{dev()};
            std::uniform_real_distribution<float> dist(0.0f, 1.0f);
            for (float& f : n->data)
                f = dist(engine);
            
            std::uniform_int_distribution op(0, OPERATOR_COUNT - 1);
            n->type = static_cast<function_t>(op(engine));
            
            return n;
        }
        
        void evaluate()
        {
            img = createImage();
            for (i32 y = 0; y < height; y++)
            {
                for (i32 x = 0; x < width; x++)
                {
                    switch (type)
                    {
                        case function_t::ADD:
                        {
                            BLT_ASSERT(has_both());
                            img->set((left->img->get(x, y) + right->img->get(x, y)) % std::numeric_limits<u8>::max(), x, y);
                            break;
                        }
                        case function_t::SUB:
                        {
                            BLT_ASSERT(has_both());
                            img->set((left->img->get(x, y) - right->img->get(x, y)) % std::numeric_limits<u8>::max(), x, y);
                            break;
                        }
                        case function_t::MUL:
                        {
                            BLT_ASSERT(has_both());
                            img->set((left->img->get(x, y) * right->img->get(x, y)) % std::numeric_limits<u8>::max(), x, y);
                            break;
                        }
                        case function_t::DIV:
                        {
                            BLT_ASSERT(has_both());
                            auto den = right->img->get(x, y);
                            if (den == 0)
                                img->set(0, x, y);
                            else
                                img->set((left->img->get(x, y) / den) % std::numeric_limits<u8>::max(), x, y);
                            break;
                        }
                        case function_t::EXP:
                        {
                            BLT_ASSERT(has_both());
                            img->set(static_cast<u8>(std::pow(left->img->get(x, y), right->img->get(x, y))), x, y);
                            break;
                        }
                        case function_t::LOG:
                        {
                            BLT_ASSERT(has_one());
                            const auto& v = get_one()->get(x, y);
                            auto v_x = v.x() == 0 ? 0.000001f : v.x();
                            auto v_y = v.y() == 0 ? 0.000001f : v.y();
                            auto v_z = v.z() == 0 ? 0.000001f : v.z();
                            img->set({std::log(v_x), std::log(v_y), std::log(v_z)}, x, y);
                            break;
                        }
                        case function_t::SQRT:
                        {
                            BLT_ASSERT(has_one());
                            const auto& v = get_one()->get(x, y);
                            img->set({std::sqrt(v.x()), std::sqrt(v.y()), std::sqrt(v.z())}, x, y);
                            break;
                        }
                        case function_t::QUAD:
                        {
                            BLT_ASSERT(has_one());
                            auto v = get_one()->get(x, y);
                            img->set(v * v, x, y);
                            break;
                        }
                        case function_t::RANDOM:
                        {
                            static std::random_device dev;
                            static std::mt19937_64 engine(dev());
                            static std::uniform_real_distribution<float> dist(0, 1);
                            img->set({dist(engine), dist(engine), dist(engine), 1.0f}, x, y);
                            break;
                        }
                        case function_t::NOISE:
                        {
                            auto v = std::numeric_limits<u8>::max() *
                                     stb_perlin_noise3(static_cast<float>(x) / ((256.0f * (data[0] + 1.0f) / 2.0f)),
                                                       static_cast<float>(y) / ((256 * (data[1] + 1) / 2)), data[2] / 2.0f + 0.5f, 0,
                                                       0, 0);
                            img->set(blt::vec4{v,v,v,1.0f}, x, y);
                            break;
                        }
                        case function_t::COLOR:
                        {
                            img->set(blt::vec4{data[0], data[1], data[2], 1.0f}, x, y);
                            break;
                        }
                        case function_t::SCALAR:
                        {
                            img->set(blt::vec4{data[0], data[0], data[0], 1.0f}, x, y);
                            break;
                        }
                    }
                }
            }
        }
        
        ~node()
        {
            destroyNode(left);
            destroyNode(right);
        }
};

class tree
{
    public:
        node* root;
    
    
};

void init()
{
    global_matrices.create_internals();
    resources.load_resources();
    renderer_2d.create();
}

void update(std::int32_t, std::int32_t)
{
    global_matrices.update();
    renderer_2d.render();
}

int main()
{
    blt::gfx::init(blt::gfx::window_data{"Window of GP test", init, update}.setSyncInterval(1));
    global_matrices.cleanup();
    resources.cleanup();
    renderer_2d.cleanup();
    blt::gfx::cleanup();
    return 0;
}
