#include "CWayland.h"
#include <wayland-client.h>
#include <array>
#include <vector>
#include <queue>
#include <memory>
#include <unordered_map>
#include <chrono>
#include <cassert>
#include <poll.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <signal.h>

namespace
{
    class ref
    {
        lean_object* ptr;

        struct steal_token {};
        
        explicit ref(lean_obj_arg ptr, steal_token) : ptr{ ptr }
        {
        }

		void clear()
        {
            if (ptr)
            {
                lean_dec(ptr);
                ptr = nullptr;
            }
        }
    
    public:
        static ref steal(lean_obj_arg ptr)
        {
            return ref{ ptr, steal_token{} };
        }

        explicit ref() : ptr{ nullptr }
        {
        }

        explicit ref(b_lean_obj_arg ptr) : ptr{ ptr }
        {
            lean_inc(ptr);
        }

        ~ref()
        {
            clear();
        }

        ref(ref const& other) : ref{ other.get() }
        {
        }

        ref(ref&& other) : ref{ std::move(other).release(), steal_token{} }
        {
        }

        ref& operator=(ref const& other)
        {
            if (ptr != other.ptr)
            {
                clear();
                ptr = other.get();
                lean_inc(ptr);
            }
            return *this;
        }

        ref& operator=(ref&& other)
        {
            if (ptr != other.ptr)
            {
                clear();
                ptr = std::move(other).release();
            }
            return *this;
        }

        void swap(ref& other)
        {
            std::swap(ptr, other.ptr);
        }

        [[nodiscard]] b_lean_obj_res get() const
        {
            return ptr;
        }

        [[nodiscard]] lean_obj_res release() &&
        {
            lean_object* result = get();
            ptr = nullptr;
            return result;
        }
    };

    char const* option_cstr_or_null(b_lean_obj_arg str)
    {
        if (lean_is_scalar(str))
        {
            return nullptr;
        }
        else
        {
            lean_object* val = lean_ctor_get(str, 0);
            return lean_string_cstr(val);
        }
    }

    lean_obj_res io_result_mk_from_errno(char const* file_name)
    {
        int err = errno;
        ref file_name_str = ref::steal(lean_mk_string(file_name));
        lean_object* error = lean_decode_io_error(err, file_name_str.get());

        return lean_io_result_mk_error(error);
    }
}

// -- Constants

uint32_t waylean_SIGINT(void)
{
    return SIGINT;
}

uint32_t waylean_SIGUSR1(void)
{
    return SIGUSR1;
}

uint32_t waylean_SIGUSR2(void)
{
    return SIGUSR2;
}

// -- Display
    
namespace
{
    lean_external_class* waylean_display_class()
    {
        static lean_external_class* klass = lean_register_external_class(
            [](void* ptr) { wl_display_disconnect(static_cast<wl_display*>(ptr)); },
            [](void*, b_lean_obj_arg fn) { }
        );
        return klass;
    }
}

// Display.new : (name : @& Option String) → IO Display
lean_obj_res waylean_display_new(b_lean_obj_arg name)
{
    char const* name_cstr = option_cstr_or_null(name);
    wl_display* disp = wl_display_connect(name_cstr);
    if (!disp) { return io_result_mk_from_errno(name_cstr ? name_cstr : "<default>"); }

    lean_object* waylean_disp = lean_alloc_external(waylean_display_class(), disp);
    return lean_io_result_mk_ok(waylean_disp);
}

// Display.roundtrip : (display : @& Display) → IO Unit
lean_obj_res waylean_display_roundtrip(b_lean_obj_arg display)
{
    wl_display* disp = static_cast<wl_display*>(lean_get_external_data(display));
    if (wl_display_roundtrip(disp) == -1) { return io_result_mk_from_errno("<unknown>"); }
    return lean_io_result_mk_ok(lean_box(0));
}

// -- Event Loop

namespace
{
    class signal_manager
    {
        std::array<int, 2> fds{ -1, -1 };
        static inline int tx_fd = -1;

        explicit signal_manager()
        {
            if (pipe(fds.data()) == -1)
            {
                perror("signal_manager::signal_manager: pipe");
                return;
            }

            if (fcntl(fds[0], F_SETFD, FD_CLOEXEC) == -1)
            {
                perror("signal_manager::signal_manager: fcntl");
            }
            if (fcntl(fds[1], F_SETFD, FD_CLOEXEC) == -1)
            {
                perror("signal_manager::signal_manager: fcntl");
            }
            int flags = fcntl(fds[0], F_GETFL, 0);
            if (flags == -1)
            {
                perror("signal_manager::signal_manager: fcntl");
            }
            else
            {
                if (fcntl(fds[0], F_SETFL, flags | O_NONBLOCK) == -1)
                {
                    perror("signal_manager::signal_manager: fcntl");
                }
            }

            struct sigaction sigpipe_action{};
            sigpipe_action.sa_handler = SIG_IGN;
            if (sigaction(SIGPIPE, &sigpipe_action, nullptr) == -1)
            {
                perror("signal_manager::signal_manager: sigaction");
            }
        }

        static signal_manager& instance()
        {
            static signal_manager res;
            return res;
        }

        static void signal_handler(int signum)
        {
            write(tx_fd, &signum, sizeof(signum));
        }

    public:
        static int get_rx_fd()
        {
            return instance().fds[0];
        }

        static void handle(int signum)
        {
            tx_fd = instance().fds[1];
            struct sigaction action{};
            action.sa_handler = signal_handler;
            action.sa_flags = SA_RESTART;
            sigemptyset(&action.sa_mask);
            if (sigaction(signum, &action, nullptr) == -1)
            {
                perror("signal_manager::handle: sigaction");
            }
        }
    };

    struct event_source
    {
        enum class type
        {
            display,
            timer,
            signal,
        } type{};
        size_t id{};
        int signum{};
    };

    struct display_source
    {
        ref display_ref{};
        wl_display* disp{};
        size_t id{};
    };

    struct timer_source
    {
        std::chrono::high_resolution_clock::time_point scheduled_time{};
        bool invalid{};
        ref message{};
        size_t id{};
    };

    bool operator>(timer_source const& lhs, timer_source const& rhs)
    {
        return lhs.scheduled_time > rhs.scheduled_time;
    }

    struct timer_source_queue : std::priority_queue<timer_source, std::vector<timer_source>, std::greater<timer_source>>
    {
        void invalidate(size_t id)
        {
            for (auto& source : c)
            {
                if (source.id == id) { source.invalid = true; }
            }
        }

        timer_source release_top()
        {
            return timer_source{ std::move(c.front()) };
        }
    };

    struct signal_source
    {
        ref message{};
        size_t id{};
    };

    struct scheduled_event
    {
        ref message{};
        size_t id{};
    };

    class event_loop
    {
        size_t next_id{ 0 };
        bool should_quit{};
        std::vector<display_source> displays;
        std::unordered_map<int, std::vector<signal_source>> signals;
        timer_source_queue timers;
        std::vector<pollfd> pollfds;
        std::deque<scheduled_event> scheduled;

        void rebuild_pollfds()
        {
            pollfds.clear();

            pollfd signal_pollfd{};
            signal_pollfd.fd = signal_manager::get_rx_fd();
            signal_pollfd.events = POLLIN;
            pollfds.emplace_back(signal_pollfd);

            for (auto& src : displays)
            {
                pollfd disp_pollfd{};
                disp_pollfd.fd = wl_display_get_fd(src.disp);
                disp_pollfd.events = POLLIN;
                pollfds.emplace_back(disp_pollfd);
            }
        }

        int get_next_timer_ms()
        {
            if (timers.empty()) { return -1; }

            auto duration = timers.top().scheduled_time - std::chrono::high_resolution_clock::now();
            auto duration_ms = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
            return std::clamp<decltype(duration_ms)>(duration_ms, 0, std::numeric_limits<int>::max());
        }

    public:
        explicit event_loop()
        {
            rebuild_pollfds();
        }

        event_source add_display(ref display)
        {
            auto res = event_source{};
            res.type = event_source::type::display;
            res.id = next_id++;

            auto source = display_source{};
            source.disp = static_cast<wl_display*>(lean_get_external_data(display.get()));
            source.display_ref = std::move(display);
            source.id = res.id;
            displays.emplace_back(std::move(source));

            rebuild_pollfds();

            return res;
        }

        event_source add_timer(unsigned delay_ms, ref message)
        {
            auto res = event_source{};
            res.type = event_source::type::timer;
            res.id = next_id++;

            auto source = timer_source{};
            source.scheduled_time = std::chrono::high_resolution_clock::now() + std::chrono::milliseconds{ delay_ms };
            source.message = std::move(message);
            source.id = res.id;
            timers.emplace(std::move(source));

            return res;
        }

        event_source add_signal(int signum, ref message)
        {
            signal_manager::handle(signum);

            auto res = event_source{};
            res.type = event_source::type::signal;
            res.id = next_id++;
            res.signum = signum;

            auto source = signal_source{};
            source.message = std::move(message);
            source.id = res.id;
            signals[signum].emplace_back(std::move(source));

            return res;
        }

        void remove(event_source const& source)
        {
            // Discharge event source long-term.
            switch (source.type)
            {
                case event_source::type::display:
                {
                    auto it = std::find_if(
                        displays.begin(), displays.end(),
                        [&](display_source const& other) { return other.id == source.id; }
                    );
                    if (it != displays.end()) { displays.erase(it); }
                    break;
                }
                case event_source::type::timer:
                {
                    timers.invalidate(source.id);
                    break;
                }
                case event_source::type::signal:
                {
                    auto& sources = signals[source.signum];
                    auto it = std::find_if(
                        sources.begin(), sources.end(),
                        [&](signal_source const& other) { return other.id == source.id; }
                    );
                    if (it != sources.end()) { sources.erase(it); }
                    break;
                }
            }

            // Discharge immediately scheduled events.
            scheduled.erase(
                std::remove_if(
                    scheduled.begin(), scheduled.end(),
                    [&](scheduled_event const& ev) { return ev.id == source.id; }
                ),
                scheduled.end()
            );
        }

        void quit()
        {
            should_quit = true;
        }

        ref poll()
        {
            while (scheduled.empty() && !should_quit)
            {
                // Otherwise, fetch any new events.
                for (auto& src : displays)
                {
                    while (wl_display_prepare_read(src.disp))
                    {
                        if (wl_display_dispatch_pending(src.disp) == -1)
                        {
                            perror("event_loop::poll: wl_display_dispatch_pending");
                        }
                    }

                    if (wl_display_flush(src.disp) == -1)
                    {
                        perror("event_loop::poll: wl_display_flush");
                    }
                }
                
                int res;
                do
                {
                    res = ::poll(pollfds.data(), pollfds.size(), get_next_timer_ms());
                }
                while (res == -1 && errno == EINTR);
                if (res == -1)
                {
                    perror("event_loop::poll: poll");
                }

                // 1. Process signals.
                int signum;
                ssize_t nread;
                while ((nread = read(signal_manager::get_rx_fd(), &signum, sizeof(signum))) > 0)
                {
                    if (nread == sizeof(signum))
                    {
                        auto it = signals.find(signum);
                        if (it != signals.end())
                        {
                            for (auto& source : it->second)
                            {
                                auto ev = scheduled_event{};
                                ev.id = source.id;
                                ev.message = source.message;
                                scheduled.emplace_back(std::move(ev));
                            }
                        }
                    }
                }
                if (nread == -1 && errno != EWOULDBLOCK && errno != EAGAIN)
                {
                    perror("event_loop::poll: read");
                }

                // 2. Process display events.
                for (auto& src : displays)
                {
                    if (wl_display_read_events(src.disp) == -1)
                    {
                        perror("event_loop::poll: wl_display_read_events");
                        continue;
                    }
                    if (wl_display_dispatch_pending(src.disp) == -1)
                    {
                        perror("event_loop::poll: wl_display_dispatch_pending");
                    }
                }

                // 3. Process expired timers.
                auto new_now = std::chrono::high_resolution_clock::now();
                while (!timers.empty())
                {
                    if (timers.top().scheduled_time > new_now) { break; }
                    
                    auto timer = timers.release_top();
                    if (!timer.invalid)
                    {
                        auto ev = scheduled_event{};
                        ev.id = timer.id;
                        ev.message = std::move(timer.message);
                        scheduled.emplace_back(std::move(ev));
                    }
                    timers.pop();
                }
            }

            if (!scheduled.empty() && !should_quit)
            {
                auto event = std::move(scheduled.front());
                scheduled.pop_front();
                auto option = ref::steal(lean_alloc_ctor(1, 1, 0));
                lean_ctor_set(option.get(), 0, std::move(event.message).release());
                return option;
            }

            return ref::steal(lean_box(0));
        }
    };

    lean_external_class* waylean_event_loop_class()
    {
        static lean_external_class* klass = lean_register_external_class(
            [](void* ptr) { delete static_cast<event_loop*>(ptr); },
            [](void*, b_lean_obj_arg fn) { }
        );
        return klass;
    }

    lean_external_class* waylean_event_source_class()
    {
        static lean_external_class* klass = lean_register_external_class(
            [](void* ptr) { delete static_cast<event_source*>(ptr); },
            [](void*, b_lean_obj_arg fn) { }
        );
        return klass;
    }
}

// EventLoop.new : IO EventLoop M
lean_obj_res waylean_event_loop_new(void)
{
    event_loop* loop = new event_loop;
    lean_object* waylean_loop = lean_alloc_external(waylean_event_loop_class(), loop);
    return lean_io_result_mk_ok(waylean_loop);
}

// EventLoop.addDisplay : (eventLoop : @& EventLoop M) → (display : @& Display M) → IO EventSource
lean_obj_res waylean_event_loop_add_display(b_lean_obj_arg loop_obj, b_lean_obj_arg display)
{
    auto loop = static_cast<event_loop*>(lean_get_external_data(loop_obj));
    auto source = new event_source{ loop->add_display(ref{ display }) };
    lean_object* waylean_source = lean_alloc_external(waylean_event_source_class(), source);
    return lean_io_result_mk_ok(waylean_source);
}

// EventLoop.addTimer : (eventLoop : @& EventLoop M) → (delayMs : UInt32) → (msg : @& M) → IO EventSource
lean_obj_res waylean_event_loop_add_timer(b_lean_obj_arg loop_obj, uint32_t delay_ms, b_lean_obj_arg message)
{
    auto loop = static_cast<event_loop*>(lean_get_external_data(loop_obj));
    auto source = new event_source{ loop->add_timer(delay_ms, ref{ message }) };
    lean_object* waylean_source = lean_alloc_external(waylean_event_source_class(), source);
    return lean_io_result_mk_ok(waylean_source);
}

// EventLoop.addSignal : (eventLoop : @& EventLoop M) → (signum : UInt32) → (msg : @& M) → IO EventSource
lean_obj_res waylean_event_loop_add_signal(b_lean_obj_arg loop_obj, uint32_t signum, b_lean_obj_arg message)
{
    auto loop = static_cast<event_loop*>(lean_get_external_data(loop_obj));
    auto source = new event_source{ loop->add_signal(static_cast<int>(signum), ref{ message }) };
    lean_object* waylean_source = lean_alloc_external(waylean_event_source_class(), source);
    return lean_io_result_mk_ok(waylean_source);
}

// EventLoop.remove : (eventLoop : @& EventLoop M) → (source : @& EventSource) → IO Unit
lean_obj_res waylean_event_loop_remove(b_lean_obj_arg loop_obj, b_lean_obj_arg source_obj)
{
    auto loop = static_cast<event_loop*>(lean_get_external_data(loop_obj));
    auto source = static_cast<event_source*>(lean_get_external_data(source_obj));
    loop->remove(*source);
    return lean_io_result_mk_ok(lean_box(0));
}

// EventLoop.quit : (eventLoop : @& EventLoop M) → IO Unit
lean_obj_res waylean_event_loop_quit(b_lean_obj_arg loop_obj)
{
    auto loop = static_cast<event_loop*>(lean_get_external_data(loop_obj));
    loop->quit();
    return lean_io_result_mk_ok(lean_box(0));
}

// EventLoop.poll : (eventLoop : @& EventLoop M) → IO (Option M)
lean_obj_res waylean_event_loop_poll(b_lean_obj_arg loop_obj)
{
    struct wl_interface x; 
    auto loop = static_cast<event_loop*>(lean_get_external_data(loop_obj));
    return lean_io_result_mk_ok(loop->poll().release());
}

// -- Type Map

namespace
{
    template<std::size_t Size = 4096>
    class arena
    {
        struct region
        {
            std::array<std::byte, Size> buffer;
            std::unique_ptr<region> prev;
            std::size_t used{ 0 };
        };

        std::unique_ptr<region> curr;

    public:
        explicit arena()
        {
            curr = std::make_unique<region>();
        }

        template<typename T>
        T* alloc_uninitialized(std::size_t count = 1)
        {
            std::size_t size = count * sizeof(T);
            if (curr->used + size > Size)
            {
                // no more space -> allocate a new region
                assert(size <= Size && "object too big for arena");
                auto new_region = std::make_unique<region>();
                new_region->prev = std::move(curr);
                curr = std::move(new_region);
            }

            std::byte* res = &curr->buffer[curr->used];
            curr->used += size;
            return std::launder(reinterpret_cast<T*>(res));
        }
    };

    struct interface
    {
        wl_interface desc;
        bool initialized;
    };

    class type_map
    {
        arena<> a;
        std::unordered_map<std::string_view, interface*> interfaces;

    public:
        bool contains_iface(std::string_view iface_name)
        {
            auto it = interfaces.find(iface_name);
            return it != interfaces.end() && it->second->initialized;
        }

        interface* resolve_iface(std::string_view iface_name)
        {
            auto it = interfaces.find(iface_name);
            if (it != interfaces.end()) { return it->second; }
            
            char* name = copy_string(iface_name);
            interface* iface = a.alloc_uninitialized<interface>();
            iface->desc = wl_interface{ 0 };
            iface->desc.name = name;
            iface->initialized = false;

            interfaces[name] = iface;
            return iface;
        }

        void register_iface(b_lean_obj_arg iface_desc)
        {
            // structure Interface where
            //   name : String
            //   version : UInt32
            //   requests : Array Message
            //   events : Array Message

            lean_object* iface_name = lean_ctor_get(iface_desc, 0);

            interface* iface = resolve_iface(lean_string_cstr(iface_name));
            if (iface->initialized) { return; }
            iface->initialized = true;
            wl_interface* desc = &iface->desc;

            lean_object* iface_requests = lean_ctor_get(iface_desc, 1);
            std::size_t request_count = lean_array_size(iface_requests);
            wl_message* requests = a.alloc_uninitialized<wl_message>(request_count);
            for (std::size_t i = 0; i < request_count; ++i)
            {
                initialize_message(&requests[i], lean_array_get_core(iface_requests, i));
            }
            desc->method_count = request_count;
            desc->methods = requests;

            lean_object* iface_events = lean_ctor_get(iface_desc, 2);
            std::size_t event_count = lean_array_size(iface_events);
            wl_message* events = a.alloc_uninitialized<wl_message>(event_count);
            for (std::size_t i = 0; i < event_count; ++i)
            {
                initialize_message(&events[i], lean_array_get_core(iface_events, i));
            }
            desc->event_count = event_count;
            desc->events = events;

            desc->version = lean_ctor_get_uint32(iface_desc, 3 * sizeof(void*));
        }

    private:
        char* copy_string(std::string_view str)
        {
            char* res = a.alloc_uninitialized<char>(str.length() + 1);
            std::copy(str.begin(), str.end(), res);
            res[str.length()] = 0;
            return res;
        }

        void initialize_message(wl_message* msg, b_lean_obj_arg msg_desc)
        {
            // structure Message where
            //   name : String
            //   signature : String
            //   types : Array (Option String)

            lean_object* msg_name = lean_ctor_get(msg_desc, 0);
            msg->name = copy_string(lean_string_cstr(msg_name));

            lean_object* msg_signature = lean_ctor_get(msg_desc, 1);
            msg->signature = copy_string(lean_string_cstr(msg_signature));

            lean_object* msg_types = lean_ctor_get(msg_desc, 2);
            std::size_t type_count = lean_array_size(msg_types);
            wl_interface const** types = a.alloc_uninitialized<wl_interface const*>(type_count);
            for (std::size_t i = 0; i < type_count; ++i)
            {
                char const* type = option_cstr_or_null(lean_array_get_core(msg_types, i));
                types[i] = type ? &resolve_iface(type)->desc : nullptr;
            }
            msg->types = types;
        }
    };

    lean_external_class* waylean_type_map_class()
    {
        static lean_external_class* klass = lean_register_external_class(
            [](void* ptr) { delete static_cast<type_map*>(ptr); },
            [](void*, b_lean_obj_arg fn) { }
        );
        return klass;
    }
}

// TypeMap.new : IO TypeMap
lean_obj_res waylean_type_map_new(void)
{
    type_map* map = new type_map;
    lean_object* waylean_map = lean_alloc_external(waylean_type_map_class(), map);
    return lean_io_result_mk_ok(waylean_map);
}

// TypeMap.contains : (typeMap : @& TypeMap) → (ifaceName : @& String) → IO Bool
lean_obj_res waylean_type_map_contains(b_lean_obj_arg map_obj, b_lean_obj_arg iface_name)
{
    auto map = static_cast<type_map*>(lean_get_external_data(map_obj));
    auto res = map->contains_iface(lean_string_cstr(iface_name));
    return lean_io_result_mk_ok(lean_box(res ? 1 : 0));
}

// TypeMap.register : (typeMap : @& TypeMap) → (ifaceDesc : @& Interface) → IO Unit
lean_obj_res waylean_type_map_register(b_lean_obj_arg map_obj, b_lean_obj_arg iface_desc)
{
    auto map = static_cast<type_map*>(lean_get_external_data(map_obj));
    map->register_iface(iface_desc);
    return lean_io_result_mk_ok(lean_box(0));
}

// -- Proxy

// Proxy.fromDisplay : (display : @& Display) → IO Proxy
lean_obj_res waylean_proxy_from_display(b_lean_obj_arg display)
{
}

// Proxy.new : (factory : @& Proxy) → (typeMap : @& TypeMap) → (ifaceName : @& String) → IO Proxy
lean_obj_res waylean_proxy_new(b_lean_obj_arg factory, b_lean_obj_arg type_map, b_lean_obj_arg iface_name)
{
}

// Proxy.send : (proxy : @& Proxy) → (opcode : UInt32) → (argumentWriter : @& ArgumentWriter) → IO Unit
lean_obj_res waylean_proxy_send(b_lean_obj_arg proxy, uint32_t opcode, b_lean_obj_arg argument_writer)
{
}

// Proxy.listen : (proxy : @& Proxy) → (handler : @& UInt32 → ArgumentReader → EventLoop M × M) → IO Unit
lean_obj_res waylean_proxy_listen(b_lean_obj_arg proxy, b_lean_obj_arg handler)
{
}