/*
    * CpuPower.cpp
    * Intel CPU power management: HWP frequency scaling, C1E promotion,
    * and the package thermal governor.
    *
    * Without this module the CPU runs at whatever P-state the firmware left
    * it in, forever, and idle cores never drop below C1. On a modern laptop
    * part (tested target: i9-13900HX, 24 cores) that means constant high
    * package power, spinning fans, and eventually the EC's thermal
    * protection (forced beeping / hard power-off on Clevo boards).
    *
    * Three mechanisms, all feature-gated by CPUID so virtual machines and
    * non-Intel parts skip them safely:
    *
    *  1. HWP (Hardware P-states, CPUID.06H:EAX[7]): the CPU autonomously
    *     picks its frequency between lowest and highest performance level,
    *     biased by the Energy/Performance Preference byte. This is what
    *     intel_pstate does on Linux and is the single biggest thermal win.
    *
    *  2. C1E promotion (MSR_POWER_CTL[1]): lets HLT idle drop core voltage
    *     and frequency to minimum instead of parking at the current ratio.
    *
    *  3. Thermal governor: reads the package digital thermal sensor every
    *     500 ms on the BSP and steps the HWP performance ceiling down when
    *     the package crosses PassiveTempC, hard-clamps toward lowest above
    *     HotTempC, and slowly recovers below ResumeTempC. This keeps the
    *     package away from TjMax (hardware PROCHOT) and far away from the
    *     EC's emergency responses.
    *
    * Copyright (c) 2026 Daniel Hammer
*/

#include "CpuPower.hpp"
#include "MSR.hpp"
#include "SmpBoot.hpp"
#include <Terminal/Terminal.hpp>
#include <atomic>

using namespace Kt;

namespace Hal {
    namespace CpuPower {

        // ============================================================
        // MSR addresses (power/thermal; used only in this module)
        // ============================================================

        static constexpr uint32_t IA32_MPERF                 = 0x0E7;
        static constexpr uint32_t IA32_APERF                 = 0x0E8;
        static constexpr uint32_t IA32_THERM_STATUS          = 0x19C;
        static constexpr uint32_t MSR_TEMPERATURE_TARGET     = 0x1A2;
        static constexpr uint32_t IA32_ENERGY_PERF_BIAS      = 0x1B0;
        static constexpr uint32_t IA32_PACKAGE_THERM_STATUS  = 0x1B1;
        static constexpr uint32_t MSR_POWER_CTL              = 0x1FC;
        static constexpr uint32_t IA32_PM_ENABLE             = 0x770;
        static constexpr uint32_t IA32_HWP_CAPABILITIES      = 0x771;
        static constexpr uint32_t IA32_HWP_REQUEST           = 0x774;

        // ============================================================
        // Governor tuning
        // ============================================================

        // Balanced energy/performance preference (0 = max performance,
        // 255 = max power saving). 0x80 matches Linux "balance_performance".
        static constexpr uint8_t  BalancedEpp       = 0x80;

        static constexpr uint8_t  PassiveTempC      = 85;   // start stepping down
        static constexpr uint8_t  HotTempC          = 94;   // step down aggressively
        static constexpr uint8_t  ResumeTempC       = 78;   // start stepping back up
        static constexpr uint32_t GovernorIntervalMs = 500;
        static constexpr uint8_t  PassiveStepDown   = 2;    // ratio units (~200 MHz)
        static constexpr uint8_t  HotStepDown       = 8;    // ratio units (~800 MHz)

        // ============================================================
        // Detected features and shared policy state
        // ============================================================

        struct Features {
            bool Intel = false;
            bool Dts = false;         // CPUID.06H:EAX[0] digital thermal sensor
            bool Ptm = false;         // CPUID.06H:EAX[6] package thermal sensor
            bool Hwp = false;         // CPUID.06H:EAX[7] hardware P-states
            bool HwpEpp = false;      // CPUID.06H:EAX[10] HWP energy/perf pref
            bool Epb = false;         // CPUID.06H:ECX[3] energy/perf bias MSR
            bool AperfMperf = false;  // CPUID.06H:ECX[0] freq feedback counters
            uint32_t MaxLeaf = 0;
        };

        static Features g_feat{};
        static bool g_hwpActive = false;
        static bool g_thermalReady = false;
        static uint8_t g_highestPerf = 0;
        static uint8_t g_lowestPerf = 0;
        static uint8_t g_epp = 0;
        static uint8_t g_tjMax = 100;
        static uint32_t g_baseMHz = 0;
        static uint32_t g_maxMHz = 0;

        // Governor output: the current performance ceiling plus an epoch
        // counter. Each CPU compares the epoch against the last one it
        // applied and rewrites its own IA32_HWP_REQUEST when they differ
        // (MSR writes only take effect on the core that executes them).
        static std::atomic<uint8_t>  g_curMaxPerf{0};
        static std::atomic<uint32_t> g_policyEpoch{0};
        static uint32_t g_appliedEpoch[Smp::MaxCPUs]{};

        static bool g_throttling = false;
        static std::atomic<uint8_t> g_lastTempC{0};
        static std::atomic<uint32_t> g_effMHz{0};
        static uint64_t g_lastGovernorMs = 0;
        static uint64_t g_lastAperf = 0;
        static uint64_t g_lastMperf = 0;

        // ============================================================
        // Helpers
        // ============================================================

        static void Cpuid(uint32_t leaf, uint32_t subleaf,
                          uint32_t& a, uint32_t& b, uint32_t& c, uint32_t& d) {
            asm volatile("cpuid"
                         : "=a"(a), "=b"(b), "=c"(c), "=d"(d)
                         : "a"(leaf), "c"(subleaf));
        }

        static void DetectFeatures() {
            uint32_t a, b, c, d;
            Cpuid(0, 0, a, b, c, d);
            g_feat.MaxLeaf = a;
            g_feat.Intel = (b == 0x756E6547 && d == 0x49656E69 && c == 0x6C65746E);
            if (!g_feat.Intel || g_feat.MaxLeaf < 6) {
                return;
            }

            Cpuid(6, 0, a, b, c, d);
            g_feat.Dts        = (a & (1u << 0)) != 0;
            g_feat.Ptm        = (a & (1u << 6)) != 0;
            g_feat.Hwp        = (a & (1u << 7)) != 0;
            g_feat.HwpEpp     = (a & (1u << 10)) != 0;
            g_feat.AperfMperf = (c & (1u << 0)) != 0;
            g_feat.Epb        = (c & (1u << 3)) != 0;

            if (g_feat.MaxLeaf >= 0x16) {
                Cpuid(0x16, 0, a, b, c, d);
                g_baseMHz = a & 0xFFFF;
                g_maxMHz = b & 0xFFFF;
            }
        }

        static uint64_t ComposeHwpRequest(uint8_t minPerf, uint8_t maxPerf, uint8_t epp) {
            // Desired = 0 selects fully autonomous frequency selection.
            return (uint64_t)minPerf
                 | ((uint64_t)maxPerf << 8)
                 | ((uint64_t)epp << 24);
        }

        // MSR_POWER_CTL is not architectural, so gate it on the digital
        // thermal sensor bit: hypervisors (KVM, TCG) mask the CPUID.06H
        // power bits, so DTS present means real Intel hardware where the
        // MSR exists.
        static void EnableC1ePromotionOnThisCore() {
            if (!g_feat.Dts) return;
            uint64_t v = ReadMSR(MSR_POWER_CTL);
            if ((v & (1ull << 1)) == 0) {
                WriteMSR(MSR_POWER_CTL, v | (1ull << 1));
            }
        }

        static uint8_t ReadPackageTempC() {
            if (g_feat.Ptm) {
                uint64_t s = ReadMSR(IA32_PACKAGE_THERM_STATUS);
                uint32_t readout = (uint32_t)((s >> 16) & 0x7F);
                return (readout <= g_tjMax) ? (uint8_t)(g_tjMax - readout) : 0;
            }
            if (g_feat.Dts) {
                uint64_t s = ReadMSR(IA32_THERM_STATUS);
                if (s & (1ull << 31)) { // reading valid
                    uint32_t readout = (uint32_t)((s >> 16) & 0x7F);
                    return (readout <= g_tjMax) ? (uint8_t)(g_tjMax - readout) : 0;
                }
            }
            return 0;
        }

        static void UpdateEffectiveFrequency() {
            if (!g_feat.AperfMperf || g_baseMHz == 0) return;

            uint64_t aperf = ReadMSR(IA32_APERF);
            uint64_t mperf = ReadMSR(IA32_MPERF);
            uint64_t dA = aperf - g_lastAperf;
            uint64_t dM = mperf - g_lastMperf;
            bool primed = (g_lastAperf | g_lastMperf) != 0;
            g_lastAperf = aperf;
            g_lastMperf = mperf;

            if (primed && dM != 0) {
                g_effMHz.store((uint32_t)((uint64_t)g_baseMHz * dA / dM),
                               std::memory_order_relaxed);
            }
        }

        static void ApplyHwpRequestOnThisCpu(uint32_t epoch) {
            auto* cpu = Smp::GetCurrentCpuData();
            if (cpu == nullptr || cpu->cpuIndex < 0 || cpu->cpuIndex >= Smp::MaxCPUs) {
                return;
            }
            g_appliedEpoch[cpu->cpuIndex] = epoch;
            WriteMSR(IA32_HWP_REQUEST,
                     ComposeHwpRequest(g_lowestPerf,
                                       g_curMaxPerf.load(std::memory_order_relaxed),
                                       g_epp));
        }

        // ============================================================
        // Public interface
        // ============================================================

        void InitializeBsp() {
            DetectFeatures();
            if (!g_feat.Intel) {
                KernelLogStream(INFO, "CpuPower")
                    << "Not an Intel CPU - power management left to firmware defaults";
                return;
            }

            if (g_feat.Dts || g_feat.Ptm) {
                uint64_t tt = ReadMSR(MSR_TEMPERATURE_TARGET);
                uint8_t tj = (uint8_t)((tt >> 16) & 0xFF);
                if (tj >= 50 && tj <= 115) {
                    g_tjMax = tj;
                }
                g_thermalReady = true;
            }

            EnableC1ePromotionOnThisCore();

            if (g_feat.Hwp) {
                // One-way until reset; harmless if firmware already set it.
                WriteMSR(IA32_PM_ENABLE, 1);
                if (ReadMSR(IA32_PM_ENABLE) & 1) {
                    uint64_t caps = ReadMSR(IA32_HWP_CAPABILITIES);
                    g_highestPerf = (uint8_t)(caps & 0xFF);
                    g_lowestPerf  = (uint8_t)((caps >> 24) & 0xFF);
                    // The EPP byte is reserved when CPUID does not
                    // advertise it; 0 (max performance bias) is the only
                    // safe value there.
                    g_epp = g_feat.HwpEpp ? BalancedEpp : 0;

                    g_curMaxPerf.store(g_highestPerf, std::memory_order_relaxed);
                    uint32_t epoch = g_policyEpoch.fetch_add(1,
                                         std::memory_order_acq_rel) + 1;
                    g_hwpActive = true;
                    ApplyHwpRequestOnThisCpu(epoch);

                    KernelLogStream(OK, "CpuPower") << "HWP enabled: perf range "
                        << base::dec << (uint64_t)g_lowestPerf << "-"
                        << (uint64_t)g_highestPerf << ", EPP=0x"
                        << base::hex << (uint64_t)g_epp;
                } else {
                    KernelLogStream(WARNING, "CpuPower")
                        << "HWP advertised but IA32_PM_ENABLE did not stick";
                }
            } else if (g_feat.Epb) {
                // Legacy knob: bias hardware decisions toward balance.
                WriteMSR(IA32_ENERGY_PERF_BIAS, 6);
                KernelLogStream(OK, "CpuPower") << "No HWP - set EPB to balanced";
            } else if (!g_thermalReady) {
                KernelLogStream(INFO, "CpuPower")
                    << "No power features exposed (hypervisor or old CPU)";
            }

            if (g_thermalReady) {
                UpdateEffectiveFrequency(); // prime APERF/MPERF baseline
                g_lastTempC.store(ReadPackageTempC(), std::memory_order_relaxed);
                KernelLogStream(OK, "CpuPower") << "Thermal governor armed: TjMax="
                    << base::dec << (uint64_t)g_tjMax << "C, passive="
                    << (uint64_t)PassiveTempC << "C, hot=" << (uint64_t)HotTempC
                    << "C, package now " << (uint64_t)g_lastTempC.load() << "C";
            }
        }

        void InitializeAp() {
            if (!g_feat.Intel) return;

            EnableC1ePromotionOnThisCore();

            if (g_hwpActive) {
                // IA32_PM_ENABLE is package-scoped on most parts, but the
                // SDM allows logical scope; setting it again is harmless.
                WriteMSR(IA32_PM_ENABLE, 1);
                ApplyHwpRequestOnThisCpu(g_policyEpoch.load(std::memory_order_acquire));
            }
        }

        void ApplyPolicyIfChanged() {
            if (!g_hwpActive) return;
            auto* cpu = Smp::GetCurrentCpuData();
            if (cpu == nullptr || cpu->cpuIndex < 0 || cpu->cpuIndex >= Smp::MaxCPUs) {
                return;
            }
            uint32_t epoch = g_policyEpoch.load(std::memory_order_acquire);
            if (g_appliedEpoch[cpu->cpuIndex] == epoch) return;
            // Also re-check C1E promotion here: POWER_CTL is per-core and
            // loses its contents across S3, and ReapplyAfterWake only runs
            // on the BSP. The epoch bump it does routes every AP through
            // this path.
            EnableC1ePromotionOnThisCore();
            ApplyHwpRequestOnThisCpu(epoch);
        }

        void ThermalTick(uint64_t nowMs) {
            if (!g_thermalReady) return;
            if (nowMs - g_lastGovernorMs < GovernorIntervalMs) return;
            g_lastGovernorMs = nowMs;

            UpdateEffectiveFrequency();

            uint8_t temp = ReadPackageTempC();
            g_lastTempC.store(temp, std::memory_order_relaxed);
            if (temp == 0 || !g_hwpActive) return;

            uint8_t cur = g_curMaxPerf.load(std::memory_order_relaxed);
            uint8_t next = cur;

            if (temp >= HotTempC) {
                next = (cur > (uint8_t)(g_lowestPerf + HotStepDown))
                     ? (uint8_t)(cur - HotStepDown) : g_lowestPerf;
            } else if (temp >= PassiveTempC) {
                next = (cur > (uint8_t)(g_lowestPerf + PassiveStepDown))
                     ? (uint8_t)(cur - PassiveStepDown) : g_lowestPerf;
            } else if (temp <= ResumeTempC && cur < g_highestPerf) {
                next = (uint8_t)(cur + 1);
            }

            if (next == cur) return;

            bool wasThrottling = g_throttling;
            g_throttling = (next < g_highestPerf);
            g_curMaxPerf.store(next, std::memory_order_relaxed);
            uint32_t epoch = g_policyEpoch.fetch_add(1, std::memory_order_acq_rel) + 1;
            ApplyHwpRequestOnThisCpu(epoch); // BSP applies immediately; APs on tick

            if (g_throttling && !wasThrottling) {
                KernelLogStream(WARNING, "CpuPower") << "Package at " << base::dec
                    << (uint64_t)temp << "C - thermal throttle engaged";
            } else if (!g_throttling && wasThrottling) {
                KernelLogStream(OK, "CpuPower") << "Package cooled to " << base::dec
                    << (uint64_t)temp << "C - thermal throttle released";
            }
        }

        void ReapplyAfterWake() {
            if (!g_feat.Intel) return;

            EnableC1ePromotionOnThisCore();

            if (g_hwpActive) {
                WriteMSR(IA32_PM_ENABLE, 1);
                // Force every CPU (including this one) to re-apply the HWP
                // request: the MSRs lost their contents across S3.
                uint32_t epoch = g_policyEpoch.fetch_add(1, std::memory_order_acq_rel) + 1;
                ApplyHwpRequestOnThisCpu(epoch);
            }

            // Reset the frequency-feedback baseline (counters restarted).
            g_lastAperf = 0;
            g_lastMperf = 0;
        }

        bool GetSnapshot(Snapshot& out) {
            out = {};
            if (!g_feat.Intel || (!g_thermalReady && !g_hwpActive)) {
                return false;
            }

            out.HwpActive   = g_hwpActive;
            out.Throttling  = g_throttling;
            out.TempC       = g_lastTempC.load(std::memory_order_relaxed);
            out.TjMaxC      = g_tjMax;
            out.HighestPerf = g_highestPerf;
            out.LowestPerf  = g_lowestPerf;
            out.CurMaxPerf  = g_curMaxPerf.load(std::memory_order_relaxed);
            out.Epp         = g_epp;
            out.BaseMHz     = g_baseMHz;
            out.MaxMHz      = g_maxMHz;
            out.EffMHz      = g_effMHz.load(std::memory_order_relaxed);
            return true;
        }

    };
};
