#include <algorithm>
#include <cctype>
#include <cerrno>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <fcntl.h>
#include <filesystem>
#include <fstream>
#include <gbm.h>
#include <iostream>
#include <print>
#include <ranges>
#include <span>
#include <stdexcept>
#include <string>
#include <string_view>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <system_error>
#include <type_traits>
#include <unistd.h>
#include <unordered_map>
#include <utility>
#include <vector>
#include <wayland-client-core.h>
#include <wayland-client-protocol.h>
#include <wayland-util.h>

#ifdef __linux__
#include <linux/openat2.h>
#endif

#include "ext-image-capture-source-v1-client-protocol.h"
#include "ext-image-copy-capture-v1-client-protocol.h"
#include "linux-dmabuf-v1-client-protocol.h"

namespace {

constexpr bool iequals(std::string_view a, std::string_view b) {
    return a.size() == b.size() && std::equal(a.begin(), a.end(), b.begin(), [](char ac, char bc) {
               return std::tolower(static_cast<unsigned char>(ac)) ==
                      std::tolower(static_cast<unsigned char>(bc));
           });
}

constexpr auto CAPTURE_FPS = 60;
constexpr auto CAPTURE_SECONDS = 5;

} // namespace

namespace wayland {

struct State {
    struct wl_display *display{};
    struct wl_registry *registry{};
    struct ext_image_copy_capture_manager_v1 *copy_capture_manager{};
    struct ext_image_copy_capture_session_v1 *copy_capture_session{};
    struct ext_image_capture_source_v1 *capture_source{};
    struct ext_output_image_capture_source_manager_v1 *capture_source_manager{};
    struct zwp_linux_dmabuf_v1 *linux_dmabuf{};

    struct OutputInfo {
        std::string name;
        std::string description;
    };
    std::unordered_map<struct wl_output *, OutputInfo> outputs;

    // We do not have owership over _dmabuf_devices.
    struct wl_array *_dmabuf_devices{};
    std::span<dev_t> dmabuf_devices;
    struct DmaBufFormat {
        uint32_t format{0};
        // We do not have owership over _modifiers.
        struct wl_array *_modifiers{};
        std::span<uint64_t> modifiers;
    } dmabuf_format;

    uint32_t width{0};
    uint32_t height{0};
    uint32_t stride{0};
    bool session_ready{false};
    bool frame_ready{false};
    std::vector<uint8_t> frame_data;

    std::error_code error;

    struct DmabufState {
        struct zwp_linux_buffer_params_v1 *params{};
        struct wl_buffer *buffer{};
        int fd{-1};
        void *data{};
        size_t size{0};
        uint32_t stride{0};
        bool buffer_done{false};
        bool buffer_failed{false};
        void *map_data{};
    } dmabuf_state;

    int drm_fd{-1};
    struct gbm_device *gbm_device{};
    struct gbm_bo *gbm_bufferobject{};

    explicit State() = default;
    State(const State &) = delete;
    State &operator=(const State &) = delete;
    State(State &&) = delete;
    State &operator=(State &&) = delete;

    ~State() {
        if (copy_capture_manager != nullptr) {
            ext_image_copy_capture_manager_v1_destroy(copy_capture_manager);
        }
        if (copy_capture_session != nullptr) {
            ext_image_copy_capture_session_v1_destroy(copy_capture_session);
        }
        if (capture_source != nullptr) {
            ext_image_capture_source_v1_destroy(capture_source);
        }
        if (capture_source_manager != nullptr) {
            ext_output_image_capture_source_manager_v1_destroy(capture_source_manager);
        }
        if (linux_dmabuf != nullptr) {
            zwp_linux_dmabuf_v1_destroy(linux_dmabuf);
        }
        for (const auto &[output, _] : outputs) {
            if (output != nullptr) {
                wl_output_destroy(output);
            }
        }
        if (registry != nullptr) {
            wl_registry_destroy(registry);
        }
        if (display != nullptr) {
            wl_display_disconnect(display);
        }
        if (gbm_bufferobject != nullptr) {
            gbm_bo_destroy(gbm_bufferobject);
        }
        if (gbm_device != nullptr) {
            gbm_device_destroy(gbm_device);
        }
    }
};

namespace error {

enum class Errc : uint8_t {
    ListenerAddFailed = 1,
    DmabufBufferFailed,
};

class ErrorCategory : public std::error_category {
public:
    [[nodiscard]] const char *name() const noexcept override { return "Wayland"; }

    [[nodiscard]] std::string message(int ev) const override {
        switch (static_cast<Errc>(ev)) {
        case Errc::ListenerAddFailed:
            return "failed to add listener to Wayland global";
        case Errc::DmabufBufferFailed:
            return "failed to create wl_buffer from dmabuf";
        default:
            // unreachable
            std::terminate();
        }
    }
};

namespace {

inline std::error_code make_error_code(Errc e) { return {static_cast<int>(e), ErrorCategory()}; }

} // namespace

} // namespace error

namespace output {

constexpr struct wl_output_listener listener = {
    .geometry = [](void * /*data*/, struct wl_output * /*output*/, int32_t /*x*/, int32_t /*y*/,
                   int32_t /*physical_width*/, int32_t /*physical_height*/, int32_t /*subpixel*/,
                   const char * /*make*/, const char * /*model*/, int32_t /*transform*/) {},
    .mode = [](void * /*data*/, struct wl_output * /*output*/, uint32_t /*flags*/, int32_t /*width*/,
               int32_t /*height*/, int32_t /*refresh*/) {},
    .done = [](void * /*data*/, struct wl_output * /*output*/) {},
    .scale = [](void * /*data*/, struct wl_output * /*output*/, int32_t /*factor*/) {},
    .name =
        [](void *data, struct wl_output *output, const char *name) {
            auto *state = static_cast<State *>(data);
            try {
                state->outputs.at(output).name = name;
            } catch (const std::out_of_range &) {
                state->error = std::make_error_code(std::errc::result_out_of_range);
            }
        },
    .description =
        [](void *data, struct wl_output *output, const char *description) {
            auto *state = static_cast<State *>(data);
            try {
                state->outputs.at(output).description = description;
            } catch (const std::out_of_range &err) {
                state->error = std::make_error_code(std::errc::result_out_of_range);
            }
        },
};

} // namespace output

namespace registry {

constexpr struct wl_registry_listener listener = {
    .global =
        [](void *data, struct wl_registry *registry, uint32_t name, const char *interface,
           uint32_t /*version*/) {
            using error::Errc;

            auto *state = static_cast<State *>(data);
            if (strcmp(interface, "wl_output") == 0) {
                auto global = static_cast<struct wl_output *>(
                    wl_registry_bind(registry, name, &wl_output_interface, 4));
                state->outputs.insert({global, State::OutputInfo{}});
                const int r = wl_output_add_listener(global, &output::listener, state);
                if (r < 0) {
                    state->error = error::make_error_code(Errc::ListenerAddFailed);
                }
            } else if (strcmp(interface, "ext_image_copy_capture_manager_v1") == 0) {
                state->copy_capture_manager = static_cast<struct ext_image_copy_capture_manager_v1 *>(
                    wl_registry_bind(registry, name, &ext_image_copy_capture_manager_v1_interface, 1));
            } else if (strcmp(interface, "ext_output_image_capture_source_manager_v1") == 0) {
                state->capture_source_manager =
                    static_cast<struct ext_output_image_capture_source_manager_v1 *>(wl_registry_bind(
                        registry, name, &ext_output_image_capture_source_manager_v1_interface, 1));
            } else if (strcmp(interface, "zwp_linux_dmabuf_v1") == 0) {
                state->linux_dmabuf = static_cast<struct zwp_linux_dmabuf_v1 *>(
                    wl_registry_bind(registry, name, &zwp_linux_dmabuf_v1_interface, 1));
            }
        },
    .global_remove =
        [](void * /*data*/, struct wl_registry * /*registry*/, uint32_t /*name*/) {
            // If we're already bound to a global that has been removed, the
            // object will remain valid until we destroy it, meaning that any
            // actions we attempt to take will effectively fail, and hence be
            // caught by other error checking logic.
        },
};

} // namespace registry

namespace session {

constexpr struct ext_image_copy_capture_session_v1_listener listener = {
    .buffer_size =
        // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
    [](void *data, struct ext_image_copy_capture_session_v1 * /*session*/, uint32_t width, uint32_t height) {
        auto *state = static_cast<State *>(data);
        state->width = width;
        state->height = height;
        state->stride = width * 4; // RGBA8888
    },
    .shm_format = [](void * /*data*/, struct ext_image_copy_capture_session_v1 * /*session*/,
                     uint32_t /*format*/) {},
    .dmabuf_device =
        [](void *data, struct ext_image_copy_capture_session_v1 * /*session*/, struct wl_array *device) {
            auto *state = static_cast<State *>(data);
            state->_dmabuf_devices = device;
            state->dmabuf_devices =
                std::span(static_cast<dev_t *>(device->data), device->size / sizeof(dev_t));
        },
    .dmabuf_format =
        [](void *data, struct ext_image_copy_capture_session_v1 * /*session*/, uint32_t format,
           struct wl_array *modifiers) {
            auto *state = static_cast<State *>(data);
            state->dmabuf_format = State::DmaBufFormat{
                .format = format,
                ._modifiers = modifiers,
                .modifiers =
                    std::span(static_cast<uint64_t *>(modifiers->data), modifiers->size / sizeof(uint64_t)),
            };
        },
    .done =
        [](void *data, struct ext_image_copy_capture_session_v1 * /*session*/) {
            auto *state = static_cast<State *>(data);
            state->session_ready = true;
        },
    .stopped = [](void * /*data*/,
                  struct ext_image_copy_capture_session_v1 * /*session*/) { /* TODO: destroy session.  */ },
};

} // namespace session

namespace linux_dmabuf {

constexpr struct zwp_linux_buffer_params_v1_listener params_listener = {
    .created =
        [](void *data, struct zwp_linux_buffer_params_v1 *params, struct wl_buffer *buffer) {
            auto *state = static_cast<wayland::State *>(data);
            state->dmabuf_state.buffer = buffer;
            state->dmabuf_state.buffer_done = true;
            // TODO: Do we need to actually destroy the params here?
            zwp_linux_buffer_params_v1_destroy(params);
        },
    .failed =
        [](void *data, struct zwp_linux_buffer_params_v1 *params) {
            auto *state = static_cast<wayland::State *>(data);
            state->dmabuf_state.buffer_failed = true;
            state->error = error::make_error_code(error::Errc::DmabufBufferFailed);
            // TODO: Do we need to actually destroy the params here?
            zwp_linux_buffer_params_v1_destroy(params);
        }};

} // namespace linux_dmabuf

namespace frame {

constexpr struct ext_image_copy_capture_frame_v1_listener listener = {
    .transform = [](void *, struct ext_image_copy_capture_frame_v1 *, uint32_t) {},
    .damage = [](void *, struct ext_image_copy_capture_frame_v1 *, int32_t, int32_t, int32_t, int32_t) {},
    .presentation_time = [](void *, struct ext_image_copy_capture_frame_v1 *, uint32_t, uint32_t,
                            uint32_t) {},
    .ready =
        [](void *data, struct ext_image_copy_capture_frame_v1 * /*frame*/) {
            auto state = static_cast<wayland::State *>(data);
            state->frame_ready = true;
        },
    .failed =
        [](void *data, struct ext_image_copy_capture_frame_v1 * /*frame*/, uint32_t /*reason*/) {
            // TODO: actually handle failure.
            std::println(std::cerr, "Debug: frame listener failed");
            auto state = static_cast<wayland::State *>(data);
            state->frame_ready = true;
        }};

} // namespace frame

} // namespace wayland

namespace std {

template <> struct is_error_code_enum<wayland::error::Errc> : true_type {};

} // namespace std

int main() {
    using std::cerr, std::print, std::println;
    using namespace wayland;
    namespace fs = std::filesystem;
    namespace ranges = std::ranges;

    int r{0};
    State state;

    state.display = wl_display_connect(nullptr);
    if (state.display == nullptr) {
        println(cerr, "Failed to connect to Wayland display: {}.", strerror(errno));
        return EXIT_FAILURE;
    }

    state.registry = wl_display_get_registry(state.display);
    if (state.registry == nullptr) {
        println(cerr, "Failed to get Wayland registry: {}.", strerror(errno));
        return EXIT_FAILURE;
    }

    r = wl_registry_add_listener(state.registry, &registry::listener, &state);
    if (r < 0) {
        println(cerr, "Failed to add listener to registry.");
        return EXIT_FAILURE;
    }

    r = wl_display_roundtrip(state.display);
    if (r < 0) {
        println(cerr, "Failed to roundtrip Wayland display.");
        return EXIT_FAILURE;
    }

    // Registry has been processed by now.
    if (state.error) {
        println(cerr, "Error: {}.", state.error.message());
        return EXIT_FAILURE;
    }

    if (state.outputs.empty()) {
        println(cerr, "Failed to find any outputs.");
        return EXIT_FAILURE;
    }

    r = wl_display_dispatch(state.display);
    if (r < 0) {
        println(cerr, "Failed to dispatch Wayland display.");
        return EXIT_FAILURE;
    }

    // All callbacks for globals have been processed by now.
    if (state.error) {
        println(cerr, "Error: {}.", state.error.message());
        return EXIT_FAILURE;
    }

    if (state.copy_capture_manager == nullptr) {
        println(cerr, "Compositor doesn't support ext-image-copy-capture-v1.");
        return EXIT_FAILURE;
    }

    if (state.capture_source_manager == nullptr) {
        println(cerr, "Compositor doesn't support ext-image-capture-source-v1.");
        return EXIT_FAILURE;
    }

    if (state.linux_dmabuf == nullptr) {
        println(cerr, "Compositor doesn't support linux-dmabuf-v1.");
        return EXIT_FAILURE;
    }

    println("Outputs:");
    for (const auto &[_, info] : state.outputs) {
        println("  {}: {}", info.name, info.description);
    }

    print("Enter output to capture: ");
    std::string input;
    std::getline(std::cin, input);

    const auto it = ranges::find_if(state.outputs.begin(), state.outputs.end(), [input](const auto &pair) {
        return iequals(input, pair.second.name) || iequals(input, pair.second.description);
    });

    if (it == state.outputs.end()) {
        println(cerr, "Output `{}` does not exist.", input);
        return EXIT_FAILURE;
    }

    auto output = it->first;

    state.capture_source =
        ext_output_image_capture_source_manager_v1_create_source(state.capture_source_manager, output);
    state.copy_capture_session =
        ext_image_copy_capture_manager_v1_create_session(state.copy_capture_manager, state.capture_source, 0);

    r = ext_image_copy_capture_session_v1_add_listener(state.copy_capture_session, &session::listener,
                                                       &state);
    if (r < 0) {
        println(cerr, "Failed to add listener to capture session.");
        return EXIT_FAILURE;
    }

    while (!state.session_ready) {
        r = wl_display_dispatch(state.display);
        if (r < 0) {
            println(cerr, "Failed to dispatch Wayland display.");
            return EXIT_FAILURE;
        }
    }

    println("Width: {}", state.width);
    println("Height: {}", state.height);
    println("Stride: {}", state.stride);
    println("dmabuf devices:");
    for (const auto &device : state.dmabuf_devices) {
        println("  {}", device);
    }
    println("dmabuf format: {}", state.dmabuf_format.format);
    println("dmabuf modifiers:");
    for (const auto &modifier : state.dmabuf_format.modifiers) {
        println("  {}", modifier);
    }

    constexpr uint64_t LAYOUT_MODIFIER = 0; // linear
    if (ranges::find(state.dmabuf_format.modifiers, LAYOUT_MODIFIER) == state.dmabuf_format.modifiers.end()) {
        println(cerr, "Error: Compositor doesn't support linear layout modifier.");
        return EXIT_FAILURE;
    }

    const dev_t drm_dev = state.dmabuf_devices.front();
    const fs::path dri_dir{"/dev/dri"};
    std::string render_node;

    if (!fs::exists(dri_dir) || !fs::is_directory(dri_dir)) {
        std::println(std::cerr, "Directory {} does not exist", dri_dir.string());
        return EXIT_FAILURE;
    }

    bool found = false;
    for (const auto &entry : fs::directory_iterator(dri_dir)) {
        if (!entry.is_character_file()) {
            continue;
        }

        const auto &path = entry.path();

        if (!path.filename().string().starts_with("renderD")) {
            continue;
        }

        struct stat st{};
        if (::stat(path.string().c_str(), &st) == 0 && st.st_rdev == drm_dev) {
            render_node = path.string();
            found = true;
            break;
        }
    }

    if (!found) {
        std::println(std::cerr, "Error: No matching DRM render node found for dev_t `{}`.", drm_dev);
        return EXIT_FAILURE;
    }

    constexpr struct open_how how{.flags = O_RDWR | O_CLOEXEC, .mode = 0, .resolve = 0};
    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-vararg)
    state.drm_fd = static_cast<int>(syscall(SYS_openat2, AT_FDCWD, render_node.c_str(), &how, sizeof(how)));
    if (state.drm_fd < 0) {
        println(cerr, "Error: failed to open DRI render node: {}.", strerror(errno));
        return EXIT_FAILURE;
    }

    state.gbm_device = gbm_create_device(state.drm_fd);
    if (state.gbm_device == nullptr) {
        println(cerr, "Error: failed to create GBM device: {}.", strerror(errno));
        return EXIT_FAILURE;
    }

    const uint64_t *modifiers = state.dmabuf_format.modifiers.data();
    size_t modifier_count = state.dmabuf_format.modifiers.size();

    println("modifier count: {}", modifier_count);
    println("modifiers:");
    for (const auto &modifier : state.dmabuf_format.modifiers) {
        println("  {}", modifier);
    }

    state.gbm_bufferobject = gbm_bo_create_with_modifiers(
        state.gbm_device, state.width, state.height, state.dmabuf_format.format, modifiers, modifier_count);
    if (state.gbm_bufferobject == nullptr) {
        println(cerr, "Error: failed to create GBM buffer {}.", strerror(errno));
        return EXIT_FAILURE;
    }

    state.dmabuf_state.fd = gbm_bo_get_fd(state.gbm_bufferobject);
    if (state.dmabuf_state.fd < 0) {
        println(cerr, "Error: failed to export GBM buffer as DMA-BUF: {}.", strerror(errno));
        gbm_bo_destroy(state.gbm_bufferobject);
        return EXIT_FAILURE;
    }

    state.dmabuf_state.data =
        gbm_bo_map(state.gbm_bufferobject, 0, 0, state.width, state.height, GBM_BO_TRANSFER_READ_WRITE,
                   &state.dmabuf_state.stride, &state.dmabuf_state.map_data);
    if (state.dmabuf_state.data == nullptr) {
        println(cerr, "Error: failed to map GBM buffer object: {}.", strerror(errno));
        return EXIT_FAILURE;
    }

    state.dmabuf_state.params = zwp_linux_dmabuf_v1_create_params(state.linux_dmabuf);
    if (state.dmabuf_state.params == nullptr) {
        println(cerr, "Error: failed to create linux-dmabuf-v1 parameters.");
        return EXIT_FAILURE;
    }

    zwp_linux_buffer_params_v1_add(state.dmabuf_state.params, state.dmabuf_state.fd, 0, 0,
                                   state.dmabuf_state.stride, LAYOUT_MODIFIER, 0);

    r = zwp_linux_buffer_params_v1_add_listener(state.dmabuf_state.params, &linux_dmabuf::params_listener,
                                                &state);
    if (r < 0) {
        println(cerr, "Failed to add listener buffer parameters.");
        return EXIT_FAILURE;
    }

    // TODO: why do we pass a signed integer here?
    zwp_linux_buffer_params_v1_create(state.dmabuf_state.params, static_cast<int32_t>(state.width),
                                      static_cast<int32_t>(state.height), state.dmabuf_format.format, 0);

    while (!state.dmabuf_state.buffer_done && !state.dmabuf_state.buffer_failed) {
        r = wl_display_roundtrip(state.display);
        if (r < 0) {
            println(cerr, "Failed to roundtrip Wayland display.");
            return EXIT_FAILURE;
        }
    }

    if (state.dmabuf_state.buffer == nullptr || state.dmabuf_state.buffer_failed) {
        println(cerr, "Error: {}.", state.error.message());
        return EXIT_FAILURE;
    }

    std::ofstream capture("capture.raw", std::ios::binary);
    if (!capture) {
        println(cerr, "Error: could not open `capture.raw` for writing: {}.", strerror(errno));
        return EXIT_FAILURE;
    }

    for (auto _ : std::views::iota(0, CAPTURE_FPS * CAPTURE_SECONDS)) {
        state.frame_ready = false;
        auto frame = ext_image_copy_capture_session_v1_create_frame(state.copy_capture_session);
        ext_image_copy_capture_frame_v1_add_listener(frame, &frame::listener, &state);

        ext_image_copy_capture_frame_v1_attach_buffer(frame, state.dmabuf_state.buffer);

        // I'm not sure why this takes a signed integer considering that all the other APIs take a
        // signed integer, but casting from an uint32_t to an int32_t is the best we can do.
        ext_image_copy_capture_frame_v1_damage_buffer(frame, 0, 0, static_cast<int32_t>(state.width),
                                                      static_cast<int32_t>(state.height));
        ext_image_copy_capture_frame_v1_capture(frame);

        while (!state.frame_ready) {
            r = wl_display_dispatch(state.display);
            if (r < 0) {
                println(cerr, "Failed to dispatch Wayland display.");
                return EXIT_FAILURE;
            }
        }

        capture.write(static_cast<const char *>(state.dmabuf_state.data),
                      state.dmabuf_state.stride * state.height);

        // TODO: are we responsible for destroying the frame?
        ext_image_copy_capture_frame_v1_destroy(frame);
    }

    return EXIT_SUCCESS;
}