#pragma once

#include <Sched/Scheduler.hpp>
#include <Ipc/Ipc.hpp>

namespace montauk::abi {

    static Sched::Process* GetRedirTarget(Sched::Process* proc) {
        if (!proc || !proc->redirected) return nullptr;
        if (proc->ioOutHandle >= 0 || proc->ioInHandle >= 0 || proc->ioKeyHandle >= 0) {
            return proc;
        }
        return (proc->parentPid >= 0) ? Sched::GetProcessByPid(proc->parentPid) : nullptr;
    }

    static bool ResolveProcessHandle(Sched::Process* proc, int handle, Ipc::HandleType expectedType,
                                     Ipc::HandleSnapshot& snapshot) {
        if (proc == nullptr || handle < 0) return false;

        int slot = Ipc::SlotForPid(proc->pid);
        if (slot < 0) return false;

        if (!snapshot.Capture(slot, handle)) return false;
        return snapshot.type == expectedType && snapshot.object != nullptr;
    }

    static Ipc::Stream* GetRedirOutStream(Sched::Process* proc, Ipc::HandleSnapshot& snapshot) {
        proc = GetRedirTarget(proc);
        if (proc == nullptr) return nullptr;

        if (!ResolveProcessHandle(proc, proc->ioOutHandle, Ipc::HandleType::Stream, snapshot)) return nullptr;
        return (Ipc::Stream*)snapshot.object;
    }

    static Ipc::Stream* GetRedirInStream(Sched::Process* proc, Ipc::HandleSnapshot& snapshot) {
        proc = GetRedirTarget(proc);
        if (proc == nullptr) return nullptr;

        if (!ResolveProcessHandle(proc, proc->ioInHandle, Ipc::HandleType::Stream, snapshot)) return nullptr;
        return (Ipc::Stream*)snapshot.object;
    }

    static Ipc::Mailbox* GetRedirKeyMailbox(Sched::Process* proc, Ipc::HandleSnapshot& snapshot) {
        proc = GetRedirTarget(proc);
        if (proc == nullptr) return nullptr;

        if (!ResolveProcessHandle(proc, proc->ioKeyHandle, Ipc::HandleType::Mailbox, snapshot)) return nullptr;
        return (Ipc::Mailbox*)snapshot.object;
    }

    static int DuplicateHandleBetweenSlots(int srcSlot, int handle, int dstSlot) {
        if (srcSlot < 0 || dstSlot < 0 || handle < 0) return -1;

        Ipc::HandleSnapshot snapshot;
        if (!snapshot.Capture(srcSlot, handle)) return -1;
        if ((snapshot.rights & Ipc::RightDup) == 0) return -1;
        return Ipc::InstallHandleForSlot(dstSlot, snapshot.object, snapshot.type, snapshot.rights);
    }

    static bool ConfigureRedirWaitsetForSlot(int slot, Sched::Process* proc) {
        if (slot < 0 || proc == nullptr) return false;

        int waitset = Ipc::CreateWaitsetHandleForSlot(slot);
        if (waitset < 0) return false;

        if (proc->ioInHandle >= 0) {
            if (Ipc::WaitsetAddHandleForSlot(slot, waitset, proc->ioInHandle,
                                             Ipc::SignalReadable | Ipc::SignalPeerClosed) < 0) {
                Ipc::CloseHandleForSlot(slot, waitset);
                return false;
            }
        }

        if (proc->ioKeyHandle >= 0) {
            if (Ipc::WaitsetAddHandleForSlot(slot, waitset, proc->ioKeyHandle,
                                             Ipc::SignalReadable | Ipc::SignalPeerClosed) < 0) {
                Ipc::CloseHandleForSlot(slot, waitset);
                return false;
            }
        }

        proc->ioWaitsetHandle = waitset;
        return true;
    }

    static int WriteAllToStream(Ipc::Stream* stream, const uint8_t* data, int len) {
        if (stream == nullptr || data == nullptr || len <= 0) return -1;

        int total = 0;
        while (total < len) {
            uint64_t observedWake = Sched::ObserveObjectWake(stream);
            int written = Ipc::StreamWrite(stream, data + total, len - total, false);
            if (written < 0) {
                return (total > 0) ? total : -1;
            }
            if (written == 0) {
                Sched::BlockOnObjectSince(stream, 0, observedWake);
                continue;
            }
            total += written;
        }
        return total;
    }
}