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

#include <LibGfx/Filter.h>
#include <LibGfx/Font/Font.h>
#include <LibGfx/SharedImageBuffer.h>
#include <LibMedia/VideoFrame.h>
#include <LibWeb/Painting/DisplayList.h>
#include <LibWeb/Painting/DisplayListResourceStorage.h>

namespace Web::Painting {

DisplayListResourceStorage::~DisplayListResourceStorage() = default;

FontResourceId DisplayListResourceStorage::add_font(Gfx::Font const& font)
{
    auto id = font.id();
    m_fonts.ensure(id, [&]() -> NonnullRefPtr<Gfx::Font const> { return font; });
    return { id };
}

ImageFrameResourceId DisplayListResourceStorage::add_image_frame(Gfx::DecodedImageFrame const& frame)
{
    auto id = frame.id();
    m_image_frames.ensure(id, [&] { return frame; });
    return { id };
}

VideoFrameResourceId DisplayListResourceStorage::add_video_frame(VideoFrameResourceId id, RefPtr<Media::VideoFrame const> frame)
{
    m_video_frames.set(id.value(), move(frame), AK::HashSetExistingEntryBehavior::Keep);
    return id;
}

DisplayListResourceId DisplayListResourceStorage::add_display_list(NonnullRefPtr<DisplayList const> display_list)
{
    auto id = display_list->id();
    m_display_lists.ensure(id, [&] { return move(display_list); });
    return { id };
}

void DisplayListResourceStorage::set_font(FontResourceId id, NonnullRefPtr<Gfx::Font const> font)
{
    m_fonts.set(id.value(), move(font));
}

void DisplayListResourceStorage::set_image_frame(ImageFrameResourceId id, Gfx::DecodedImageFrame frame)
{
    m_image_frames.set(id.value(), move(frame));
}

void DisplayListResourceStorage::set_video_frame(VideoFrameResourceId id, RefPtr<Media::VideoFrame const> frame)
{
    m_video_frames.set(id.value(), move(frame));
}

void DisplayListResourceStorage::set_display_list(DisplayListResourceId id, NonnullRefPtr<DisplayList const> display_list)
{
    m_display_lists.set(id.value(), move(display_list));
}

static ReadonlyBytes inline_data(ReadonlyBytes payload, DisplayListDataSpan span)
{
    VERIFY(static_cast<size_t>(span.offset) + span.size <= payload.size());
    return payload.slice(span.offset, span.size);
}

void DisplayListResourceStorage::append_referenced_resources_from(
    DisplayListResourceStorage const& source,
    ReadonlyBytes command_bytes)
{
    auto referenced_resources = source.collect_referenced_resources(command_bytes);
    for (auto id : referenced_resources.fonts)
        add_font(source.font(id));
    for (auto id : referenced_resources.image_frames)
        add_image_frame(source.image_frame(id));
    for (auto id : referenced_resources.video_frames)
        add_video_frame(id, source.video_frame(id));
    for (auto id : referenced_resources.display_lists)
        add_display_list(source.display_list(id));
}

void DisplayListResourceStorage::collect_referenced_resources(
    ReadonlyBytes command_bytes,
    DisplayListResourceSet& referenced_resources) const
{
    auto add_display_list_resource = [&](DisplayListResourceId id, Optional<ReadonlyBytes> command_bytes_to_collect) {
        if (referenced_resources.display_lists.set(id, AK::HashSetExistingEntryBehavior::Keep) != HashSetResult::InsertedNewEntry)
            return;
        auto const& nested_display_list = display_list(id);
        collect_referenced_resources(command_bytes_to_collect.value_or(nested_display_list.command_bytes()), referenced_resources);
    };

    DisplayList::for_each_command_header(command_bytes, [&](DisplayListCommandHeader const& header, ReadonlyBytes payload) {
        visit_display_list_command(header.type, payload, [&](auto const& command) {
            if constexpr (requires { command.font_id; })
                referenced_resources.fonts.set(command.font_id, AK::HashSetExistingEntryBehavior::Keep);
            if constexpr (requires { command.frame_id; })
                referenced_resources.image_frames.set(command.frame_id, AK::HashSetExistingEntryBehavior::Keep);
            if constexpr (requires { command.video_frame_id; })
                referenced_resources.video_frames.set(command.video_frame_id, AK::HashSetExistingEntryBehavior::Keep);
            if constexpr (requires { command.paint_style; command.paint_kind; }) {
                if (command.paint_kind == decltype(command.paint_kind)::PaintStyle
                    && command.paint_style.type == DisplayListPaintStyleType::Pattern)
                    add_display_list_resource(command.paint_style.pattern_tile_display_list_id, {});
            }
            if constexpr (requires { command.backdrop_filter_data; }) {
                if (command.has_backdrop_filter) {
                    Gfx::deserialize_filter(inline_data(payload, command.backdrop_filter_data), [&](u64 image_id) {
                        referenced_resources.image_frames.set(ImageFrameResourceId { image_id }, AK::HashSetExistingEntryBehavior::Keep);
                        return image_frame(ImageFrameResourceId { image_id });
                    });
                }
            }
            if constexpr (requires { command.filter_data; }) {
                if (command.has_filter) {
                    Gfx::deserialize_filter(inline_data(payload, command.filter_data), [&](u64 image_id) {
                        referenced_resources.image_frames.set(ImageFrameResourceId { image_id }, AK::HashSetExistingEntryBehavior::Keep);
                        return image_frame(ImageFrameResourceId { image_id });
                    });
                }
            }
            if constexpr (requires { command.display_list_id; }) {
                if constexpr (requires { command.command_bytes; })
                    add_display_list_resource(command.display_list_id, inline_data(payload, command.command_bytes));
                else
                    add_display_list_resource(command.display_list_id, {});
            }
        });
    });
}

DisplayListResourceSet DisplayListResourceStorage::collect_referenced_resources(ReadonlyBytes command_bytes) const
{
    DisplayListResourceSet referenced_resources;
    collect_referenced_resources(command_bytes, referenced_resources);
    return referenced_resources;
}

DisplayListResourceSet DisplayListResourceStorage::collect_referenced_resources(DisplayList const& display_list) const
{
    return collect_referenced_resources(display_list.command_bytes());
}

DisplayListResourceTransaction DisplayListResourceStorage::create_transaction(
    DisplayListResourceSet const& previous,
    DisplayListResourceSet const& current) const
{
    DisplayListResourceTransaction transaction;

    for (auto id : current.fonts) {
        if (!previous.fonts.contains(id))
            transaction.fonts.append({ id, font(id) });
    }
    for (auto id : current.image_frames) {
        if (!previous.image_frames.contains(id))
            transaction.image_frames.append({ id, image_frame(id) });
    }
    for (auto id : current.video_frames) {
        if (!previous.video_frames.contains(id))
            transaction.video_frames.append({ id, video_frame(id) });
    }
    for (auto id : current.display_lists) {
        if (!previous.display_lists.contains(id))
            transaction.display_lists.append(display_list(id));
    }

    for (auto id : previous.fonts) {
        if (!current.fonts.contains(id))
            transaction.font_ids_to_remove.append(id);
    }
    for (auto id : previous.image_frames) {
        if (!current.image_frames.contains(id))
            transaction.image_frame_ids_to_remove.append(id);
    }
    for (auto id : previous.video_frames) {
        if (!current.video_frames.contains(id))
            transaction.video_frame_ids_to_remove.append(id);
    }
    for (auto id : previous.display_lists) {
        if (!current.display_lists.contains(id))
            transaction.display_list_ids_to_remove.append(id);
    }
    return transaction;
}

void DisplayListResourceStorage::apply_transaction(DisplayListResourceTransaction&& transaction)
{
    for (auto& font : transaction.fonts)
        set_font(font.id, move(font.font));
    for (auto& frame : transaction.image_frames)
        set_image_frame(frame.id, move(frame.frame));
    for (auto& video_frame : transaction.video_frames)
        add_video_frame(video_frame.id, move(video_frame.frame));
    for (auto& display_list : transaction.display_lists)
        add_display_list(move(display_list));

    for (auto id : transaction.font_ids_to_remove)
        m_fonts.remove(id.value());
    for (auto id : transaction.image_frame_ids_to_remove)
        m_image_frames.remove(id.value());
    for (auto id : transaction.video_frame_ids_to_remove)
        m_video_frames.remove(id.value());
    for (auto id : transaction.display_list_ids_to_remove)
        m_display_lists.remove(id.value());
}

void DisplayListResourceStorage::retain_only(DisplayListResourceSet const& resource_set)
{
    m_fonts.remove_all_matching([&](auto id, auto const&) { return !resource_set.fonts.contains(FontResourceId { id }); });
    m_image_frames.remove_all_matching([&](auto id, auto const&) { return !resource_set.image_frames.contains(ImageFrameResourceId { id }); });
    m_video_frames.remove_all_matching([&](auto id, auto const&) { return !resource_set.video_frames.contains(VideoFrameResourceId { id }); });
    m_display_lists.remove_all_matching([&](auto id, auto const&) { return !resource_set.display_lists.contains(DisplayListResourceId { id }); });
}

void DisplayListResourceStorage::update_video_frame(VideoFrameResourceId frame_id, NonnullRefPtr<Media::VideoFrame const> frame)
{
    m_video_frames.set(frame_id.value(), move(frame));
}

void DisplayListResourceStorage::clear_video_frame(VideoFrameResourceId frame_id)
{
    if (m_video_frames.contains(frame_id.value()))
        m_video_frames.set(frame_id.value(), nullptr);
}

void DisplayListResourceStorage::update_compositor_surface(CompositorSurfaceId surface_id, Gfx::SharedImage&& shared_image)
{
    auto shared_image_buffer = Gfx::SharedImageBuffer::import_from_shared_image(move(shared_image));
    m_compositor_surfaces.set(surface_id.value(), Gfx::DecodedImageFrame { *shared_image_buffer.bitmap() });
}

void DisplayListResourceStorage::clear_compositor_surface(CompositorSurfaceId surface_id)
{
    m_compositor_surfaces.remove(surface_id.value());
}

}
