/*
    * Keyboard.hpp
    * SYS_ISKEYAVAILABLE, SYS_GETKEY, SYS_GETCHAR syscalls
    * Copyright (c) 2026 Daniel Hammer
*/

#pragma once
#include <Sched/Scheduler.hpp>
#include <Drivers/PS2/Keyboard.hpp>
#include <Libraries/Memory.hpp>

#include "Common.hpp"

namespace montauk::abi {
    // ---- Key text ----------------------------------------------------------
    // Keyboard layouts are applied in userspace (desktop, terminal), which send
    // KeyEvents carrying the Unicode codepoint. Binaries built before that send
    // 9-byte KeyEventV1s whose `ascii` is a Windows-1252 byte; both kinds share
    // a redirected child's key mailbox, told apart by message length.

    static uint32_t KeyCp1252ToUnicode(uint8_t b) {
        static constexpr uint16_t high[32] = {
            0x20AC, 0x81,   0x201A, 0x0192, 0x201E, 0x2026, 0x2020, 0x2021,
            0x02C6, 0x2030, 0x0160, 0x2039, 0x0152, 0x8D,   0x017D, 0x8F,
            0x90,   0x2018, 0x2019, 0x201C, 0x201D, 0x2022, 0x2013, 0x2014,
            0x02DC, 0x2122, 0x0161, 0x203A, 0x0153, 0x9D,   0x017E, 0x0178,
        };
        return (b >= 0x80 && b <= 0x9F) ? high[b - 0x80] : b;
    }

    // The character a key event types, or 0 for control/navigation keys.
    static uint32_t KeyTextCodepoint(const KeyEvent& ev) {
        if (ev.codepoint) return ev.codepoint;
        uint8_t b = (uint8_t)ev.ascii;
        if (b < 0x20 || b == 0x7F) return 0;
        return KeyCp1252ToUnicode(b);
    }

    static int KeyEncodeUtf8(uint32_t cp, uint8_t* out) {
        if (cp > 0x10FFFF || (cp >= 0xD800 && cp <= 0xDFFF)) cp = 0xFFFD;
        if (cp < 0x80) { out[0] = (uint8_t)cp; return 1; }
        if (cp < 0x800) {
            out[0] = (uint8_t)(0xC0 | (cp >> 6));
            out[1] = (uint8_t)(0x80 | (cp & 0x3F));
            return 2;
        }
        if (cp < 0x10000) {
            out[0] = (uint8_t)(0xE0 | (cp >> 12));
            out[1] = (uint8_t)(0x80 | ((cp >> 6) & 0x3F));
            out[2] = (uint8_t)(0x80 | (cp & 0x3F));
            return 3;
        }
        out[0] = (uint8_t)(0xF0 | (cp >> 18));
        out[1] = (uint8_t)(0x80 | ((cp >> 12) & 0x3F));
        out[2] = (uint8_t)(0x80 | ((cp >> 6) & 0x3F));
        out[3] = (uint8_t)(0x80 | (cp & 0x3F));
        return 4;
    }

    static void KeyFromDriver(const Drivers::PS2::Keyboard::KeyEvent& k, KeyEvent* out) {
        memset(out, 0, sizeof(KeyEvent));
        out->scancode = k.Scancode;
        out->ascii    = k.Ascii;
        out->pressed  = k.Pressed;
        out->shift    = k.Shift;
        out->ctrl     = k.Ctrl;
        out->alt      = k.Alt;
        out->capslock = k.CapsLock;
        out->altgr    = k.AltGr;
        out->extended = k.Extended;
        // The kernel's own scancode table is US ASCII.
        uint8_t a = (uint8_t)k.Ascii;
        out->codepoint = (a >= 0x20 && a < 0x7F) ? a : 0;
    }

    // Receive one key from a redirected child's mailbox as a full KeyEvent,
    // upgrading legacy 9-byte messages. Returns MailboxRecv's result.
    static int KeyMailboxRecv(Ipc::Mailbox* mailbox, KeyEvent* out) {
        memset(out, 0, sizeof(KeyEvent));
        uint16_t len = sizeof(KeyEvent);
        int rc = Ipc::MailboxRecv(mailbox, nullptr, out, &len, true);
        if (rc > 0 && len < sizeof(KeyEvent)) {
            memset((uint8_t*)out + len, 0, sizeof(KeyEvent) - len);
            out->codepoint = KeyTextCodepoint(*out);
        }
        return rc;
    }

    // Turn a pressed key into the next getchar byte. Characters above ASCII are
    // returned as UTF-8, one byte per call; the rest is parked in the process.
    static bool KeyTakeChar(Sched::Process* proc, const KeyEvent& ev, char* out) {
        if (!ev.pressed) return false;
        uint32_t cp = KeyTextCodepoint(ev);
        if (cp >= 0x80 && proc != nullptr) {
            int n = KeyEncodeUtf8(cp, proc->getcharUtf8);
            proc->getcharUtf8Len = (uint8_t)n;
            proc->getcharUtf8Pos = 1;
            *out = (char)proc->getcharUtf8[0];
            return true;
        }
        if (ev.ascii != 0 && ev.ascii != KEY_ASCII_SUBSTITUTE) {
            *out = ev.ascii;
            return true;
        }
        return false;
    }

    static bool KeyPendingChar(Sched::Process* proc, char* out) {
        if (proc == nullptr || proc->getcharUtf8Pos >= proc->getcharUtf8Len) return false;
        *out = (char)proc->getcharUtf8[proc->getcharUtf8Pos++];
        return true;
    }

    // ---- Syscalls ------------------------------------------------------------

    static bool Sys_IsKeyAvailable() {
        auto* proc = Sched::GetCurrentProcessPtr();
        if (proc && proc->redirected) {
            Ipc::HandleSnapshot inputSnapshot;
            Ipc::HandleSnapshot mailboxSnapshot;
            Ipc::Stream* input = GetRedirInStream(proc, inputSnapshot);
            Ipc::Mailbox* mailbox = GetRedirKeyMailbox(proc, mailboxSnapshot);
            if (input != nullptr || mailbox != nullptr) {
                if (input != nullptr && Ipc::StreamHasData(input)) return true;
                return mailbox != nullptr && Ipc::MailboxHasMessage(mailbox);
            }
        }
        return Drivers::PS2::Keyboard::IsKeyAvailable();
    }

    // Blocking read of the next key event, in the full format.
    static void ReadKeyEvent(KeyEvent* outEvent) {
        auto* proc = Sched::GetCurrentProcessPtr();
        if (proc && proc->redirected) {
            Ipc::HandleSnapshot snapshot;
            Ipc::Mailbox* mailbox = GetRedirKeyMailbox(proc, snapshot);
            if (mailbox != nullptr) {
                for (;;) {
                    uint64_t observedWake = Sched::ObserveObjectWake(mailbox);
                    int rc = KeyMailboxRecv(mailbox, outEvent);
                    if (rc > 0) return;
                    if (rc < 0) {
                        memset(outEvent, 0, sizeof(KeyEvent));
                        return;
                    }
                    Sched::BlockOnObjectSince(mailbox, 0, observedWake);
                }
            }
        }
        KeyFromDriver(Drivers::PS2::Keyboard::GetKey(), outEvent);
    }

    // SYS_GETKEY: legacy 9-byte event. A new-format event's first 9 bytes are
    // already a valid KeyEventV1 (ascii holds the 1252 byte or SUB).
    static void Sys_GetKey(KeyEventV1* outEvent) {
        if (outEvent == nullptr) return;
        KeyEvent ev;
        ReadKeyEvent(&ev);
        memcpy(outEvent, &ev, sizeof(KeyEventV1));
    }

    // SYS_GETKEY_EX: size-versioned. The caller's size was validated by the
    // dispatcher (>= sizeof(KeyEventV1), writable).
    static int Sys_GetKeyEx(void* outEvent, uint64_t size) {
        KeyEvent ev;
        ReadKeyEvent(&ev);
        uint64_t n = size < sizeof(KeyEvent) ? size : sizeof(KeyEvent);
        memcpy(outEvent, &ev, n);
        if (size > n) memset((uint8_t*)outEvent + n, 0, size - n);
        return Success;
    }

    static char Sys_GetChar() {
        auto* proc = Sched::GetCurrentProcessPtr();
        char pending;
        if (KeyPendingChar(proc, &pending)) return pending;
        if (proc && proc->redirected) {
            Ipc::HandleSnapshot inputSnapshot;
            Ipc::HandleSnapshot mailboxSnapshot;
            Ipc::Stream* input = GetRedirInStream(proc, inputSnapshot);
            Ipc::Mailbox* mailbox = GetRedirKeyMailbox(proc, mailboxSnapshot);
            if (input != nullptr || mailbox != nullptr) {
                for (;;) {
                    uint64_t inputWake = input ? Sched::ObserveObjectWake(input) : 0;
                    uint64_t mailboxWake = mailbox ? Sched::ObserveObjectWake(mailbox) : 0;
                    if (input != nullptr) {
                        uint8_t c = 0;
                        int rc = Ipc::StreamRead(input, &c, 1, true);
                        if (rc > 0) return (char)c;
                        if (rc < 0 && mailbox == nullptr) return 0;
                    }

                    if (mailbox != nullptr) {
                        KeyEvent ev;
                        int rc = KeyMailboxRecv(mailbox, &ev);
                        if (rc > 0) {
                            char c;
                            if (KeyTakeChar(proc, ev, &c)) return c;
                            continue;
                        }
                        if (rc < 0 && input == nullptr) return 0;
                    }

                    if (proc->ioWaitsetHandle >= 0) {
                        Ipc::WaitsetReady ready{};
                        if (Ipc::WaitsetWaitHandle(proc->ioWaitsetHandle, &ready, ~0ULL) < 0) {
                            return 0;
                        }
                    } else if (input != nullptr) {
                        Sched::BlockOnObjectSince(input, 0, inputWake);
                    } else if (mailbox != nullptr) {
                        Sched::BlockOnObjectSince(mailbox, 0, mailboxWake);
                    } else {
                        return 0;
                    }
                }
            }
        }
        return Drivers::PS2::Keyboard::GetChar();
    }

    /*
     * Non-blocking counterpart of Sys_GetChar: returns the next character, or
     * 0 when none is pending. IsKeyAvailable is not enough to make GetChar
     * safe to call -- a queued key-release or modifier event reports as
     * "available" but carries no ascii, and GetChar would go on to block. So
     * the drain has to happen on this side of the syscall boundary.
     */
    static char Sys_GetCharNb() {
        auto* proc = Sched::GetCurrentProcessPtr();
        char pending;
        if (KeyPendingChar(proc, &pending)) return pending;
        if (proc && proc->redirected) {
            Ipc::HandleSnapshot inputSnapshot;
            Ipc::HandleSnapshot mailboxSnapshot;
            Ipc::Stream* input = GetRedirInStream(proc, inputSnapshot);
            Ipc::Mailbox* mailbox = GetRedirKeyMailbox(proc, mailboxSnapshot);
            if (input != nullptr || mailbox != nullptr) {
                if (input != nullptr) {
                    uint8_t c = 0;
                    if (Ipc::StreamRead(input, &c, 1, true) > 0) return (char)c;
                }

                while (mailbox != nullptr) {
                    KeyEvent ev;
                    if (KeyMailboxRecv(mailbox, &ev) <= 0) break;
                    char c;
                    if (KeyTakeChar(proc, ev, &c)) return c;
                }
                return 0;
            }
        }

        while (Drivers::PS2::Keyboard::IsKeyAvailable()) {
            auto k = Drivers::PS2::Keyboard::GetKey();
            if (k.Pressed && k.Ascii != 0) return k.Ascii;
        }
        return 0;
    }
};