#include "iwa/resource/bitmap.hpp"

#include <cmath>
#include "iwa/log.hpp"
#include "iwa/util/vkutil.hpp"

namespace iwa
{
enum class InterpolationMethod
{
    NEAREST
};

class BitmapViewBase
{
protected:
    Bitmap* mBitmap;
protected:
    explicit inline BitmapViewBase(Bitmap* bitmap_) : mBitmap(bitmap_) {}
public:
    BitmapViewBase(const BitmapViewBase&) = delete;
    BitmapViewBase(BitmapViewBase&&) = default;

    BitmapViewBase& operator=(const BitmapViewBase&) = delete;
    BitmapViewBase& operator=(BitmapViewBase&&) = default;

    virtual ~BitmapViewBase() = default; // just to make clang happy, should be optimised out again
    [[nodiscard]] virtual glm::vec4 getPixel(unsigned x, unsigned y) const = 0;
    [[nodiscard]] virtual glm::vec4 getPixelRelative(float x, float y, InterpolationMethod interpolation = InterpolationMethod::NEAREST) const = 0;
    [[nodiscard]] virtual std::vector<glm::vec4> getPixels(unsigned x, unsigned y, unsigned width, unsigned height) const = 0;
    virtual void fill(const glm::vec4& color) = 0;
    virtual void copyChannels(const Bitmap& other, const std::vector<ChannelMapping>& mappings) = 0;
    virtual void multiply(const glm::vec4& color) = 0;
};

template<typename TData, int numComponents, bool normalized = std::is_floating_point_v<TData>>
class BitmapView : public BitmapViewBase
{
private:
    static constexpr std::size_t pixelSize = sizeof(TData) * numComponents;
public:
    explicit inline BitmapView(Bitmap* bitmap_) : BitmapViewBase(bitmap_) {}
    TData* getPixelRaw(unsigned x, unsigned y)
    {
        const std::size_t offset = x + y * mBitmap->getSize().width;
        return reinterpret_cast<TData*>(&mBitmap->getData()[offset * pixelSize]);
    }

    [[nodiscard]] const TData* getPixelRaw(unsigned x, unsigned y) const
    {
        const std::size_t offset = x + y * mBitmap->getSize().width;
        return reinterpret_cast<const TData*>(&mBitmap->getData()[offset * pixelSize]);
    }

    [[nodiscard]] TData convertChannel(float channel) const
    {
        if constexpr (normalized) {
            return static_cast<TData>(channel);
        } else {
            return static_cast<TData>(255 * channel);
        }
    }

    [[nodiscard]] float convertChannelBack(TData channel) const
    {
        if constexpr (normalized) {
            return static_cast<float>(channel);
        }
        else {
            return channel / 255.f;
        }
    }

    [[nodiscard]] TData clampChannel(TData value)
    {
        if constexpr (normalized) {
            return std::clamp(value, TData(0), TData(1));
        }
        else {
            return std::clamp(value, TData(0), TData(255));
        }
    }

    void convertColor(const glm::vec4& color, std::array<TData, numComponents>& outConverted) const
    {
        for (int comp = 0; comp < numComponents; ++comp)
        {
            outConverted[comp] = convertChannel(color[comp]);
        }
    }

    [[nodiscard]] glm::vec4 convertColorBack(const TData* raw) const
    {
        glm::vec4 result;
        for (int comp = 0; comp < numComponents; ++comp)
        {
            result[comp] = convertChannelBack(raw[comp]);
        }
        return result;
    }

    // overrides
    [[nodiscard]] glm::vec4 getPixel(unsigned x, unsigned y) const override
    {
        const TData* pixelRaw = getPixelRaw(x, y);
        return convertColorBack(pixelRaw);
    }

    [[nodiscard]] glm::vec4 getPixelRelative(float x, float y, InterpolationMethod interpolation = InterpolationMethod::NEAREST) const override
    {
        const vk::Extent2D size = mBitmap->getSize();
        (void) interpolation; // TODO
        const unsigned myX = std::clamp(static_cast<unsigned>(std::round(x * static_cast<float>(size.width))), 0u, size.width - 1);
        const unsigned myY = std::clamp(static_cast<unsigned>(std::round(y * static_cast<float>(size.height))), 0u, size.height - 1);

        return getPixel(myX, myY);
    }

    [[nodiscard]] std::vector<glm::vec4> getPixels(unsigned x, unsigned y, unsigned width, unsigned height) const override
    {
        std::vector<glm::vec4> pixels;
        pixels.resize(static_cast<std::size_t>(width) * height);

        for (unsigned ypos = y; ypos < y + height; ++ypos)
        {
            for (unsigned xpos = x; xpos < x + width; ++xpos)
            {
                pixels[xpos + height * ypos] = getPixel(xpos, ypos);
            }
        }
        return pixels;
    }

    void fill(const glm::vec4& color) override
    {
        std::array<TData, numComponents> convertedColor; // NOLINT(cppcoreguidelines-pro-type-member-init)
        convertColor(color, convertedColor);

        for (unsigned y = 0; y < mBitmap->getSize().height; ++y)
        {
            for (unsigned x = 0; x < mBitmap->getSize().width; ++x)
            {
                TData* pixel = getPixelRaw(x, y);
                for (int comp = 0; comp < numComponents; ++comp)
                {
                    pixel[comp] = convertedColor[comp];
                }
            }
        }
    };

    // TODO: maybe introduce a "dual-mView" class that includes both types at once?
    void copyChannels(const Bitmap& other, const std::vector<ChannelMapping>& mappings) override
    {
        const vk::Extent2D size = mBitmap->getSize();
        const std::vector<glm::vec4> otherPixels = other.getAllPixels();

        if(otherPixels.size() == static_cast<std::size_t>(size.width) * size.height)
        {
            for (unsigned y = 0; y < size.height; ++y)
            {
                for (unsigned x = 0; x < size.width; ++x)
                {
                    const glm::vec4 color = otherPixels[x + y * size.height];
                    for (const ChannelMapping& mapping : mappings)
                    {
                        assert(static_cast<int>(mapping.to) < numComponents);
                        TData* myPixel = getPixelRaw(x, y);
                        myPixel[static_cast<int>(mapping.to)] = convertChannel(color[static_cast<int>(mapping.from)]);
                    }
                }
            }
        }
        else
        {
            const vk::Extent2D otherSize = other.getSize();
            const float factorX = static_cast<float>(otherSize.width) / static_cast<float>(size.width);
            const float factorY = static_cast<float>(otherSize.height) / static_cast<float>(size.height);

            for (unsigned y = 0; y < size.height; ++y)
            {
                for (unsigned x = 0; x < size.width; ++x)
                {
                    const unsigned otherX = std::clamp(static_cast<unsigned>(std::round(factorX * static_cast<float>(x))), 0u, otherSize.width - 1);
                    const unsigned otherY = std::clamp(static_cast<unsigned>(std::round(factorY * static_cast<float>(x))), 0u, otherSize.height - 1);
                    const glm::vec4 color = otherPixels[otherX + otherY * size.height];
                    for (const ChannelMapping& mapping : mappings)
                    {
                        assert(static_cast<int>(mapping.to) < numComponents);
                        TData* myPixel = getPixelRaw(x, y);
                        myPixel[static_cast<int>(mapping.to)] = convertChannel(color[static_cast<int>(mapping.from)]);
                    }
                }
            }
        }
    }

    void multiply(const glm::vec4& color) override
    {
        for (unsigned y = 0; y < mBitmap->getSize().height; ++y)
        {
            for (unsigned x = 0; x < mBitmap->getSize().width; ++x)
            {
                TData* pixel = getPixelRaw(x, y);
                for (int comp = 0; comp < numComponents; ++comp)
                {
                    pixel[comp] = clampChannel(static_cast<TData>(pixel[comp] * color[comp]));
                }
            }
        }
    }
};

void Bitmap::createView()
{
    switch(mFormat)
    {
        case vk::Format::eR8Uint:
        case vk::Format::eR8Unorm:
        case vk::Format::eR8Srgb:
            mView = std::make_unique<BitmapView<std::uint8_t, 1>>(this);
            break;
        case vk::Format::eR8G8Uint:
        case vk::Format::eR8G8Unorm:
        case vk::Format::eR8G8Srgb:
            mView = std::make_unique<BitmapView<std::uint8_t, 2>>(this);
            break;
        case vk::Format::eR8G8B8Uint:
        case vk::Format::eR8G8B8Unorm:
        case vk::Format::eR8G8B8Srgb:
            mView = std::make_unique<BitmapView<std::uint8_t, 3>>(this);
            break;
        case vk::Format::eR8G8B8A8Uint:
        case vk::Format::eR8G8B8A8Unorm:
        case vk::Format::eR8G8B8A8Srgb:
            mView = std::make_unique<BitmapView<std::uint8_t, 4>>(this);
            break;
        case vk::Format::eR32Sfloat:
            mView = std::make_unique<BitmapView<float, 1>>(this);
            break;
        case vk::Format::eR32G32Sfloat:
            mView = std::make_unique<BitmapView<float, 2>>(this);
            break;
        case vk::Format::eR32G32B32Sfloat:
            mView = std::make_unique<BitmapView<float, 3>>(this);
            break;
        case vk::Format::eR32G32B32A32Sfloat:
            mView = std::make_unique<BitmapView<float, 4>>(this);
            break;
        default:
            logAndDie("Missing format for bitmap mView!");
    }
}

//
// public functions
//
Bitmap::Bitmap(BitmapCreationArgs args, ObjectPtr<BaseObject> owner)
        : super_t(std::move(owner)), mFormat(args.format), mSize(args.size)
{
    const std::size_t dataSize = static_cast<std::size_t>(args.size.width) * args.size.height * vkFormatSize(args.format);
    if (args.initialData.empty())
    {
        mData.resize(dataSize);
    }
    else
    {
        MIJIN_ASSERT(args.initialData->byteSize() == dataSize, "Bitmap initial data size is invalid.");
        mData = std::move(*args.initialData);
    }

    createView();
}

Bitmap::~Bitmap() // NOLINT(modernize-use-equals-default)
{
    // needs to be implemented for the mView to be deleted correctly
}

glm::vec4 Bitmap::getPixel(unsigned int x, unsigned int y) const
{
    return mView->getPixel(x, y);
}

std::vector<glm::vec4> Bitmap::getPixels(unsigned x, unsigned y, unsigned width, unsigned height) const
{
    return mView->getPixels(x, y, width, height);
}

void Bitmap::fill(const glm::vec4& color)
{
    mView->fill(color);
}

void Bitmap::copyChannels(const Bitmap& other, const std::vector<ChannelMapping>& mappings)
{
    mView->copyChannels(other, mappings);
}

void Bitmap::multiply(const glm::vec4& color)
{
    mView->multiply(color);
}
} // namespace iwa