/*
    * 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/SmpBoot.hpp>
#include <Io/IoPort.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Sched/Scheduler.hpp>
#include <Drivers/Net/E1000E.hpp>
#include <Drivers/USB/Xhci.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.
    // APs use a coarser 10 ms scheduler tick to avoid waking idle cores
    // 1000 times per second with no useful work to do.
    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;
    // Keep a bounded periodic wake while otherwise tickless so wall-clock
    // accounting and typematic repeat still advance even on a fully idle BSP.
    static constexpr uint32_t BSP_IDLE_MAX_INTERVAL_MS = 50;
    static constexpr uint32_t BSP_INTERACTIVE_IDLE_MAX_INTERVAL_MS = 1;

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

    static bool g_schedEnabled = false;
    static volatile bool g_bspIdleOneShotArmed = false;
    static uint32_t g_bspIdleOneShotMs = 0;
    static uint32_t g_bspIdleInitialCount = 0;
    static std::atomic<uint64_t> g_interactiveUntilMs{0};

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

    static void ProgramBspIdleOneShotTimer(uint32_t intervalMs) {
        g_bspIdleOneShotMs = intervalMs;
        g_bspIdleInitialCount = CountForIntervalMs(intervalMs);
        g_bspIdleOneShotArmed = true;

        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_DIVIDE, DIVIDE_BY_16);
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_LVT,
            (Hal::IRQ_VECTOR_BASE + Hal::IRQ_TIMER));
        Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_INITIAL, g_bspIdleInitialCount);
    }

    // Timer IRQ handler: BSP handles timekeeping and the few timer-driven
    // fallbacks that are still required; APs only run scheduler accounting.
    static void TimerHandler(uint8_t) {
        auto* cpu = Smp::GetCurrentCpuData();
        uint32_t schedElapsedMs = (cpu->cpuIndex == 0) ? BSP_TICK_INTERVAL_MS : AP_TICK_INTERVAL_MS;

        if (cpu->cpuIndex == 0) {
            if (g_bspIdleOneShotArmed) {
                schedElapsedMs = g_bspIdleOneShotMs;
                g_bspIdleOneShotArmed = false;
                ProgramBspPeriodicTimer();
            }

            g_tickCount.fetch_add(schedElapsedMs, std::memory_order_relaxed);

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

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

    // 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).
        g_bspIdleOneShotArmed = false;
        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 NoteInteractiveActivity(uint32_t durationMs) {
        uint64_t until = GetMilliseconds() + durationMs;
        uint64_t current = g_interactiveUntilMs.load(std::memory_order_relaxed);

        while (until > current &&
               !g_interactiveUntilMs.compare_exchange_weak(
                   current, until, std::memory_order_relaxed, std::memory_order_relaxed)) {
        }
    }

    static bool IsInteractiveActivityActive() {
        return GetMilliseconds() < g_interactiveUntilMs.load(std::memory_order_relaxed);
    }

    void EnableSchedulerTick() {
        g_schedEnabled = true;
    }

    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 a coarser periodic timer on APs. The scheduler still gets
        // a 10 ms time slice, but idle APs stop taking 1000 timer interrupts/sec.
        ProgramTimer(true, AP_TICK_INTERVAL_MS);
    }

    void IdleOnce(bool hasMwait, volatile uint64_t* monitorAddr) {
        auto* cpu = Smp::GetCurrentCpuData();
        if (cpu == nullptr || cpu->cpuIndex != 0 || !g_schedEnabled || g_ticksPerMs == 0) {
            Hal::CpuIdle::Wait(BSP_TICK_INTERVAL_MS, hasMwait, monitorAddr);
            return;
        }

        Sched::RunBspMaintenance();

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

        bool interactive = IsInteractiveActivityActive();
        uint32_t waitMs = interactive
            ? BSP_INTERACTIVE_IDLE_MAX_INTERVAL_MS
            : BSP_IDLE_MAX_INTERVAL_MS;

        uint64_t now = GetTicks();
        uint64_t nextDeadline = Sched::GetNextDeadlineTick();
        uint64_t keyboardDeadline = Drivers::USB::HidKeyboard::GetNextTickDeadline();
        if (keyboardDeadline != 0 &&
            (nextDeadline == 0 || keyboardDeadline < nextDeadline)) {
            nextDeadline = keyboardDeadline;
        }

        if (nextDeadline != 0) {
            if (nextDeadline <= now) {
                waitMs = BSP_TICK_INTERVAL_MS;
            } else {
                uint64_t untilDeadline = nextDeadline - now;
                if (untilDeadline < waitMs) {
                    waitMs = (uint32_t)untilDeadline;
                }
            }
        }

        if (waitMs == 0) {
            waitMs = BSP_TICK_INTERVAL_MS;
        }

        asm volatile("cli" ::: "memory");
        ProgramBspIdleOneShotTimer(waitMs);
        if (interactive) {
            Hal::CpuIdle::WaitShallowWithInterruptsDisabled(hasMwait, monitorAddr);
        } else {
            Hal::CpuIdle::WaitWithInterruptsDisabled(waitMs, hasMwait, monitorAddr);
        }

        if (g_bspIdleOneShotArmed) {
            uint32_t currentCount = Hal::LocalApic::ReadRegister(Hal::LocalApic::REG_TIMER_CURRENT);
            uint32_t elapsedTicks = (g_bspIdleInitialCount > currentCount)
                ? (g_bspIdleInitialCount - currentCount)
                : 0;
            uint32_t elapsedMs = (g_ticksPerMs > 0) ? (elapsedTicks / g_ticksPerMs) : 0;
            if (elapsedMs > 0) {
                g_tickCount.fetch_add(elapsedMs, std::memory_order_relaxed);
            }

            g_bspIdleOneShotArmed = false;
            Hal::LocalApic::WriteRegister(Hal::LocalApic::REG_TIMER_INITIAL, 0);
            ProgramBspPeriodicTimer();
        }
        asm volatile("sti" ::: "memory");
    }

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