#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <map>
#include <set>
#include <string>
#include <drm_fourcc.h>
#include <wayfire/core.hpp>
#include <wayfire/output-layout.hpp>
#include <wayfire/output.hpp>
#include <wayfire/plugin.hpp>
#include <wayfire/render-manager.hpp>
#include <wayfire/nonstd/wlroots-full.hpp>
#include <wayfire/plugins/common/shared-core-data.hpp>
#include <wayfire/plugins/ipc/ipc-method-repository.hpp>

/* A small on-demand summary of the actual scene below SCREENSAVER. No
 * screenshot protocol, wallpaper paths, pixel grids or continuous render hook. */
class oledsaver_sampler_plugin_t : public wf::plugin_interface_t
{
    using clock = std::chrono::steady_clock;
    wf::shared_data::ref_ptr_t<wf::ipc::method_repository_t> repo;
    struct sample_t
    {
        clock::time_point when;
        wf::geometry_t geometry;
        double luminance;
    };
    std::map<std::string, sample_t> cache;

    static double linear(double c)
    {
        return c <= 0.04045 ? c / 12.92 : std::pow((c + 0.055) / 1.055, 2.4);
    }

    wf::json_t sample()
    {
        wf::json_t response;
        response["version"] = 1;
        response["outputs"] = wf::json_t{};
        std::set<std::string> present;
        for (auto output : wf::get_core().output_layout->get_outputs())
        {
            const std::string name = output->handle->name;
            present.insert(name);
            wf::json_t result;
            const auto geometry = output->get_layout_geometry();
            const auto now = clock::now();
            if (output->is_inhibited() || geometry.width <= 0 || geometry.height <= 0)
            {
                cache.erase(name);
                result["error"] = "output unavailable or locked";
                response["outputs"][name] = result;
                continue;
            }

            auto cached = cache.find(name);
            if (cached != cache.end() && cached->second.geometry == geometry &&
                now - cached->second.when < std::chrono::seconds(1))
            {
                result["luminance"] = cached->second.luminance;
                result["cached"] = true;
                response["outputs"][name] = result;
                continue;
            }

            const double scale = 64.0 / std::max(geometry.width, geometry.height);
            const int width = std::max(1, int(std::ceil(geometry.width * scale)));
            const int height = std::max(1, int(std::ceil(geometry.height * scale)));
            wf::auxilliary_buffer_t desktop, buffer;
            // Custom scene effects can use output pixel sizes internally. Compose
            // at the output's native size, then downsample on the GPU; only the
            // tiny summary is read back. Both buffers are released after sampling.
            const wf::dimensions_t native{output->handle->width, output->handle->height};
            bool ok = desktop.allocate(native) != wf::buffer_reallocation_result_t::FAILED &&
                buffer.allocate({width, height}) != wf::buffer_reallocation_result_t::FAILED;
            if (ok)
            {
                wf::render_target_t target{desktop};
                target.geometry = geometry;
                target.scale = output->handle->scale;
                target.wl_transform = output->handle->transform;
                ok = output->render->render_desktop(target);
                if (ok)
                {
                    buffer.get_renderbuffer().blit(desktop, {0, 0, native.width, native.height},
                        {0, 0, width, height});
                }
            }

            // ABGR8888 is RGBA byte order on little-endian systems. Reading words
            // keeps the channel extraction independent of host byte order.
            std::array<uint32_t, 64 * 64> pixels{};
            auto texture = ok ? buffer.get_texture() : nullptr;
            if (texture)
            {
                wlr_texture_read_pixels_options options{};
                options.data = pixels.data();
                options.format = DRM_FORMAT_ABGR8888;
                options.stride = width * sizeof(uint32_t);
                ok = wlr_texture_read_pixels(texture, &options);
            } else
            {
                ok = false;
            }

            if (ok)
            {
                double luminance = 0.0;
                for (int i = 0; i < width * height; ++i)
                {
                    const auto pixel = pixels[i];
                    luminance += 0.2126 * linear((pixel & 255) / 255.0) +
                        0.7152 * linear(((pixel >> 8) & 255) / 255.0) +
                        0.0722 * linear(((pixel >> 16) & 255) / 255.0);
                }

                luminance /= width * height;
                cache[name] = {now, geometry, luminance};
                result["luminance"] = luminance;
                result["cached"] = false;
                result["width"] = width;
                result["height"] = height;
            } else
            {
                cache.erase(name);
                result["error"] = "desktop sample failed";
            }

            response["outputs"][name] = result;
        }

        for (auto it = cache.begin(); it != cache.end();)
        {
            if (!present.count(it->first)) { it = cache.erase(it); } else { ++it; }
        }

        return response;
    }

  public:
    void init() override
    {
        repo->register_method("oledsaver/sample_desktop", [this] (wf::json_t) { return sample(); });
    }

    void fini() override
    {
        repo->unregister_method("oledsaver/sample_desktop");
        cache.clear();
    }
};

DECLARE_WAYFIRE_PLUGIN(oledsaver_sampler_plugin_t);