/*
    * Timer.cpp
    * Periodic tick, timekeeping and idle-loop servicing
    * Copyright (c) 2025 Daniel Hammer
*/

#include "Timer.hpp"
#include <atomic>
#include <Arch/Cpu.hpp>
#include <Arch/CpuPower.hpp>
#include <Arch/Interrupts.hpp>
#include <Arch/Smp.hpp>
#include <Arch/Timer.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Sched/Scheduler.hpp>
#include <Drivers/Net/E1000.hpp>
#include <Drivers/Net/E1000E.hpp>
#include <Drivers/Net/Genet.hpp>
#include <Drivers/Net/Wifi/Wifi.hpp>
#include <Drivers/Audio/IntelHda.hpp>
#include <Drivers/USB/Xhci.hpp>
#include <Drivers/Bluetooth/Bluetooth.hpp>
#include <Drivers/Serial/BtUart.hpp>
#include <Drivers/USB/HidKeyboard.hpp>

using namespace Kt;

namespace Timekeeping {
    // The BSP keeps a 1 ms tick for timekeeping and sleep deadlines.
    // Running APs use a 10 ms scheduler tick. Idle APs mask it entirely and
    // rely on reschedule IPIs, avoiding periodic package wakeups.
    static constexpr uint32_t BSP_TICK_INTERVAL_MS = 1;
    static constexpr uint32_t BSP_TIMER_HZ = 1000 / BSP_TICK_INTERVAL_MS;
    static constexpr uint32_t AP_TICK_INTERVAL_MS = 10;

    // Global state
    static std::atomic<uint64_t> g_tickCount{0};

    // Set once Arch::Timer is calibrated and the BSP tick is running.
    static bool g_timerReady = false;

    static std::atomic<bool> g_schedEnabled{false};

    // Timer IRQ handler: BSP runs timekeeping + scheduler accounting; APs
    // run scheduler accounting only.
    static void TimerHandler(uint8_t, bool interruptedUser) {
        auto* cpu = Smp::GetCurrentCpuData();
        uint32_t schedElapsedMs = (cpu->cpuIndex == 0) ? BSP_TICK_INTERVAL_MS : AP_TICK_INTERVAL_MS;

        if (cpu->cpuIndex == 0) {
            g_tickCount.fetch_add(schedElapsedMs, std::memory_order_relaxed);

            if (Drivers::Net::E1000E::RequiresPolling()) {
                Drivers::Net::E1000E::Poll();
            }
            Drivers::USB::HidKeyboard::Tick();

            // Bottom halves normally run from idle, but a continuously busy
            // desktop may never enter idle. Force one bounded idle-context
            // pass for USB immediately, and for network RX at most every
            // 10 ms (important for E1000E's polling fallback).
            // UART Bluetooth (Raspberry Pi) has no xHCI behind it: its
            // received bytes wait in the UART ring for the same pass.
            bool usbWork = Drivers::USB::Xhci::HasDeferredWork() ||
                           Drivers::Serial::BtUart::HasRxData();
            bool audioWork = Drivers::Audio::IntelHda::HasDeferredWork();
            bool netWork = Drivers::Net::E1000::HasDeferredWork() ||
                           Drivers::Net::E1000E::HasDeferredWork() ||
                           Drivers::Net::Genet::HasDeferredWork();
            // A forced trip into the idle bottom-half context is safe only when
            // this timer interrupted ring 3. If it interrupted a syscall, that
            // suspended kernel stack may own one of the process-context mutexes
            // used below; the reserved idle context would then spin on the lock
            // while refusing to schedule its owner, deadlocking the CPU.
            if (interruptedUser &&
                (usbWork || audioWork ||
                 (netWork && (g_tickCount.load(std::memory_order_relaxed) % 10 == 0)))) {
                Sched::RequestKernelWork();
            }
        }

        if (g_schedEnabled) {
            Sched::Tick(schedElapsedMs, interruptedUser);
        }
    }

    void TimerInitialize() {
        // Register before the hardware starts so the first tick is counted.
        Arch::RegisterIrqHandler(Arch::IRQ_TIMER, TimerHandler);

        if (!Arch::Timer::InitializeBsp(BSP_TICK_INTERVAL_MS)) {
            KernelLogStream(ERROR, "Timer") << "No usable tick timer";
            return;
        }
        g_timerReady = true;

        KernelLogStream(OK, "Timer") << "Tick started: BSP " << base::dec
            << (uint64_t)BSP_TIMER_HZ << " Hz periodic";
    }

    void TimerReinitialize() {
        if (!g_timerReady) {
            KernelLogStream(ERROR, "Timer") << "Cannot reinit timer - not calibrated";
            return;
        }

        // Reprogram the timer hardware (its state was lost during S3).
        // The calibration and the IRQ handler registration survive in RAM.
        Arch::Timer::StartPeriodic(BSP_TICK_INTERVAL_MS);

        KernelLogStream(OK, "Timer") << "Timer restarted after S3 resume";
    }

    uint64_t GetTicks() {
        return g_tickCount.load(std::memory_order_relaxed);
    }

    uint64_t GetMilliseconds() {
        return g_tickCount.load(std::memory_order_relaxed);  // 1 tick = 1 ms at 1000 Hz
    }

    void EnableSchedulerTick() {
        g_schedEnabled.store(true, std::memory_order_release);
    }

    void TimerInitializeAP() {
        // Use the BSP's calibration directly: all cores share the same
        // timer clock, and this avoids calibration contention during AP boot.
        if (!g_timerReady) return;

        // Configure the 10 ms scheduler timer for running APs. Their idle loop
        // masks it after initialization and rearms it when dispatching work.
        Arch::Timer::StartPeriodic(AP_TICK_INTERVAL_MS);
    }

    void TimerEnterApIdle() {
        auto* cpu = Smp::GetCurrentCpuData();
        if (cpu == nullptr || cpu->cpuIndex == 0 || !g_timerReady) return;

        Arch::Timer::Stop();
    }

    void TimerLeaveApIdle() {
        auto* cpu = Smp::GetCurrentCpuData();
        if (cpu == nullptr || cpu->cpuIndex == 0 || !g_timerReady) return;

        // Reprogram the period as well as unmasking. A deep idle state
        // may have stopped the local timer at an arbitrary point in its period.
        Arch::Timer::StartPeriodic(AP_TICK_INTERVAL_MS);
    }

    void ServiceDeferredWork() {
        auto* cpu = Smp::GetCurrentCpuData();
        if (cpu == nullptr || cpu->currentSlot >= 0) return;
        cpu->kernelWorkPending = false;

        // These bottom halves deliberately run outside a process because they
        // enter subsystems with process-context mutexes. Do not let a timer or
        // reschedule IPI switch this idle stack to a process while such a lock
        // remains held: the process could immediately wait on that same lock,
        // leaving the suspended idle owner with no opportunity to release it.
        // The scheduler and reschedule-IPI path both honor this reservation.
        bool wasReserved = cpu->reservedForKernelWork;
        cpu->reservedForKernelWork = true;

        // Drain USB work only when the MSI path has actually queued something.
        // Bluetooth shares this controller, so service its protocol queues in
        // the same pass after xHCI has delivered completion callbacks.
        bool usbWork = Drivers::USB::Xhci::HasDeferredWork();
        if (usbWork) {
            Drivers::USB::Xhci::ProcessDeferredWork();
        }

        // NIC hard IRQs only acknowledge/mask and queue RX work. Dispatching
        // Ethernet/TCP/UDP here keeps process-context IPC mutexes out of IRQs.
        if (Drivers::Net::E1000::HasDeferredWork())
            Drivers::Net::E1000::ProcessDeferredWork();
        if (Drivers::Net::E1000E::HasDeferredWork())
            Drivers::Net::E1000E::ProcessDeferredWork();
        if (Drivers::Net::Genet::HasDeferredWork())
            Drivers::Net::Genet::ProcessDeferredWork();

        // HDA completion IRQs only acknowledge/mask. Resampling and DMA-ring
        // refill are far too expensive for hard interrupt context.
        if (cpu->cpuIndex == 0 && Drivers::Audio::IntelHda::HasDeferredWork()) {
            Drivers::Audio::IntelHda::ProcessDeferredWork();
        }

        // Complete any boot-deferred Bluetooth bring-up (Intel firmware
        // download) here instead of on the boot path: the download takes
        // seconds and used to stall kmain before the first process spawned.
        // Cheap no-op unless an adapter is waiting; self-claiming, and safe
        // to preempt (the scheduler saves/resumes the idle context).
        bool bluetoothWork = usbWork ||
            Drivers::Bluetooth::HasDeferredWork();
        if (bluetoothWork)
            Drivers::Bluetooth::ServiceDeferredInit();

        // Service Bluetooth inbound traffic (the headset's SDP/AVRCP queries,
        // AVDTP commands, ACL flow-control credits) whenever a core idles.
        // Without this, BT events are only processed while some syscall happens
        // to be pumping a wait loop -- a peer that queries us between audio
        // writes got silence (observed: Bose re-dialing SDP during playback,
        // queries never answered).  Self-serializing and a cheap no-op when
        // the adapter is down.
        if (bluetoothWork)
            Drivers::Bluetooth::ServiceEvents();

        // Wi-Fi mirrors the Bluetooth split: the firmware load needs the
        // ramdisk, and the RX/notification ring must be drained outside hard
        // interrupt context (the MSI handler only latches a flag).
        if (Drivers::Net::Wifi::HasDeferredWork()) {
            Drivers::Net::Wifi::ServiceDeferredInit();
            Drivers::Net::Wifi::ServiceEvents();
        }

        // Thermal policy records transitions during BSP maintenance; print
        // them from this explicitly non-interrupt idle path.
        Arch::CpuPower::ServiceDeferredDiagnostics();

        cpu->reservedForKernelWork = wasReserved;
    }

    void IdleOnce() {
        auto* cpu = Smp::GetCurrentCpuData();

        // HaltWithInterruptsDisabled() deliberately returns with IF clear so
        // the next readiness check can lead into an atomic enable+halt.  Deferred
        // work cannot inherit that state: xHCI/Bluetooth service contains
        // wall-clock waits and depends on the BSP timer and device MSIs.  If
        // it runs with IF=0, GetMilliseconds() never advances on the BSP and
        // the completion interrupt cannot arrive, producing a permanent idle-
        // loop freeze.  Keep IRQs enabled for all active idle-context work and
        // disable them again only immediately before halting below.
        Arch::EnableInterrupts();

        ServiceDeferredWork();

        if (cpu == nullptr || cpu->cpuIndex != 0 || !g_schedEnabled || !g_timerReady) {
            // Non-BSP or pre-scheduler fallback: plain HLT keeps the LAPIC
            // timer running in C1 on hardware where MWAIT promotes to deeper
            // states that freeze the timer (no ARAT in practice).
            Arch::DisableInterrupts();
            Arch::HaltWithInterruptsDisabled();
            return;
        }

        Sched::RunBspMaintenance();

        if (Sched::HasReadyProcesses()) {
            Sched::Schedule();
            return;
        }

        // We deliberately DO NOT reprogram the LAPIC timer to a one-shot here.
        // Some laptop CPUs let MWAIT/HLT drop into a deep C-state where the
        // LAPIC timer stops firing, even though CPUID.06H:EAX[ARAT] claims
        // it doesn't. Symptom: the BSP enters idle and never wakes from its
        // timer; only external IRQs (PS/2, USB MSI) bring it out, and even
        // those don't advance g_tickCount because the one-shot's current
        // count is frozen. By leaving the boot-time periodic 1ms timer in
        // place, TimerHandler runs every tick and g_tickCount keeps moving
        // -- which is what RunBspMaintenance needs to wake sleeping procs.
        //
        // HLT (vs MWAIT) is the second half: MWAIT EAX=0 is a *hint* and
        // the CPU may go deeper than C1; HLT generally stays in C1, which
        // is the deepest state guaranteed to keep the LAPIC timer alive.
        Arch::DisableInterrupts();
        Arch::HaltWithInterruptsDisabled();
    }

    void BusyWaitUs(uint64_t us) {
        for (uint64_t i = 0; i < us; i++)
            Arch::IoDelay();
    }

    void BusyWaitMs(uint64_t ms) {
        if (ms == 0) return;

        bool ticking = g_timerReady && Arch::InterruptsEnabled();

        // Either clock alone is long enough, so stop at whichever says so
        // first. The tick is the accurate one; the IoDelay spins bound the wait
        // when the tick is not running yet (or is masked), so this can never
        // hang on a timer that does not fire. One extra tick of slack because
        // the first may already be nearly over.
        uint64_t start = g_tickCount.load(std::memory_order_relaxed);
        uint64_t delayBudget = (ms + 1) * 1000;
        for (uint64_t spins = 0; spins < delayBudget; spins++) {
            if (ticking &&
                g_tickCount.load(std::memory_order_relaxed) - start >= ms + 1)
                return;
            Arch::IoDelay();
        }
    }

    void Sleep(uint64_t ms) {
        if (ms == 0) return;
        // Use BlockForSleep to properly deschedule the process instead
        // of busy-waiting with Schedule(). This frees the CPU for real
        // work and avoids unnecessary scheduling overhead.
        Sched::BlockForSleep(ms);
    }
};