/*
    * ApicTimer.cpp
    * Local APIC timer: PIT-calibrated periodic tick for timekeeping
    * Copyright (c) 2025 Daniel Hammer
*/

#include "ApicTimer.hpp"
#include <atomic>
#include <ACPI/CpuIdle.hpp>
#include <Hal/Apic/Apic.hpp>
#include <Hal/Apic/Interrupts.hpp>
#include <Hal/Cpu.hpp>
#include <Hal/CpuPower.hpp>
#include <Hal/SmpBoot.hpp>
#include <Io/IoPort.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/Wifi/Wifi.hpp>
#include <Drivers/Audio/IntelHda.hpp>
#include <Drivers/USB/Xhci.hpp>
#include <Drivers/USB/Bluetooth/Bluetooth.hpp>
#include <Drivers/USB/HidKeyboard.hpp>

using namespace Kt;

namespace Timekeeping {
    // PIT constants
    static constexpr uint32_t PIT_FREQUENCY = 1193182;     // PIT oscillator frequency in Hz
    static constexpr uint16_t PIT_CHANNEL2_DATA = 0x42;
    static constexpr uint16_t PIT_COMMAND = 0x43;
    static constexpr uint16_t PIT_GATE_PORT = 0x61;

    // APIC timer LVT mode bits
    static constexpr uint32_t LVT_MASKED   = (1 << 16);
    static constexpr uint32_t LVT_PERIODIC = (1 << 17);

    // APIC timer divide configuration values
    static constexpr uint32_t DIVIDE_BY_16 = 0x03;

    // 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};
    static uint32_t g_ticksPerMs = 0;

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

    static uint32_t CountForIntervalMs(uint32_t intervalMs) {
        return g_ticksPerMs * intervalMs;
    }

    static void ProgramTimer(bool periodic, uint32_t intervalMs) {
        uint32_t lvt = (Hal::IRQ_VECTOR_BASE + Hal::IRQ_TIMER);
        if (periodic) {
            lvt |= LVT_PERIODIC;
        }

        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_DIVIDE, DIVIDE_BY_16);
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_LVT, lvt);
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_INITIAL,
            CountForIntervalMs(intervalMs));
    }

    static void ProgramBspPeriodicTimer() {
        ProgramTimer(true, BSP_TICK_INTERVAL_MS);
    }

    // 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).
            bool usbWork = Drivers::USB::Xhci::HasDeferredWork();
            bool audioWork = Drivers::Audio::IntelHda::HasDeferredWork();
            bool netWork = Drivers::Net::E1000::HasDeferredWork() ||
                           Drivers::Net::E1000E::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);
        }
    }

    // Use PIT channel 2 to create a precise delay for calibration.
    // Returns the number of APIC timer ticks that elapsed during ~10ms.
    static uint32_t CalibratePit() {
        // We want to measure ~10ms worth of APIC timer ticks.
        // PIT count for 10ms: PIT_FREQUENCY / 100 = 11931
        uint16_t pitCount = PIT_FREQUENCY / 100;

        // Set up PIT channel 2 in one-shot mode (mode 0), lobyte/hibyte
        // Command: channel 2 (bits 7-6 = 10), lobyte/hibyte (bits 5-4 = 11),
        //          mode 0 (bits 3-1 = 000), binary (bit 0 = 0) => 0xB0
        Io::Out8(0xB0, PIT_COMMAND);

        // Disable PIT channel 2 gate to prepare
        uint8_t gate = Io::In8(PIT_GATE_PORT);
        gate &= ~0x01;   // Disable gate (bit 0 = 0)
        gate &= ~0x02;   // Disable speaker (bit 1 = 0)
        Io::Out8(gate, PIT_GATE_PORT);

        // Load the count value (low byte first, then high byte)
        Io::Out8((uint8_t)(pitCount & 0xFF), PIT_CHANNEL2_DATA);
        Io::IoPortWait();
        Io::Out8((uint8_t)(pitCount >> 8), PIT_CHANNEL2_DATA);

        // Set APIC timer to max count with divide-by-16, one-shot (masked so no IRQ)
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_DIVIDE, DIVIDE_BY_16);
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_LVT, LVT_MASKED);
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_INITIAL, 0xFFFFFFFF);

        // Enable PIT channel 2 gate to start counting
        gate = Io::In8(PIT_GATE_PORT);
        gate |= 0x01;    // Enable gate (bit 0 = 1)
        Io::Out8(gate, PIT_GATE_PORT);

        // Wait for PIT channel 2 output to go high (bit 5 of port 0x61)
        while (!(Io::In8(PIT_GATE_PORT) & 0x20)) {
            asm volatile("pause");
        }

        // Read how many APIC timer ticks elapsed
        uint32_t currentCount = Hal::LocalApic::ReadRegister(Hal::LocalApic::REG_TIMER_CURRENT);
        uint32_t elapsed = 0xFFFFFFFF - currentCount;

        // Stop the APIC timer
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_INITIAL, 0);

        return elapsed;
    }

    void ApicTimerInitialize() {
        KernelLogStream(INFO, "Timer") << "Calibrating APIC timer using PIT channel 2";

        // Calibrate: measure APIC ticks over ~10ms
        uint32_t ticksIn10ms = CalibratePit();
        g_ticksPerMs = ticksIn10ms / 10;

        if (g_ticksPerMs == 0) {
            KernelLogStream(ERROR, "Timer") << "APIC timer calibration failed (0 ticks/ms)";
            return;
        }

        uint64_t timerFreqHz = (uint64_t)g_ticksPerMs * 1000;

        KernelLogStream(OK, "Timer") << "APIC timer: " << base::dec << (uint64_t)g_ticksPerMs
            << " ticks/ms (" << timerFreqHz << " Hz, divide-by-16)";

        // Register IRQ handler for timer (IRQ 0 = vector 32)
        Hal::RegisterIrqHandler(Hal::IRQ_TIMER, TimerHandler);

        // Configure APIC timer: periodic mode, vector 32
        ProgramBspPeriodicTimer();

        KernelLogStream(OK, "Timer") << "APIC timer started: BSP " << base::dec
            << (uint64_t)BSP_TIMER_HZ << " Hz periodic, initial count="
            << (uint64_t)CountForIntervalMs(BSP_TICK_INTERVAL_MS);
    }

    void ApicTimerReinitialize() {
        if (g_ticksPerMs == 0) {
            KernelLogStream(ERROR, "Timer") << "Cannot reinit APIC timer - not calibrated";
            return;
        }

        // Reprogram the APIC timer registers (they were lost during S3).
        // The calibrated g_ticksPerMs value is still valid (it's in RAM).
        // The IRQ handler registration also survives (it's a function pointer array in RAM).
        ProgramBspPeriodicTimer();

        KernelLogStream(OK, "Timer") << "APIC 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 ApicTimerInitializeAP() {
        // Use the BSP's calibrated ticks-per-ms value directly.
        // All cores share the same bus clock, so the APIC timer rate is
        // identical. This avoids PIT contention during AP boot.
        if (g_ticksPerMs == 0) return;

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

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

        uint32_t lvt = Hal::LocalApic::ReadRegister(Hal::LocalApic::REG_TIMER_LVT);
        if ((lvt & LVT_MASKED) == 0) {
            Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_LVT,
                                          lvt | LVT_MASKED);
        }
    }

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

        // Reprogram the initial count as well as unmasking. A deep idle state
        // may have stopped the local timer at an arbitrary point in its period.
        ProgramTimer(true, 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();

        // 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::USB::Bluetooth::HasDeferredWork();
        if (bluetoothWork)
            Drivers::USB::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::USB::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.
        Hal::CpuPower::ServiceDeferredDiagnostics();

        cpu->reservedForKernelWork = wasReserved;
    }

    void IdleOnce(bool hasMwait, volatile uint64_t* monitorAddr) {
        (void)hasMwait;
        (void)monitorAddr;
        auto* cpu = Smp::GetCurrentCpuData();

        // HaltWithInterruptsDisabled() deliberately returns with IF clear so
        // the next readiness check can lead into an atomic STI+HLT.  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 entering HLT below.
        asm volatile("sti" ::: "memory");

        ServiceDeferredWork();

        if (cpu == nullptr || cpu->cpuIndex != 0 || !g_schedEnabled || g_ticksPerMs == 0) {
            // 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).
            asm volatile("cli" ::: "memory");
            Hal::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.
        asm volatile("cli" ::: "memory");
        Hal::HaltWithInterruptsDisabled();
    }

    void BusyWaitUs(uint64_t us) {
        // An I/O write to port 0x80 is an ISA-bus cycle of roughly a
        // microsecond on every chipset since the 486, and never shorter.
        for (uint64_t i = 0; i < us; i++)
            Io::IoPortWait();
    }

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

        uint64_t flags;
        asm volatile("pushfq; pop %0" : "=r"(flags));
        bool ticking = g_ticksPerMs != 0 && (flags & (1 << 9));

        // Either clock alone is long enough, so stop at whichever says so
        // first. The tick is the accurate one; the port writes 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 portBudget = (ms + 1) * 1000;
        for (uint64_t spins = 0; spins < portBudget; spins++) {
            if (ticking &&
                g_tickCount.load(std::memory_order_relaxed) - start >= ms + 1)
                return;
            Io::IoPortWait();
        }
    }

    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);
    }
};