/*
 * Copyright (c) 2024-2025, Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <LibGfx/Bitmap.h>
#include <LibGfx/PaintingSurface.h>
#include <LibGfx/SharedImageBuffer.h>
#include <LibGfx/SkiaUtils.h>

#include <core/SkCanvas.h>
#include <core/SkColorSpace.h>
#include <core/SkImage.h>
#include <core/SkPaint.h>
#include <core/SkRect.h>
#include <core/SkSurface.h>
#include <gpu/ganesh/GrBackendSurface.h>
#include <gpu/ganesh/GrDirectContext.h>
#include <gpu/ganesh/SkSurfaceGanesh.h>

#ifdef AK_OS_MACOS
#    include <gpu/ganesh/mtl/GrMtlBackendSurface.h>
#elif defined(USE_VULKAN_DMABUF_IMAGES)
#    include <LibGfx/VulkanImage.h>
#    include <gpu/ganesh/vk/GrVkBackendSurface.h>
#    include <gpu/ganesh/vk/GrVkTypes.h>
#endif

namespace Gfx {

struct PaintingSurface::Impl {
    RefPtr<SkiaBackendContext> context;
    IntSize size;
    sk_sp<SkSurface> surface;
    RefPtr<Bitmap> bitmap;
};

#if defined(AK_OS_MACOS) || defined(USE_VULKAN_DMABUF_IMAGES)
static GrSurfaceOrigin origin_to_sk_origin(PaintingSurface::Origin origin)
{
    switch (origin) {
    case PaintingSurface::Origin::BottomLeft:
        return kBottomLeft_GrSurfaceOrigin;
    default:
        VERIFY_NOT_REACHED();
    case PaintingSurface::Origin::TopLeft:
        return kTopLeft_GrSurfaceOrigin;
    }
}
#endif

#ifdef USE_VULKAN_DMABUF_IMAGES
static SkColorType vk_format_to_sk_color_type(VkFormat format)
{
    switch (format) {
    case VK_FORMAT_B8G8R8A8_UNORM:
        return kBGRA_8888_SkColorType;
    // add more as needed
    default:
        VERIFY_NOT_REACHED();
        return kUnknown_SkColorType;
    }
}

static void release_vulkan_image(void* context)
{
    VulkanImage* image = static_cast<VulkanImage*>(context);
    image->unref();
}

NonnullRefPtr<PaintingSurface> PaintingSurface::create_from_vkimage(NonnullRefPtr<SkiaBackendContext> context, NonnullRefPtr<VulkanImage> vulkan_image, Origin origin)
{
    IntSize size(vulkan_image->info.extent.width, vulkan_image->info.extent.height);
    GrVkImageInfo info = {
        .fImage = vulkan_image->image,
        .fAlloc = {}, // we're managing the memory ourselves
        .fImageTiling = vulkan_image->info.tiling,
        .fImageLayout = vulkan_image->info.layout,
        .fFormat = vulkan_image->info.format,
        .fImageUsageFlags = vulkan_image->info.usage,
        .fSampleCount = 1,
        .fLevelCount = 1,
        .fCurrentQueueFamily = VK_QUEUE_FAMILY_IGNORED,
        .fProtected = skgpu::Protected::kNo,
        .fYcbcrConversionInfo = {},
        .fSharingMode = vulkan_image->info.sharing_mode,
    };
    GrBackendRenderTarget rt = GrBackendRenderTargets::MakeVk(size.width(), size.height(), info);
    // Note, we're implicitly giving Skia a reference to vulkan_image. It will eventually be released by the callback function.
    vulkan_image->ref();
    sk_sp<SkSurface> surface = SkSurfaces::WrapBackendRenderTarget(context->sk_context(), rt, origin_to_sk_origin(origin), vk_format_to_sk_color_type(vulkan_image->info.format),
        SkColorSpace::MakeSRGB(), nullptr, release_vulkan_image, vulkan_image.ptr());
    return adopt_ref(*new PaintingSurface(make<Impl>(context, size, surface, nullptr)));
}
#endif

NonnullRefPtr<PaintingSurface> PaintingSurface::create_with_size(IntSize size, BitmapFormat color_type, AlphaType alpha_type, RefPtr<SkiaBackendContext> context)
{
    auto sk_color_type = to_skia_color_type(color_type);
    auto sk_alpha_type = to_skia_alpha_type(color_type, alpha_type);
    auto image_info = SkImageInfo::Make(size.width(), size.height(), sk_color_type, sk_alpha_type, SkColorSpace::MakeSRGB());

    if (context) {
        auto surface = SkSurfaces::RenderTarget(context->sk_context(), skgpu::Budgeted::kNo, image_info);
        if (surface)
            return adopt_ref(*new PaintingSurface(make<Impl>(context, size, surface, nullptr)));
        dbgln("Unable to create GPU surface for size {}x{}, falling back to CPU", size.width(), size.height());
        context = nullptr;
    }

    auto bitmap = Bitmap::create(color_type, alpha_type, size).value();
    auto surface = SkSurfaces::WrapPixels(image_info, bitmap->begin(), bitmap->pitch());
    VERIFY(surface);
    return adopt_ref(*new PaintingSurface(make<Impl>(context, size, surface, bitmap)));
}

NonnullRefPtr<PaintingSurface> PaintingSurface::wrap_bitmap(Bitmap& bitmap)
{
    auto color_type = to_skia_color_type(bitmap.format());
    auto alpha_type = to_skia_alpha_type(bitmap.format(), bitmap.alpha_type());
    auto size = bitmap.size();
    auto image_info = SkImageInfo::Make(bitmap.width(), bitmap.height(), color_type, alpha_type, SkColorSpace::MakeSRGB());
    auto surface = SkSurfaces::WrapPixels(image_info, bitmap.begin(), bitmap.pitch());
    return adopt_ref(*new PaintingSurface(make<Impl>(RefPtr<SkiaBackendContext> {}, size, surface, bitmap)));
}

#ifdef AK_OS_MACOS
NonnullRefPtr<PaintingSurface> PaintingSurface::create_from_shared_image_buffer(SharedImageBuffer& shared_image_buffer, NonnullRefPtr<SkiaBackendContext> context, Origin origin)
{
    auto const& iosurface_handle = shared_image_buffer.iosurface_handle();
    auto metal_texture = context->metal_context().create_texture_from_iosurface(iosurface_handle);
    IntSize const size { metal_texture->width(), metal_texture->height() };
    auto image_info = SkImageInfo::Make(size.width(), size.height(), kBGRA_8888_SkColorType, kPremul_SkAlphaType, SkColorSpace::MakeSRGB());
    GrMtlTextureInfo mtl_info;
    mtl_info.fTexture = sk_ret_cfp(metal_texture->texture());
    auto backend_render_target = GrBackendRenderTargets::MakeMtl(metal_texture->width(), metal_texture->height(), mtl_info);
    auto surface = SkSurfaces::WrapBackendRenderTarget(context->sk_context(), backend_render_target, origin_to_sk_origin(origin), kBGRA_8888_SkColorType, SkColorSpace::MakeSRGB(), nullptr);
    return adopt_ref(*new PaintingSurface(make<Impl>(context, size, surface, nullptr)));
}
#endif

PaintingSurface::PaintingSurface(NonnullOwnPtr<Impl>&& impl)
    : m_impl(move(impl))
{
}

PaintingSurface::~PaintingSurface()
{
    m_impl->surface = nullptr;
}

NonnullRefPtr<Bitmap> PaintingSurface::snapshot_bitmap() const
{
    auto bitmap = MUST(Bitmap::create(BitmapFormat::BGRA8888, AlphaType::Premultiplied, size()));
    read_into_bitmap(*bitmap);
    return bitmap;
}

SharedImage PaintingSurface::snapshot_into_shared_image() const
{
    auto shared_image_buffer = SharedImageBuffer::create(size());
    read_into_bitmap(*shared_image_buffer.bitmap());
    return shared_image_buffer.export_shared_image();
}

void PaintingSurface::read_into_bitmap(Bitmap& bitmap, IntPoint source_position) const
{
    auto color_type = to_skia_color_type(bitmap.format());
    auto alpha_type = to_skia_alpha_type(bitmap.format(), bitmap.alpha_type());
    auto image_info = SkImageInfo::Make(bitmap.width(), bitmap.height(), color_type, alpha_type, SkColorSpace::MakeSRGB());
    SkPixmap const pixmap(image_info, bitmap.begin(), bitmap.pitch());
    m_impl->surface->readPixels(pixmap, source_position.x(), source_position.y());
}

void PaintingSurface::write_from_bitmap(Bitmap const& bitmap)
{
    auto color_type = to_skia_color_type(bitmap.format());
    auto alpha_type = to_skia_alpha_type(bitmap.format(), bitmap.alpha_type());
    auto image_info = SkImageInfo::Make(bitmap.width(), bitmap.height(), color_type, alpha_type, SkColorSpace::MakeSRGB());
    SkPixmap const pixmap(image_info, bitmap.begin(), bitmap.pitch());
    m_impl->surface->writePixels(pixmap, 0, 0);
}

void PaintingSurface::copy_from_surface(PaintingSurface& source)
{
    source.flush();

    auto image = source.m_impl->surface->makeImageSnapshot();
    if (!image)
        return;

    SkPaint paint;
    paint.setBlendMode(SkBlendMode::kSrc);
    canvas().drawImageRect(
        image.get(),
        SkRect::MakeIWH(image->width(), image->height()),
        SkRect::MakeIWH(size().width(), size().height()),
        SkSamplingOptions {},
        &paint,
        SkCanvas::kStrict_SrcRectConstraint);
}

IntSize PaintingSurface::size() const
{
    return m_impl->size;
}

IntRect PaintingSurface::rect() const
{
    return { {}, m_impl->size };
}

SkCanvas& PaintingSurface::canvas() const
{
    return *m_impl->surface->getCanvas();
}

SkSurface& PaintingSurface::sk_surface() const
{
    return *m_impl->surface;
}

void PaintingSurface::notify_content_will_change()
{
    m_impl->surface->notifyContentWillChange(SkSurface::kDiscard_ContentChangeMode);
}

template<>
sk_sp<SkImage> PaintingSurface::sk_image_snapshot() const
{
    return m_impl->surface->makeImageSnapshot();
}

RefPtr<SkiaBackendContext> PaintingSurface::skia_backend_context() const
{
    return m_impl->context;
}

void PaintingSurface::flush()
{
    if (on_flush)
        on_flush(*this);
}

}