/*
    * CpuIdle.cpp
    * ACPI-guided CPU idle-state selection
    * Copyright (c) 2026 Daniel Hammer
*/

#include "CpuIdle.hpp"
#include <ACPI/FADT.hpp>
#include <ACPI/AML/AmlInterpreter.hpp>
#include <Hal/Cpu.hpp>
#include <Io/IoPort.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>

using namespace Kt;

namespace Hal {
    namespace CpuIdle {

        enum class IdleEntryKind : uint8_t {
            FixedHardware,
            IoPort,
        };

        struct IdleState {
            bool     Valid = false;
            uint8_t  CType = 0;
            uint32_t LatencyUs = 0;
            uint32_t PowerMw = 0;
            IdleEntryKind EntryKind = IdleEntryKind::FixedHardware;
            uint16_t IoPort = 0;
            uint8_t  WidthBits = 0;
        };

        static constexpr int MaxIdleStates = 8;
        static constexpr uint8_t GAS_SYSTEM_IO = 0x01;
        static constexpr uint8_t GAS_FIXED_HARDWARE = 0x7F;
        static constexpr uint32_t ResidencyLatencyRatio = 8;
        static constexpr uint32_t IoStateMinPredictedMs = 20;
        static constexpr uint16_t LegacyC2MaxLatencyUs = 100;

        static IdleState g_idleStates[MaxIdleStates]{};
        static int g_idleStateCount = 0;
        static volatile uint64_t g_fallbackMonitor = 0;

        static uint64_t ReadLe64(const uint8_t* data) {
            uint64_t value = 0;
            for (int i = 0; i < 8; i++) {
                value |= (uint64_t)data[i] << (i * 8);
            }
            return value;
        }

        static bool IsName(const char* actual, const char* expected) {
            for (int i = 0; i < 4; i++) {
                if (actual[i] != expected[i]) {
                    return false;
                }
            }
            return true;
        }

        static uint32_t DecodePkgLength(const uint8_t* aml, uint32_t length, uint32_t& pos) {
            if (pos >= length) return 0;

            uint8_t lead = aml[pos];
            uint32_t byteCount = (lead >> 6) & 0x03;

            if (byteCount == 0) {
                pos++;
                return lead & 0x3F;
            }

            uint32_t value = lead & 0x0F;
            pos++;

            for (uint32_t i = 0; i < byteCount && pos < length; i++) {
                value |= (uint32_t)aml[pos++] << (4 + 8 * i);
            }

            return value;
        }

        static bool DecodeInteger(const uint8_t* aml, uint32_t length, uint32_t& pos, uint64_t& value) {
            if (pos >= length) return false;

            uint8_t op = aml[pos];
            switch (op) {
                case AML::ZeroOp:
                    value = 0;
                    pos++;
                    return true;
                case AML::OneOp:
                    value = 1;
                    pos++;
                    return true;
                case AML::BytePrefix:
                    if (pos + 1 >= length) return false;
                    value = aml[pos + 1];
                    pos += 2;
                    return true;
                case AML::WordPrefix:
                    if (pos + 2 >= length) return false;
                    value = (uint16_t)aml[pos + 1] | ((uint16_t)aml[pos + 2] << 8);
                    pos += 3;
                    return true;
                case AML::DWordPrefix:
                    if (pos + 4 >= length) return false;
                    value = (uint32_t)aml[pos + 1]
                          | ((uint32_t)aml[pos + 2] << 8)
                          | ((uint32_t)aml[pos + 3] << 16)
                          | ((uint32_t)aml[pos + 4] << 24);
                    pos += 5;
                    return true;
                case AML::QWordPrefix:
                    if (pos + 8 >= length) return false;
                    value = ReadLe64(&aml[pos + 1]);
                    pos += 9;
                    return true;
                default:
                    return false;
            }
        }

        static bool DecodeBufferObject(const uint8_t* aml, uint32_t length, uint32_t& pos,
                                       const uint8_t*& data, uint32_t& dataLength) {
            if (pos >= length || aml[pos] != AML::BufferOp) return false;

            pos++;
            uint32_t pkgStart = pos;
            uint32_t pkgLen = DecodePkgLength(aml, length, pos);
            uint32_t end = pkgStart + pkgLen;
            if (end > length) end = length;

            uint64_t declaredLength = 0;
            if (!DecodeInteger(aml, end, pos, declaredLength)) {
                return false;
            }

            data = &aml[pos];
            dataLength = (uint32_t)(end - pos);
            if (declaredLength < dataLength) {
                dataLength = (uint32_t)declaredLength;
            }

            pos = end;
            return true;
        }

        static bool DecodePackageObject(const uint8_t* aml, uint32_t length, uint32_t& pos,
                                        const uint8_t*& data, uint32_t& dataLength) {
            if (pos >= length) return false;
            if (aml[pos] != AML::PackageOp && aml[pos] != AML::VarPackageOp) return false;

            pos++;
            uint32_t pkgStart = pos;
            uint32_t pkgLen = DecodePkgLength(aml, length, pos);
            uint32_t end = pkgStart + pkgLen;
            if (end > length) end = length;

            data = &aml[pos];
            dataLength = (uint32_t)(end - pos);
            pos = end;
            return true;
        }

        static const char* EntryKindName(const IdleState& state) {
            return (state.EntryKind == IdleEntryKind::IoPort) ? "ioport" : "fixed";
        }

        static bool ParseGas(const uint8_t* buffer, uint32_t length, IdleState& state) {
            if (length < 12) {
                return false;
            }

            uint8_t addressSpace = buffer[0];
            uint8_t bitWidth = buffer[1];
            uint8_t accessSize = buffer[3];
            uint64_t address = ReadLe64(&buffer[4]);

            if (addressSpace == GAS_FIXED_HARDWARE) {
                state.EntryKind = IdleEntryKind::FixedHardware;
                return true;
            }

            if (addressSpace != GAS_SYSTEM_IO || address == 0 || address > 0xFFFF) {
                return false;
            }

            uint8_t widthBits = bitWidth;
            if (widthBits == 0) {
                switch (accessSize) {
                    case 1: widthBits = 8; break;
                    case 2: widthBits = 16; break;
                    case 3: widthBits = 32; break;
                    default: widthBits = 8; break;
                }
            }

            if (widthBits != 8 && widthBits != 16 && widthBits != 32) {
                return false;
            }

            state.EntryKind = IdleEntryKind::IoPort;
            state.IoPort = (uint16_t)address;
            state.WidthBits = widthBits;
            return true;
        }

        static bool ParseIdleStatePackage(const uint8_t* aml, uint32_t length, IdleState& state) {
            if (length < 2) return false;

            uint32_t pos = 1; // skip package element count byte
            const uint8_t* regData = nullptr;
            uint32_t regLength = 0;
            uint64_t type = 0;
            uint64_t latency = 0;
            uint64_t power = 0;

            if (!DecodeBufferObject(aml, length, pos, regData, regLength)) {
                return false;
            }
            if (!DecodeInteger(aml, length, pos, type)) {
                return false;
            }
            if (!DecodeInteger(aml, length, pos, latency)) {
                return false;
            }
            if (!DecodeInteger(aml, length, pos, power)) {
                return false;
            }

            state = {};
            state.CType = (uint8_t)type;
            state.LatencyUs = (uint32_t)latency;
            state.PowerMw = (uint32_t)power;

            if (!ParseGas(regData, regLength, state)) {
                return false;
            }

            state.Valid = true;
            return true;
        }

        static bool ParseCstObject(const char* path, const AML::Object& obj) {
            if (obj.Type != AML::ObjectType::Package || obj.Buffer.Length < 2) {
                return false;
            }

            uint32_t pos = 1; // package element count byte
            uint64_t declaredStateCount = 0;
            if (!DecodeInteger(obj.Buffer.Data, obj.Buffer.Length, pos, declaredStateCount)) {
                return false;
            }

            int parsedCount = 0;
            for (uint64_t i = 0; i < declaredStateCount && parsedCount < MaxIdleStates; i++) {
                const uint8_t* nestedData = nullptr;
                uint32_t nestedLength = 0;
                if (!DecodePackageObject(obj.Buffer.Data, obj.Buffer.Length, pos, nestedData, nestedLength)) {
                    break;
                }

                IdleState state{};
                if (!ParseIdleStatePackage(nestedData, nestedLength, state)) {
                    continue;
                }

                g_idleStates[parsedCount++] = state;
            }

            if (parsedCount == 0) {
                return false;
            }

            g_idleStateCount = parsedCount;

            KernelLogStream(OK, "CpuIdle") << "Loaded " << base::dec
                << (uint64_t)g_idleStateCount << " ACPI idle state(s) from " << path;

            for (int i = 0; i < g_idleStateCount; i++) {
                const auto& state = g_idleStates[i];
                KernelLogStream(INFO, "CpuIdle") << "C" << base::dec << (uint64_t)state.CType
                    << " latency=" << (uint64_t)state.LatencyUs << "us"
                    << " power=" << (uint64_t)state.PowerMw << "mW"
                    << " entry=" << EntryKindName(state);
            }

            return true;
        }

        static bool LoadLegacyPblkIdleStates(const FADT::ParsedFADT& fadt) {
            if (!fadt.Valid ||
                fadt.WorstC2Latency == 0 ||
                fadt.WorstC2Latency > LegacyC2MaxLatencyUs) {
                return false;
            }

            auto& interp = AML::GetInterpreter();
            if (!interp.IsInitialized()) {
                return false;
            }

            auto& ns = interp.GetNamespace();
            char path[256];

            for (int32_t nodeIndex = 0; nodeIndex < ns.NodeCount(); nodeIndex++) {
                auto* node = ns.GetNode(nodeIndex);
                if (node == nullptr || node->Obj.Type != AML::ObjectType::Processor) {
                    continue;
                }

                uint32_t pblk = node->Obj.Processor.PblkAddr;
                uint8_t pblkLen = node->Obj.Processor.PblkLen;
                if (pblk == 0 || pblk > 0xFFFF || pblkLen < 5) {
                    continue;
                }

                IdleState c1{};
                c1.Valid = true;
                c1.CType = 1;
                c1.LatencyUs = 1;
                c1.EntryKind = IdleEntryKind::FixedHardware;

                IdleState c2{};
                c2.Valid = true;
                c2.CType = 2;
                c2.LatencyUs = fadt.WorstC2Latency;
                c2.EntryKind = IdleEntryKind::IoPort;
                c2.IoPort = (uint16_t)(pblk + 4); // ACPI legacy P_LVL2 register
                c2.WidthBits = 8;

                g_idleStates[0] = c1;
                g_idleStates[1] = c2;
                g_idleStateCount = 2;

                ns.GetNodePath(nodeIndex, path, sizeof(path));
                KernelLogStream(OK, "CpuIdle") << "Using legacy ACPI P_BLK C2 idle state from "
                    << path << " latency=" << base::dec << (uint64_t)c2.LatencyUs
                    << "us port=" << base::hex << (uint64_t)c2.IoPort;
                return true;
            }

            return false;
        }

        static const IdleState* SelectIdleState(uint32_t predictedIdleMs, bool hasMwait) {
            uint64_t predictedIdleUs = (uint64_t)predictedIdleMs * 1000;
            const IdleState* best = nullptr;

            for (int i = 0; i < g_idleStateCount; i++) {
                const auto& state = g_idleStates[i];
                if (!state.Valid) continue;

                // Keep the first implementation conservative: skip ACPI I/O
                // states deeper than C2 until cache/bus-master coordination is
                // implemented.
                if (state.EntryKind == IdleEntryKind::IoPort && state.CType > 2) {
                    continue;
                }
                if (state.EntryKind == IdleEntryKind::IoPort &&
                    predictedIdleMs < IoStateMinPredictedMs) {
                    continue;
                }

                if (state.CType > 1 && state.LatencyUs != 0 &&
                    predictedIdleUs < (uint64_t)state.LatencyUs * ResidencyLatencyRatio) {
                    continue;
                }

                if (state.EntryKind == IdleEntryKind::FixedHardware && !hasMwait &&
                    state.CType > 1) {
                    continue;
                }

                if (best == nullptr || state.CType >= best->CType) {
                    best = &state;
                }
            }

            return best;
        }

        static void FallbackWait(bool hasMwait, volatile uint64_t* monitorAddr) {
            volatile uint64_t* addr = (monitorAddr != nullptr) ? monitorAddr : &g_fallbackMonitor;
            if (hasMwait) {
                Hal::IdleWait(addr);
            } else {
                asm volatile("hlt");
            }
        }

        static void FallbackWaitWithInterruptsDisabled(bool hasMwait,
                                                       volatile uint64_t* monitorAddr) {
            volatile uint64_t* addr = (monitorAddr != nullptr) ? monitorAddr : &g_fallbackMonitor;
            if (hasMwait) {
                Hal::IdleWaitWithInterruptsDisabled(addr);
            } else {
                Hal::HaltWithInterruptsDisabled();
            }
        }

        static void EnterIoIdleWithInterruptsDisabled(const IdleState& state) {
            switch (state.WidthBits) {
                case 16: {
                    uint16_t value;
                    asm volatile("sti\n\tinw %w1, %w0\n\tcli"
                                 : "=a"(value) : "Nd"(state.IoPort) : "memory");
                    break;
                }
                case 32: {
                    uint32_t value;
                    asm volatile("sti\n\tinl %w1, %0\n\tcli"
                                 : "=a"(value) : "Nd"(state.IoPort) : "memory");
                    break;
                }
                default: {
                    uint8_t value;
                    asm volatile("sti\n\tinb %w1, %b0\n\tcli"
                                 : "=a"(value) : "Nd"(state.IoPort) : "memory");
                    break;
                }
            }
        }

        void Initialize(ACPI::CommonSDTHeader* xsdt) {
            g_idleStateCount = 0;
            for (int i = 0; i < MaxIdleStates; i++) {
                g_idleStates[i] = {};
            }

            FADT::ParsedFADT fadt{};
            bool fadtReady = xsdt != nullptr && FADT::Parse(xsdt, fadt) && fadt.Valid;
            if (fadtReady && fadt.SMI_CommandPort != 0 && fadt.CStateControl != 0) {
                Io::Out8(fadt.CStateControl, (uint16_t)fadt.SMI_CommandPort);
                KernelLogStream(INFO, "CpuIdle") << "Advertised ACPI _CST support";
            }

            auto& interp = AML::GetInterpreter();
            if (!interp.IsInitialized()) {
                KernelLogStream(INFO, "CpuIdle") << "AML interpreter not ready - using MWAIT/HLT";
                return;
            }

            auto& ns = interp.GetNamespace();
            char path[256];

            for (int32_t nodeIndex = 0; nodeIndex < ns.NodeCount(); nodeIndex++) {
                auto* node = ns.GetNode(nodeIndex);
                if (node == nullptr) continue;
                if (!IsName(node->Name, "_CST")) continue;

                ns.GetNodePath(nodeIndex, path, sizeof(path));

                AML::Object obj{};
                if (!interp.EvaluateObject(path, obj)) {
                    continue;
                }

                if (ParseCstObject(path, obj)) {
                    return;
                }
            }

            if (fadtReady && LoadLegacyPblkIdleStates(fadt)) {
                return;
            }

            KernelLogStream(INFO, "CpuIdle")
                << "ACPI _CST unavailable or unsupported - using safe MWAIT/HLT idle fallback";
        }

        void Wait(uint32_t predictedIdleMs, bool hasMwait, volatile uint64_t* monitorAddr) {
            const IdleState* state = SelectIdleState(predictedIdleMs, hasMwait);
            if (state == nullptr) {
                FallbackWait(hasMwait, monitorAddr);
                return;
            }

            if (state->EntryKind == IdleEntryKind::IoPort) {
                switch (state->WidthBits) {
                    case 16:
                        (void)Io::In16(state->IoPort);
                        return;
                    case 32:
                        (void)Io::In32(state->IoPort);
                        return;
                    default:
                        (void)Io::In8(state->IoPort);
                        return;
                }
            }

            FallbackWait(hasMwait, monitorAddr);
        }

        void WaitWithInterruptsDisabled(uint32_t predictedIdleMs, bool hasMwait,
                                        volatile uint64_t* monitorAddr) {
            const IdleState* state = SelectIdleState(predictedIdleMs, hasMwait);
            if (state != nullptr && state->EntryKind == IdleEntryKind::IoPort) {
                EnterIoIdleWithInterruptsDisabled(*state);
                return;
            }

            FallbackWaitWithInterruptsDisabled(hasMwait, monitorAddr);
        }

        void WaitShallowWithInterruptsDisabled(bool hasMwait, volatile uint64_t* monitorAddr) {
            FallbackWaitWithInterruptsDisabled(hasMwait, monitorAddr);
        }

    };
};