// SPDX-License-Identifier: Apache-2.0
#include <vtbackend/Image.hpp>

#include <crispy/Point.hpp>
#include <crispy/StrongLRUHashtable.hpp>

#include <algorithm>
#include <memory>
#include <ranges>
#include <utility>

// clang-format off
#if (__cpp_lib_simd >= 202411L)
    #include <simd>
    namespace simd = std;
    #define VTBACKEND_SIMD_FOUND 1
#elif __has_include(<experimental/simd>) && !defined(__APPLE__) && !defined(__FreeBSD__)
    #include <experimental/simd>
    namespace simd = std::experimental;
    #define VTBACKEND_SIMD_FOUND 1
#endif
// clang-format on

using std::make_shared;
using std::ostream;
using std::shared_ptr;
using std::string;

namespace vtbackend
{

ImageStats& ImageStats::get()
{
    static ImageStats stats {};
    return stats;
}

Image::~Image()
{
    --ImageStats::get().instances;
    _onImageRemove(this);
}

RasterizedImage::~RasterizedImage()
{
    --ImageStats::get().rasterized;
}

ImageFragment::~ImageFragment()
{
    --ImageStats::get().fragments;
}

ImagePool::ImagePool(OnImageRemove onImageRemove, ImageId nextImageId):
    _nextImageId { nextImageId },
    _imageNameToImageCache { crispy::StrongHashtableSize { 1024 },
                             crispy::LRUCapacity { 100 },
                             "ImagePool name-to-image mappings" },
    _onImageRemove { std::move(onImageRemove) }
{
}

static constexpr ImageSize computeTargetSize(ImageResize resizePolicy,
                                             ImageSize imageSize,
                                             ImageSize gridSize) noexcept
{
    auto const imageWidth = unbox(imageSize.width);
    auto const imageHeight = unbox(imageSize.height);
    auto const gridWidth = unbox(gridSize.width);
    auto const gridHeight = unbox(gridSize.height);

    switch (resizePolicy)
    {
        case ImageResize::NoResize:
            return {
                .width = Width::cast_from(imageWidth),
                .height = Height::cast_from(imageHeight),
            };
        case ImageResize::ResizeToFit: {
            auto const scale = std::min(static_cast<double>(gridWidth) / static_cast<double>(imageWidth),
                                        static_cast<double>(gridHeight) / static_cast<double>(imageHeight));
            return {
                .width = Width::cast_from(static_cast<int>(imageWidth * scale)),
                .height = Height::cast_from(static_cast<int>(imageHeight * scale)),
            };
        }
        case ImageResize::ResizeToFill: {
            auto const scale = std::max(static_cast<double>(gridWidth) / static_cast<double>(imageWidth),
                                        static_cast<double>(gridHeight) / static_cast<double>(imageHeight));
            return {
                .width = Width::cast_from(static_cast<int>(imageWidth * scale)),
                .height = Height::cast_from(static_cast<int>(imageHeight * scale)),
            };
        }
        case ImageResize::StretchToFill: {
            return {
                .width = Width::cast_from(gridWidth),
                .height = Height::cast_from(gridHeight),
            };
        }
    }
    std::unreachable();
}

namespace
{
    struct TopLeft
    {
        int x {};
        int y {};
    };
} // namespace

static constexpr TopLeft computeTargetTopLeftOffset(ImageAlignment alignmentPolicy,
                                                    ImageSize targetSize,
                                                    ImageSize gridSize) noexcept
{
    auto const gridWidth = unbox<int>(gridSize.width);
    auto const gridHeight = unbox<int>(gridSize.height);
    auto const paramWidth = unbox<int>(targetSize.width);
    auto const paramHeight = unbox<int>(targetSize.height);
    switch (alignmentPolicy)
    {
        case ImageAlignment::TopStart: return { .x = 0, .y = 0 };
        case ImageAlignment::TopCenter: return { .x = (gridWidth - paramWidth) / 2, .y = 0 };
        case ImageAlignment::TopEnd: return { .x = gridWidth - paramWidth, .y = 0 };
        case ImageAlignment::MiddleStart: return { .x = 0, .y = (gridHeight - paramHeight) / 2 };
        case ImageAlignment::MiddleCenter:
            return { .x = (gridWidth - paramWidth) / 2, .y = (gridHeight - paramHeight) / 2 };
        case ImageAlignment::MiddleEnd:
            return { .x = gridWidth - paramWidth, .y = (gridHeight - paramHeight) / 2 };
        case ImageAlignment::BottomStart: return { .x = 0, .y = gridHeight - paramHeight };
        case ImageAlignment::BottomCenter:
            return { .x = (gridWidth - paramWidth) / 2, .y = gridHeight - paramHeight };
        case ImageAlignment::BottomEnd: return { .x = gridWidth - paramWidth, .y = gridHeight - paramHeight };
    }
    std::unreachable();
}

#ifdef VTBACKEND_SIMD_FOUND
namespace
{
    struct SimdContext
    {
        int width;
        int cellX;
        int xOffset;
        int yOffset;
        int paramWidth;
        int paramHeight;
        int imageWidth;
        int imageHeight;
        int subOffsetX;
        int subOffsetY;
        int subWidth;
        int subHeight;
        std::vector<uint8_t> const& imageData;
        uint32_t defaultColor;
    };

    void fillFragmentSimd(int& x, uint8_t*& target, SimdContext const& context, int globalY, bool yInBounds)
    {
        using FloatV = simd::native_simd<float>;
        using IntV = simd::rebind_simd_t<int, FloatV>;
        constexpr int SimdWidth = FloatV::size();

        for (; x + SimdWidth <= context.width; x += SimdWidth) // SIMD loop
        {
            auto const globalXVec = IntV([](int i) { return i; }) + (context.cellX + x);

            // Check bounds
            // X bounds depend on vector, Y bounds are scalar for this row
            auto const xInBounds =
                (globalXVec >= context.xOffset) && (globalXVec < (context.xOffset + context.paramWidth));

            if (yInBounds)
            {
                // Check if ALL X are in bounds
                if (simd::all_of(xInBounds))
                {
                    // Fully in bounds
                    auto const sourceXVec =
                        context.subOffsetX
                        + simd::static_simd_cast<int>(
                            (simd::static_simd_cast<float>(globalXVec) - static_cast<float>(context.xOffset))
                            * static_cast<float>(context.subWidth) / static_cast<float>(context.paramWidth));

                    auto const sourceY = context.subOffsetY
                                         + static_cast<int>((globalY - context.yOffset)
                                                            * static_cast<double>(context.subHeight)
                                                            / static_cast<double>(context.paramHeight));

                    for (int i = 0; i < SimdWidth; ++i)
                    {
                        auto const sourceIndex =
                            ((static_cast<size_t>(sourceY) * static_cast<size_t>(context.imageWidth))
                             + static_cast<size_t>(sourceXVec[i]))
                            * 4;
                        if (sourceIndex + 4 <= context.imageData.size())
                            *(uint32_t*) target = *(uint32_t*) &context.imageData[sourceIndex];
                        else
                            *(uint32_t*) target = context.defaultColor;
                        target += 4;
                    }
                }
                else if (simd::none_of(xInBounds))
                {
                    // None (X out of bounds)
                    for (int i = 0; i < SimdWidth; ++i)
                    {
                        *(uint32_t*) target = context.defaultColor;
                        target += 4;
                    }
                }
                else
                {
                    // Mixed X
                    for (int i = 0; i < SimdWidth; ++i)
                    {
                        if (xInBounds[i])
                        {
                            auto const globalX = context.cellX + x + i;
                            auto const sourceX =
                                context.subOffsetX
                                + static_cast<int>((globalX - context.xOffset)
                                                   * static_cast<double>(context.subWidth)
                                                   / static_cast<double>(context.paramWidth));
                            auto const sourceY =
                                context.subOffsetY
                                + static_cast<int>((globalY - context.yOffset)
                                                   * static_cast<double>(context.subHeight)
                                                   / static_cast<double>(context.paramHeight));
                            auto const sourceIndex =
                                ((static_cast<size_t>(sourceY) * static_cast<size_t>(context.imageWidth))
                                 + static_cast<size_t>(sourceX))
                                * 4;
                            if (sourceIndex + 4 <= context.imageData.size())
                                *(uint32_t*) target = *(uint32_t*) &context.imageData[sourceIndex];
                            else
                                *(uint32_t*) target = context.defaultColor;
                        }
                        else
                            *(uint32_t*) target = context.defaultColor;
                        target += 4;
                    }
                }
            }
            else
            {
                // Y out of bounds -> All out
                for (int i = 0; i < SimdWidth; ++i)
                {
                    *(uint32_t*) target = context.defaultColor;
                    target += 4;
                }
            }
        }
    }
} // namespace
#endif

FragmentPlacement RasterizedImage::fragmentPlacement(CellLocation pos, ImageSize targetCellSize) const
{
    // Deliberately mirrors fragment()'s setup: the two must agree on which source texel a target
    // pixel comes from, or swapping a CPU resample for a GPU one would move the image.
    auto const cellSize = targetCellSize.area() > 0 ? targetCellSize : _cellSize;
    auto const gridSize = ImageSize {
        .width = Width::cast_from(unbox<int>(_cellSpan.columns) * unbox<int>(cellSize.width)),
        .height = Height::cast_from(unbox<int>(_cellSpan.lines) * unbox<int>(cellSize.height)),
    };

    auto const subWidth =
        _imageSubSize.width.value > 0 ? unbox<int>(_imageSubSize.width) : unbox<int>(_image->width());
    auto const subHeight =
        _imageSubSize.height.value > 0 ? unbox<int>(_imageSubSize.height) : unbox<int>(_image->height());
    auto const subOffsetX = static_cast<int>(_imageOffset.x.value);
    auto const subOffsetY = static_cast<int>(_imageOffset.y.value);
    auto const imageWidth = unbox<int>(_image->width());
    auto const imageHeight = unbox<int>(_image->height());
    auto const effectiveImageSize = ImageSize { Width(subWidth), Height(subHeight) };
    auto const targetSize = computeTargetSize(_resizePolicy, effectiveImageSize, gridSize);
    auto const paramWidth = unbox<int>(targetSize.width);
    auto const paramHeight = unbox<int>(targetSize.height);
    auto const [xOffset, yOffset] = computeTargetTopLeftOffset(_alignmentPolicy, targetSize, gridSize);

    if (paramWidth <= 0 || paramHeight <= 0 || imageWidth <= 0 || imageHeight <= 0)
        return {};

    auto const cellX = unbox<int>(pos.column) * unbox<int>(cellSize.width);
    auto const cellY = unbox<int>(pos.line) * unbox<int>(cellSize.height);

    // Intersect the cell with the image's target region, in grid pixels.
    auto const left = std::max(cellX, xOffset);
    auto const top = std::max(cellY, yOffset);
    auto const right = std::min(cellX + unbox<int>(cellSize.width), xOffset + paramWidth);
    auto const bottom = std::min(cellY + unbox<int>(cellSize.height), yOffset + paramHeight);

    if (left >= right || top >= bottom)
        return {}; // wholly in the gap

    // The very map fragment() applies per pixel, evaluated at the rectangle's edges instead.
    auto const sourceEdge = [](int target, int offset, int extent, int subOffset, int subExtent) {
        return static_cast<double>(subOffset) + ((static_cast<double>(target - offset) * subExtent) / extent);
    };
    auto const sourceLeft = sourceEdge(left, xOffset, paramWidth, subOffsetX, subWidth);
    auto const sourceRight = sourceEdge(right, xOffset, paramWidth, subOffsetX, subWidth);
    auto const sourceTop = sourceEdge(top, yOffset, paramHeight, subOffsetY, subHeight);
    auto const sourceBottom = sourceEdge(bottom, yOffset, paramHeight, subOffsetY, subHeight);

    return FragmentPlacement {
        .hasImage = true,
        .targetX = left - cellX,
        .targetY = top - cellY,
        .targetSize = ImageSize { Width::cast_from(right - left), Height::cast_from(bottom - top) },
        .sourceX = static_cast<float>(sourceLeft / imageWidth),
        .sourceY = static_cast<float>(sourceTop / imageHeight),
        .sourceWidth = static_cast<float>((sourceRight - sourceLeft) / imageWidth),
        .sourceHeight = static_cast<float>((sourceBottom - sourceTop) / imageHeight),
        .coversCell = left == cellX && top == cellY && right == cellX + unbox<int>(cellSize.width)
                      && bottom == cellY + unbox<int>(cellSize.height),
    };
}

Image::Data RasterizedImage::fragment(CellLocation pos, ImageSize targetCellSize) const
{
    auto const cellSize = targetCellSize.area() > 0 ? targetCellSize : _cellSize;
    auto const gridSize = ImageSize {
        .width = Width::cast_from(unbox<int>(_cellSpan.columns) * unbox<int>(cellSize.width)),
        .height = Height::cast_from(unbox<int>(_cellSpan.lines) * unbox<int>(cellSize.height)),
    };

    // Use sub-region dimensions if specified, otherwise use full image.
    auto const subWidth =
        _imageSubSize.width.value > 0 ? unbox<int>(_imageSubSize.width) : unbox<int>(_image->width());
    auto const subHeight =
        _imageSubSize.height.value > 0 ? unbox<int>(_imageSubSize.height) : unbox<int>(_image->height());
    auto const subOffsetX = static_cast<int>(_imageOffset.x.value);
    auto const subOffsetY = static_cast<int>(_imageOffset.y.value);
    auto const imageWidth = unbox<int>(_image->width());
    auto const effectiveImageSize = ImageSize { Width(subWidth), Height(subHeight) };
    auto const targetSize = computeTargetSize(_resizePolicy, effectiveImageSize, gridSize);
    auto const paramWidth = unbox<int>(targetSize.width);
    auto const paramHeight = unbox<int>(targetSize.height);
    auto const [xOffset, yOffset] = computeTargetTopLeftOffset(_alignmentPolicy, targetSize, gridSize);

    // The pixel offset of the top-left corner of the current cell in the global grid system
    auto const cellX = unbox<int>(pos.column) * unbox<int>(cellSize.width);
    auto const cellY = unbox<int>(pos.line) * unbox<int>(cellSize.height);

    auto fragmentData = Image::Data {};
    fragmentData.resize(cellSize.area() * 4); // RGBA
    uint8_t* target = fragmentData.data();

#ifdef VTBACKEND_SIMD_FOUND
    auto const imageHeight = unbox<int>(_image->height());
    auto const simdContext = SimdContext {
        .width = unbox<int>(cellSize.width),
        .cellX = cellX,
        .xOffset = xOffset,
        .yOffset = yOffset,
        .paramWidth = paramWidth,
        .paramHeight = paramHeight,
        .imageWidth = imageWidth,
        .imageHeight = imageHeight,
        .subOffsetX = subOffsetX,
        .subOffsetY = subOffsetY,
        .subWidth = subWidth,
        .subHeight = subHeight,
        .imageData = _image->data(),
        .defaultColor = _defaultColor.value,
    };
#endif

    // Iterate over every pixel in the CELL
    for (int const y: std::views::iota(0, unbox<int>(cellSize.height)))
    {
        auto const globalY = cellY + y;
        bool const yInBounds = (globalY >= yOffset) && (globalY < (yOffset + paramHeight));
        int x = 0;

#ifdef VTBACKEND_SIMD_FOUND
        fillFragmentSimd(x, target, simdContext, globalY, yInBounds);
#endif

        // Scalar epilogue
        for (; x < unbox<int>(cellSize.width); ++x)
        {
            auto const globalX = cellX + x;
            if (globalX >= xOffset && globalX < xOffset + paramWidth && yInBounds)
            {
                auto const sourceX = subOffsetX
                                     + static_cast<int>((globalX - xOffset) * static_cast<double>(subWidth)
                                                        / static_cast<double>(paramWidth));
                auto const sourceY = subOffsetY
                                     + static_cast<int>((globalY - yOffset) * static_cast<double>(subHeight)
                                                        / static_cast<double>(paramHeight));
                auto const sourceIndex = ((static_cast<size_t>(sourceY) * static_cast<size_t>(imageWidth))
                                          + static_cast<size_t>(sourceX))
                                         * 4;
                if (sourceIndex + 4 <= _image->data().size())
                    *(uint32_t*) target = *(uint32_t*) &_image->data()[sourceIndex];
                else
                    *(uint32_t*) target = _defaultColor.value;
            }
            else
                *(uint32_t*) target = _defaultColor.value;
            target += 4;
        }
    }

    return fragmentData;
}

shared_ptr<Image const> ImagePool::create(ImageFormat format, ImageSize size, Image::Data&& data)
{
    // TODO: This operation should be idempotent, i.e. if that image has been created already, return a
    // reference to that.
    auto const id = _nextImageId++;
    // The remover prunes the id index before the user's callback runs. It captures the
    // index by shared_ptr, never the pool: an image's last reference may outlive the
    // pool member order-wise, and it may drop on any thread.
    auto image = make_shared<Image>(
        id, format, std::move(data), size, [index = _idIndex, remover = _onImageRemove](Image const* image) {
            {
                auto const _ = std::lock_guard { index->mutex };
                // Prune only if this dying image still owns its slot. After the uint32
                // id counter wraps, a live image may have re-taken this id via
                // insert_or_assign (see create() below); its weak_ptr is unexpired, so
                // erasing by id unconditionally would drop the LIVE entry. Our own slot's
                // weak_ptr is already expired once the destructor's remover runs, so
                // expired() precisely selects the entry to remove.
                if (auto const it = index->images.find(image->id().value);
                    it != index->images.end() && it->second.expired())
                    index->images.erase(it);
            }
            if (remover)
                remover(image);
        });
    {
        auto const _ = std::lock_guard { _idIndex->mutex };
        // insert_or_assign, not emplace: after the uint32 id counter wraps, the
        // reused id may still carry a stale index entry (an expired weak_ptr whose
        // remover has not yet pruned it). emplace would no-op on the existing key,
        // leaving the new live image unindexed — findImageById would then miss it
        // and callers would wrongly treat a present image as gone.
        _idIndex->images.insert_or_assign(id.value, image);
    }
    return image;
}

shared_ptr<Image const> ImagePool::findImageById(ImageId id) const
{
    auto const _ = std::lock_guard { _idIndex->mutex };
    if (auto const it = _idIndex->images.find(id.value); it != _idIndex->images.end())
        return it->second.lock();
    return {};
}

shared_ptr<RasterizedImage> rasterize(shared_ptr<Image const> image,
                                      ImageAlignment alignmentPolicy,
                                      ImageResize resizePolicy,
                                      RGBAColor defaultColor,
                                      GridSize cellSpan,
                                      ImageSize cellSize,
                                      ImageLayer layer)
{
    return make_shared<RasterizedImage>(
        std::move(image), alignmentPolicy, resizePolicy, defaultColor, cellSpan, cellSize, layer);
}

void ImagePool::link(string const& name, shared_ptr<Image const> imageRef)
{
    _imageNameToImageCache.emplace(name, std::move(imageRef));
}

shared_ptr<Image const> ImagePool::findImageByName(string const& name) const noexcept
{
    if (auto const* imageRef = _imageNameToImageCache.tryGet(name))
        return *imageRef;

    return {};
}

void ImagePool::unlink(string const& name)
{
    _imageNameToImageCache.remove(name);
}

void ImagePool::clear()
{
    _imageNameToImageCache.clear();
    {
        auto const _ = std::lock_guard { _idIndex->mutex };
        _idIndex->images.clear();
    }
}

void ImagePool::inspect(ostream& os) const
{
    os << "Image pool:\n";
    os << std::format("global image stats: {}\n", ImageStats::get());
    _imageNameToImageCache.inspect(os);
}

} // namespace vtbackend
