#include <AssetManager/AssetManager.hpp>
#include <FluidRegistry/FluidRegistry.hpp>

FluidRegistry fluidRegistry;

FluidRegistry::~FluidRegistry()
{
    for (auto& [name, fluid] : fluids)
    {
        delete fluid;
    }
}

void FluidRegistry::registerAllFluids()
{
    registerFluid<Fluid>(
        "main:water",
        "Water",
        "fluid/water",
        0.012f,
        BarrelTier::Wooden,
        CapsuleTier::Tin,
        Sound::waterIn,
        Sound::waterOut
    )
        .withMapColor(glm::vec4(0.0f, 0.5f, 1.0f, 1.0f))
        .withFlowRate(16);

    registerFluid<Fluid>(
        "main:oil",
        "Crude Oil",
        "fluid/oil",
        0.02f,
        BarrelTier::Steel,
        CapsuleTier::Tin,
        Sound::oilIn,
        Sound::oilOut
    )
        .withMapColor(glm::vec4(0.1f, 0.1f, 0.1f, 1.0f))
        .withFlowRate(4);

    registerFluid<Fluid>(
        "main:petrol",
        "Petrol",
        "fluid/petrol",
        0.0f,
        BarrelTier::Steel,
        CapsuleTier::Tin
    );

    registerFluid<Fluid>(
        "main:refu",
        "Refined Uranium",
        "fluid/refu",
        0.0f,
        BarrelTier::Lead,
        CapsuleTier::Lead
    );

    registerFluid<Fluid>(
        "main:plutonium",
        "Plutonium",
        "fluid/plutonium",
        0.0f,
        BarrelTier::Lead,
        CapsuleTier::Lead
    );
}

Fluid* FluidRegistry::getFluidByName(std::string name)
{
    if (!fluids.contains(name))
    {
        return nullptr;
    }

    return fluids[name];
}

void FluidRegistry::addLoadingTasks(std::vector<std::function<void (std::atomic<unsigned int> &)>> &tasks)
{
    for (auto& [_, theFluid] : fluids)
    {
        Fluid *fluid = theFluid;
        tasks.push_back([fluid](std::atomic<unsigned int> &tasksDone)
        {
            std::string iconPath = assets.resolve(formatString("blocktex/$+.png", fluid->textureName));
            Bitmap iconBitmap = Bitmap::fromFileOrDefault(iconPath);

            std::shared_ptr<Vk::Device> device = GraphicsContext::instance->getDevice();

            std::shared_ptr<Vk::CommandBuffer> cmdbuf = GraphicsContext::instance->createAsyncCommandBuffer();
            cmdbuf->reset();
            cmdbuf->begin();

            fluid->icon.widthValue = fluid->icon.heightValue = 8;
            
            fluid->icon.addComponent(device, *cmdbuf, iconBitmap.data());

            fluid->icon.buildMetadata(device, *cmdbuf);
            fluid->icon.buildDescriptorSet(device);

            GraphicsContext::onNextFrame([cmdbuf, fluid, &tasksDone](GraphicsContext &)
            {
                std::shared_ptr<Vk::Fence> fence = std::make_shared<Vk::Fence>(GraphicsContext::instance->getDevice());
                cmdbuf->end();
                cmdbuf->submit(fence.get());
                fence->wait();

                fluid->icon.finalize();
            
                tasksDone++;
            });
        });
    }
}

int FluidRegistry::getTotalFluids()
{
    return (int) fluids.size();
}

void FluidRegistry::forEachFluid(std::function<void(Fluid*)> callback)
{
    for (const auto &[_, fluid] : fluids)
    {
        callback(fluid);
    }
}