/*
 *  <Short Description>
 *  Copyright (C) 2024  Brett Terpstra
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#ifndef GP_IMAGE_TEST_IMAGE_H
#define GP_IMAGE_TEST_IMAGE_H

#include <blt/std/types.h>
#include <array>
#include <optional>
#include "blt/std/assert.h"
#include "blt/std/utility.h"
#include "blt/std/allocator.h"
#include "blt/math/vectors.h"
#include <variant>
#include <config.h>

using image_data_t = std::array<blt::vec3, width * height>;
inline blt::bump_allocator<> img_allocator(8192);

class image
{
    private:
        mutable std::variant<image_data_t*, blt::vec3, float> data = nullptr;
    public:
        image()
        {
            data = img_allocator.allocate<image_data_t>();
            img_allocator.construct(std::get<image_data_t*>(data));
        }
        
        explicit image(const blt::vec3& v): data(v)
        {}
        
        explicit image(float f): data(f)
        {}
        
        image(const image& copy)
        {
            std::visit(blt::lambda_visitor{
                    [&](const image_data_t* d) {
                        data = img_allocator.allocate<image_data_t>();
                        img_allocator.construct(std::get<image_data_t*>(data));
                        
                        for (size_t i = 0; i < d->size(); i++)
                            std::get<image_data_t*>(data)[i] = d[i];
                    },
                    [&](const blt::vec3& v) {
                        this->data = v;
                    },
                    [&](float f) {
                        this->data = f;
                    }
            }, copy.data);
        }
        
        image(image&& move) noexcept
        {
            std::visit(blt::lambda_visitor{
                    [&](image_data_t* d) {
                        this->data = d;
                    },
                    [&](const blt::vec3& v) {
                        this->data = v;
                    },
                    [&](float f) {
                        this->data = f;
                    }
            }, move.data);
            
            move.data = nullptr;
        }
        
        image& operator=(const image& copy)
        {
            if (&copy == this)
                return *this;
            if (std::holds_alternative<image_data_t*>(data))
            {
                BLT_INFO("Copy deallocate");
                img_allocator.deallocate(std::get<image_data_t*>(data), 1);
            }
            std::visit(blt::lambda_visitor{
                    [&](const image_data_t* d) {
                        data = img_allocator.allocate<image_data_t>();
                        img_allocator.construct(std::get<image_data_t*>(data));
                        
                        for (size_t i = 0; i < d->size(); i++)
                            std::get<image_data_t*>(data)[i] = d[i];
                    },
                    [&](const blt::vec3& v) {
                        this->data = v;
                    },
                    [&](float f) {
                        this->data = f;
                    }
            }, copy.data);
            return *this;
        };
        
        image& operator=(image&& move) noexcept
        {
            std::swap(move.data, data);
            return *this;
        }
        
        image_data_t& getData()
        {
            if (std::holds_alternative<image_data_t*>(data))
                return *std::get<image_data_t*>(data);
            else
            {
                blt::vec3 color = std::visit(blt::lambda_visitor{
                        [](const image_data_t*) -> blt::vec3 {
                            BLT_ASSERT("There has been an error in the matrix. You should not have gotten here!");
                            return {};
                        },
                        [](const blt::vec3& v) -> blt::vec3 {
                            return v;
                        },
                        [](float f) -> blt::vec3 {
                            return blt::vec3{f, f, f};
                        }
                }, data);
                
                data = img_allocator.allocate<image_data_t>();
                img_allocator.construct(std::get<image_data_t*>(data));
                
                for (auto& v : *std::get<image_data_t*>(data))
                    v = color;
                return *std::get<image_data_t*>(data);
            }
        }
        
        [[nodiscard]] const image_data_t& getData() const
        {
            if (std::holds_alternative<image_data_t*>(data))
                return *std::get<image_data_t*>(data);
            else
            {
                blt::vec3 color = std::visit(blt::lambda_visitor{
                        [](const image_data_t*) -> blt::vec3 {
                            BLT_ASSERT("There has been an error in the matrix. You should not have gotten here!");
                            return {};
                        },
                        [](const blt::vec3& v) -> blt::vec3 {
                            return v;
                        },
                        [](float f) -> blt::vec3 {
                            return blt::vec3{f, f, f};
                        }
                }, data);
                
                data = img_allocator.allocate<image_data_t>();
                img_allocator.construct(std::get<image_data_t*>(data));
                
                for (auto& v : *std::get<image_data_t*>(data))
                    v = color;
                return *std::get<image_data_t*>(data);
            }
        }
        
        [[nodiscard]] blt::vec3 get(blt::i32 x = 0, blt::i32 y = 0) const
        {
            return std::visit(blt::lambda_visitor{
                    [x, y](const image_data_t* d) -> blt::vec3 {
                        if (x < 0 || y < 0 || x >= width || y >= height)
                            return blt::vec3{0, 0, 0};
                        return (*d)[y * height + x];
                    },
                    [](const blt::vec3& v) -> blt::vec3 {
                        return v;
                    },
                    [](float f) -> blt::vec3 {
                        return blt::vec3{f, f, f};
                    }
            }, data);
        }
        
        void set(const blt::vec3& c, blt::i32 x, blt::i32 y)
        {
            if (!std::holds_alternative<image_data_t*>(data))
            {
                data = img_allocator.allocate<image_data_t>();
                img_allocator.construct(std::get<image_data_t*>(data));
            }
            (*std::get<image_data_t*>(data))[y * height + x] = c;
        }
        
        ~image()
        {
            if (std::holds_alternative<image_data_t*>(data))
            {
                BLT_TRACE("%p", std::get<image_data_t*>(data));
                img_allocator.deallocate(std::get<image_data_t*>(data), 1);
            }
        }
};

#endif //GP_IMAGE_TEST_IMAGE_H