/*
 * Copyright (c) 2023, Andreas Kling <andreas@ladybird.org>
 * Copyright (c) 2024-2025, stasoid <stasoid@yahoo.com>
 * Copyright (c) 2025, ayeteadoe <ayeteadoe@gmail.com>
 * Copyright (c) 2025, Ryszard Goc <ryszardgoc@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/Assertions.h>
#include <AK/Atomic.h>
#include <AK/Diagnostics.h>
#include <AK/HashMap.h>
#include <AK/NonnullOwnPtr.h>
#include <AK/Time.h>
#include <AK/Windows.h>
#include <LibCore/EventLoopImplementationWindows.h>
#include <LibCore/Notifier.h>
#include <LibCore/ThreadEventQueue.h>
#include <LibCore/TimeoutSet.h>
#include <LibCore/Timer.h>
#include <LibSync/Mutex.h>
#include <LibSync/MutexProtected.h>

struct OwnHandle {
    HANDLE handle = NULL;

    explicit OwnHandle(HANDLE h = NULL)
        : handle(h)
    {
    }

    OwnHandle(OwnHandle&& h)
    {
        handle = h.handle;
        h.handle = NULL;
    }

    // This operation can only be done when handle is NULL
    OwnHandle& operator=(OwnHandle&& other)
    {
        VERIFY(!handle);
        if (this == &other)
            return *this;
        handle = other.handle;
        other.handle = NULL;
        return *this;
    }

    ~OwnHandle()
    {
        if (handle)
            CloseHandle(handle);
    }

    bool operator==(OwnHandle const& h) const { return handle == h.handle; }
    bool operator==(HANDLE h) const { return handle == h; }
};

template<>
struct Traits<OwnHandle> : DefaultTraits<OwnHandle> {
    static unsigned hash(OwnHandle const& h) { return Traits<HANDLE>::hash(h.handle); }
};
template<>
constexpr bool IsHashCompatible<HANDLE, OwnHandle> = true;

namespace Core {

enum class CompletionType : u8 {
    Wake,
    Timer,
    Notifer,
    Process,
};

struct CompletionPacket {
    CompletionType type;
};

struct EventLoopWake final : CompletionPacket {
    OwnHandle wait_packet;
    OwnHandle wait_event;
};

// All timers of a thread share one waitable timer, armed for the earliest deadline. Own bookkeeping (rather than a
// kernel timer object per Core timer) is used because the kernel delivers simultaneously signaled wait completion
// packets in LIFO order, while timers that are due together must fire in registration order.
struct EventLoopMasterTimer final : CompletionPacket {

    ~EventLoopMasterTimer()
    {
        CancelWaitableTimer(timer.handle);
        if (wait_packet_associated) {
            NTSTATUS status = g_system.NtCancelWaitCompletionPacket(wait_packet.handle, TRUE);
            VERIFY(NT_SUCCESS(status));
        }
    }

    OwnHandle timer;
    OwnHandle wait_packet;
    bool wait_packet_associated { false };
};

class EventLoopTimer final : public EventLoopTimeout {
public:
    EventLoopTimer() = default;

    void reload(MonotonicTime const& now) { m_fire_time = now + interval; }

    virtual void fire(TimeoutSet& timeout_set, MonotonicTime current_time) override
    {
        auto strong_owner = owner.strong_ref();

        if (!strong_owner)
            return;

        if (should_reload) {
            MonotonicTime next_fire_time = m_fire_time + interval;
            if (next_fire_time <= current_time) {
                next_fire_time = current_time + interval;
            }
            m_fire_time = next_fire_time;
            if (next_fire_time != current_time) {
                timeout_set.schedule_absolute(this);
            } else {
                // NOTE: Unfortunately we need to treat timeouts with the zero interval in a
                //       special way. TimeoutSet::schedule_absolute for them will result in an
                //       infinite loop. TimeoutSet::schedule_relative, on the other hand, will do a
                //       correct thing of scheduling them for the next iteration of the loop.
                m_duration = {};
                timeout_set.schedule_relative(this);
            }
        }

        ThreadEventQueue::current().post_event(strong_owner, Event::Type::Timer);
    }

    AK::Duration interval;
    bool should_reload { false };
    WeakPtr<EventReceiver> owner;
};

struct EventLoopNotifier final : CompletionPacket {

    ~EventLoopNotifier()
    {
    }

    Notifier* notifier;
    OwnHandle wait_packet;
    OwnHandle wait_event;
};

struct EventLoopProcess final : CompletionPacket {
    ~EventLoopProcess() = default;

    OwnHandle process;
    pid_t pid;
    Function<void(pid_t)> exit_handler;
    OwnHandle jobobject;
};

struct ThreadData {
    static ThreadData* the()
    {
        thread_local OwnPtr<ThreadData> thread_data = make<ThreadData>();
        if (thread_data)
            return &*thread_data;
        return nullptr;
    }

    ThreadData()
        : master_timer(make<EventLoopMasterTimer>())
        , wake_data(make<EventLoopWake>())
    {
        wake_data->type = CompletionType::Wake;
        wake_data->wait_event.handle = CreateEvent(NULL, FALSE, FALSE, NULL);

        // Consider a way for different event loops to have a different number of threads
        iocp.handle = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 1);
        VERIFY(iocp.handle);

        NTSTATUS status = g_system.NtCreateWaitCompletionPacket(&wake_data->wait_packet.handle, GENERIC_READ | GENERIC_WRITE, NULL);
        VERIFY(NT_SUCCESS(status));
        status = g_system.NtAssociateWaitCompletionPacket(wake_data->wait_packet.handle, iocp.handle, wake_data->wait_event.handle, wake_data.ptr(), NULL, 0, 0, NULL);
        VERIFY(NT_SUCCESS(status));

        master_timer->type = CompletionType::Timer;
        // A high-resolution timer avoids rounding short deadlines up to the default ~16ms timer tick.
        master_timer->timer.handle = CreateWaitableTimerExW(NULL, NULL, CREATE_WAITABLE_TIMER_HIGH_RESOLUTION, TIMER_ALL_ACCESS);
        if (!master_timer->timer.handle)
            master_timer->timer.handle = CreateWaitableTimerExW(NULL, NULL, 0, TIMER_ALL_ACCESS);
        VERIFY(master_timer->timer.handle);
        status = g_system.NtCreateWaitCompletionPacket(&master_timer->wait_packet.handle, GENERIC_READ | GENERIC_WRITE, NULL);
        VERIFY(NT_SUCCESS(status));
    }
    ~ThreadData()
    {
        NTSTATUS status = g_system.NtCancelWaitCompletionPacket(wake_data->wait_packet.handle, TRUE);
        VERIFY(NT_SUCCESS(status));
    }

    OwnHandle iocp;

    // These are only used to register and unregister. The event loop doesn't access these.
    HashMap<Notifier*, NonnullOwnPtr<EventLoopNotifier>> notifiers;

    // Owns the timer allocations; ids stay valid until unregister_timer even after a one-shot timer fires.
    HashMap<intptr_t, NonnullOwnPtr<EventLoopTimer>> timers;
    TimeoutSet timeouts;
    // The deadline the master timer is currently armed for, to skip redundant SetWaitableTimer calls.
    Optional<MonotonicTime> armed_deadline;
    NonnullOwnPtr<EventLoopMasterTimer> master_timer;

    // The wake completion packet is posted to the thread's event loop to wake it.
    NonnullOwnPtr<EventLoopWake> wake_data;
};

static Sync::MutexProtected<HashMap<pid_t, NonnullOwnPtr<EventLoopProcess>>> s_processes;

// Arms (or disarms) the thread's shared waitable timer for the earliest pending deadline.
static void arm_master_timer(ThreadData& thread_data)
{
    auto& master_timer = *thread_data.master_timer;

    auto earliest = thread_data.timeouts.next_timer_expiration();

    if (!earliest.has_value()) {
        if (thread_data.armed_deadline.has_value()) {
            CancelWaitableTimer(master_timer.timer.handle);
            thread_data.armed_deadline = {};
        }
        return;
    }

    if (thread_data.armed_deadline == earliest)
        return;

    LARGE_INTEGER due_time = {};
    // Measured in 0.1μs intervals; negative means relative to now. Round up so the timer doesn't fire
    // before the deadline (and wake the loop with nothing due), and clamp to at least one interval so
    // an already-passed deadline stays a relative time.
    due_time.QuadPart = -AK::max<i64>(((*earliest - MonotonicTime::now()).to_nanoseconds() + 99) / 100, 1);
    BOOL succeeded = SetWaitableTimer(master_timer.timer.handle, &due_time, 0, NULL, NULL, FALSE);
    VERIFY(succeeded);
    thread_data.armed_deadline = earliest;

    if (!master_timer.wait_packet_associated) {
        NTSTATUS status = g_system.NtAssociateWaitCompletionPacket(master_timer.wait_packet.handle, thread_data.iocp.handle, master_timer.timer.handle, &master_timer, NULL, 0, 0, NULL);
        VERIFY(NT_SUCCESS(status));
        master_timer.wait_packet_associated = true;
    }
}

EventLoopImplementationWindows::EventLoopImplementationWindows()
    : m_wake_event(ThreadData::the()->wake_data->wait_event.handle)
{
    VERIFY(m_wake_event);
}

EventLoopImplementationWindows::~EventLoopImplementationWindows()
{
}

int EventLoopImplementationWindows::exec()
{
    for (;;) {
        if (m_exit_requested)
            return m_exit_code;
        pump(PumpMode::WaitForEvents);
    }
    VERIFY_NOT_REACHED();
}

static constexpr bool debug_event_loop = false;

size_t EventLoopImplementationWindows::pump(PumpMode pump_mode)
{
    auto& event_queue = ThreadEventQueue::current();
    auto* thread_data = ThreadData::the();

    // NOTE: The number of entries to dequeue is to be optimized. Ideally we always dequeue all outstanding packets,
    // but we don't want to increase the cost of each pump unnecessarily. If more than one entry is never dequeued
    // at once, we could switch to using GetQueuedCompletionStatus which directly returns the values.
    constexpr ULONG entry_count = 32;
    OVERLAPPED_ENTRY entries[entry_count];
    ULONG entries_removed = 0;

    bool has_pending_events = event_queue.has_pending_events();
    DWORD timeout = 0;
    if (!has_pending_events && pump_mode == PumpMode::WaitForEvents)
        timeout = INFINITE;

    BOOL success = GetQueuedCompletionStatusEx(thread_data->iocp.handle, entries, entry_count, &entries_removed, timeout, FALSE);
    dbgln_if(debug_event_loop, "Event loop dequed {} events", entries_removed);

    if (success) {
        for (ULONG i = 0; i < entries_removed; i++) {
            auto& entry = entries[i];
            auto* packet = reinterpret_cast<CompletionPacket*>(entry.lpCompletionKey);

            if (packet->type == CompletionType::Wake) {
                auto* wake_data = static_cast<EventLoopWake*>(packet);
                NTSTATUS status = g_system.NtAssociateWaitCompletionPacket(wake_data->wait_packet.handle, thread_data->iocp.handle, wake_data->wait_event.handle, wake_data, NULL, 0, 0, NULL);
                VERIFY(NT_SUCCESS(status));
                continue;
            }
            if (packet->type == CompletionType::Timer) {
                auto* master_timer = static_cast<EventLoopMasterTimer*>(packet);
                master_timer->wait_packet_associated = false;
                thread_data->armed_deadline = {};

                // TimeoutSet fires due timers in registration order for equal deadlines. (This cannot rely on
                // kernel wakeup order: simultaneously signaled wait completion packets are delivered in LIFO
                // order.)
                auto now = MonotonicTime::now();
                thread_data->timeouts.fire_expired(now);
                // Zero-interval reloads are rescheduled relative to defer them to the next pump.
                thread_data->timeouts.absolutize_relative_timeouts(now);

                arm_master_timer(*thread_data);
                continue;
            }
            if (packet->type == CompletionType::Notifer) {
                auto* notifier_data = static_cast<EventLoopNotifier*>(packet);
                event_queue.post_event(notifier_data->notifier, Core::Event::Type::NotifierActivation);
                NTSTATUS status = g_system.NtAssociateWaitCompletionPacket(notifier_data->wait_packet.handle, thread_data->iocp.handle, notifier_data->wait_event.handle, notifier_data, NULL, 0, 0, NULL);
                VERIFY(NT_SUCCESS(status));
                continue;
            }
            if (packet->type == CompletionType::Process) {
                auto* process_data = static_cast<EventLoopProcess*>(packet);
                pid_t const process_id = process_data->pid;
                // NOTE: This may seem like the incorrect parameter, but https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-jobobject_associate_completion_port
                // states that this field represents the event type indicator
                DWORD const event_type = entry.dwNumberOfBytesTransferred;
                if (reinterpret_cast<intptr_t>(entry.lpOverlapped) == process_id && (event_type == JOB_OBJECT_MSG_EXIT_PROCESS || event_type == JOB_OBJECT_MSG_ABNORMAL_EXIT_PROCESS)) {
                    Optional<NonnullOwnPtr<EventLoopProcess>> owned_process = s_processes.with_locked([&](auto& processes) {
                        return processes.take(process_id);
                    });
                    if (owned_process.has_value())
                        owned_process.release_value()->exit_handler(process_id);
                }
                continue;
            }
            VERIFY_NOT_REACHED();
        }
    } else {
        DWORD error = GetLastError();
        switch (error) {
        case WAIT_TIMEOUT:
            break;
        default:
            dbgln("GetQueuedCompletionStatusEx failed with unexpected error: {}", Error::from_windows_error(error));
            VERIFY_NOT_REACHED();
        }
    }

    return event_queue.process();
}

void EventLoopImplementationWindows::quit(int code)
{
    m_exit_requested = true;
    m_exit_code = code;
}

void EventLoopImplementationWindows::wake()
{
    SetEvent(m_wake_event);
}

static int notifier_type_to_network_event(NotificationType type)
{
    switch (type) {
    case NotificationType::Read:
        return FD_READ | FD_CLOSE | FD_ACCEPT;
    case NotificationType::Write:
        return FD_WRITE;
    default:
        dbgln("This notification type is not implemented: {}", (int)type);
        VERIFY_NOT_REACHED();
    }
}

void EventLoopManagerWindows::register_notifier(Notifier& notifier)
{
    auto* thread_data = ThreadData::the();
    auto& notifiers = thread_data->notifiers;

    if (notifiers.contains(&notifier))
        return;

    HANDLE event = CreateEvent(NULL, FALSE, FALSE, NULL);
    VERIFY(event);
    int rc = WSAEventSelect(notifier.fd(), event, notifier_type_to_network_event(notifier.type()));
    VERIFY(!rc);

    auto notifier_data = make<EventLoopNotifier>();
    notifier_data->type = CompletionType::Notifer;
    notifier_data->notifier = &notifier;
    notifier_data->wait_event.handle = event;
    NTSTATUS status = g_system.NtCreateWaitCompletionPacket(&notifier_data->wait_packet.handle, GENERIC_READ | GENERIC_WRITE, NULL);
    VERIFY(NT_SUCCESS(status));
    status = g_system.NtAssociateWaitCompletionPacket(notifier_data->wait_packet.handle, thread_data->iocp.handle, event, notifier_data.ptr(), NULL, 0, 0, NULL);
    VERIFY(NT_SUCCESS(status));
    notifiers.set(&notifier, move(notifier_data));
}

void EventLoopManagerWindows::unregister_notifier(Notifier& notifier)
{
    auto* thread_data = ThreadData::the();
    VERIFY(thread_data);

    auto& notifiers = thread_data->notifiers;
    auto maybe_notifier_data = notifiers.take(&notifier);
    if (!maybe_notifier_data.has_value())
        return;
    auto notifier_data = move(maybe_notifier_data.value());
    // We are removing the signalled packets since the caller no longer expects them
    NTSTATUS status = g_system.NtCancelWaitCompletionPacket(notifier_data->wait_packet.handle, TRUE);
    VERIFY(NT_SUCCESS(status));
    // TODO: Reuse the data structure
}

intptr_t EventLoopManagerWindows::register_timer(EventReceiver& object, int milliseconds, bool should_reload)
{
    VERIFY(milliseconds >= 0);
    auto* thread_data = ThreadData::the();
    VERIFY(thread_data);

    auto timer = make<EventLoopTimer>();
    timer->owner = object.make_weak_ptr();
    timer->interval = AK::Duration::from_milliseconds(milliseconds);
    timer->should_reload = should_reload;
    timer->reload(MonotonicTime::now());
    thread_data->timeouts.schedule_absolute(timer.ptr());

    // Ids are process-wide unique so a stale id (e.g. one unregistered from the wrong thread) can never
    // match another thread's timer.
    static Atomic<intptr_t> next_timer_id { 1 };
    auto timer_id = next_timer_id.fetch_add(1, AK::MemoryOrder::memory_order_relaxed);
    thread_data->timers.set(timer_id, move(timer));
    arm_master_timer(*thread_data);
    return timer_id;
}

void EventLoopManagerWindows::unregister_timer(intptr_t timer_id)
{
    if (auto* thread_data = ThreadData::the()) {
        auto timer = thread_data->timers.take(timer_id);
        if (!timer.has_value())
            return;
        if ((*timer)->is_scheduled())
            thread_data->timeouts.unschedule(timer->ptr());
        arm_master_timer(*thread_data);
    }
}

int EventLoopManagerWindows::register_signal([[maybe_unused]] int signal_number, [[maybe_unused]] Function<void(int)> handler)
{
    dbgln("Core::EventLoopManagerWindows::register_signal() is not implemented");
    VERIFY_NOT_REACHED();
}

void EventLoopManagerWindows::unregister_signal([[maybe_unused]] int handler_id)
{
    dbgln("Core::EventLoopManagerWindows::unregister_signal() is not implemented");
    VERIFY_NOT_REACHED();
}

void EventLoopManagerWindows::register_process(pid_t pid, ESCAPING Function<void(pid_t)> exit_handler)
{
    auto* thread_data = ThreadData::the();
    VERIFY(thread_data);

    s_processes.with_locked([&](auto& processes) {
        if (processes.contains(pid))
            return;

        HANDLE process_handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
        VERIFY(process_handle);

        HANDLE job_object_handle = CreateJobObject(nullptr, nullptr);
        VERIFY(job_object_handle);

        BOOL succeeded = AssignProcessToJobObject(job_object_handle, process_handle);
        VERIFY(succeeded);

        auto process_data = make<EventLoopProcess>();
        process_data->type = CompletionType::Process;
        process_data->process.handle = process_handle;
        process_data->pid = pid;
        process_data->exit_handler = move(exit_handler);
        process_data->jobobject.handle = job_object_handle;

        JOBOBJECT_ASSOCIATE_COMPLETION_PORT joacp = { .CompletionKey = process_data.ptr(), .CompletionPort = thread_data->iocp.handle };
        succeeded = SetInformationJobObject(job_object_handle, JobObjectAssociateCompletionPortInformation, &joacp, sizeof(JOBOBJECT_ASSOCIATE_COMPLETION_PORT));
        VERIFY(succeeded);

        processes.set(pid, move(process_data));
    });
}

void EventLoopManagerWindows::unregister_process(pid_t pid)
{
    auto maybe_process = s_processes.with_locked([&](auto& processes) {
        return processes.take(pid);
    });
    if (!maybe_process.has_value())
        return;

    auto process_data = maybe_process.release_value();
    JOBOBJECT_ASSOCIATE_COMPLETION_PORT joacp = { .CompletionKey = process_data, .CompletionPort = nullptr };
    BOOL succeeded = SetInformationJobObject(process_data->jobobject.handle, JobObjectAssociateCompletionPortInformation, &joacp, sizeof(JOBOBJECT_ASSOCIATE_COMPLETION_PORT));
    VERIFY(succeeded);
}

void EventLoopManagerWindows::did_post_event()
{
}

NonnullOwnPtr<EventLoopImplementation> EventLoopManagerWindows::make_implementation()
{
    return make<EventLoopImplementationWindows>();
}

}