/*
    * AmlParser.cpp
    * AML bytecode parser — S5 extraction (brute-force) and interpreter init
    * Copyright (c) 2026 Daniel Hammer
*/

#include "AmlParser.hpp"
#include "AmlInterpreter.hpp"
#include <ACPI/ACPI.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Memory/HHDM.hpp>
#include <Memory/Heap.hpp>

using namespace Kt;

namespace Hal {
    namespace AML {

        // ============================================================================

        // Legacy S5 extraction (brute-force scan)

        // ============================================================================
        // Kept for fast S5 extraction during early boot before the full
        // interpreter is loaded.

        static constexpr uint8_t NameOp_     = 0x08;
        static constexpr uint8_t PackageOp_  = 0x12;
        static constexpr uint8_t ZeroOp_     = 0x00;
        static constexpr uint8_t OneOp_      = 0x01;
        static constexpr uint8_t OnesOp_     = 0xFF;
        static constexpr uint8_t BytePrefix_ = 0x0A;
        static constexpr uint8_t WordPrefix_ = 0x0B;
        static constexpr uint8_t DWordPrefix_= 0x0C;

        static uint32_t DecodePkgLength(const uint8_t* aml, uint32_t* pos) {
            uint8_t lead = aml[*pos];
            uint32_t byteCount = (lead >> 6) & 0x03;

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

            uint32_t length = lead & 0x0F;
            (*pos)++;

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

            return length;
        }

        static uint32_t DecodeIntegerLegacy(const uint8_t* aml, uint32_t* pos) {
            uint8_t op = aml[*pos];

            switch (op) {
                case ZeroOp_:
                    (*pos)++;
                    return 0;
                case OneOp_:
                    (*pos)++;
                    return 1;
                case OnesOp_:
                    (*pos)++;
                    return 0xFFFFFFFF;
                case BytePrefix_: {
                    (*pos)++;
                    uint8_t val = aml[*pos];
                    (*pos)++;
                    return val;
                }
                case WordPrefix_: {
                    (*pos)++;
                    uint16_t val = aml[*pos] | ((uint16_t)aml[*pos + 1] << 8);
                    *pos += 2;
                    return val;
                }
                case DWordPrefix_: {
                    (*pos)++;
                    uint32_t val = aml[*pos]
                                 | ((uint32_t)aml[*pos + 1] << 8)
                                 | ((uint32_t)aml[*pos + 2] << 16)
                                 | ((uint32_t)aml[*pos + 3] << 24);
                    *pos += 4;
                    return val;
                }
                default:
                    (*pos)++;
                    return 0;
            }
        }

        S5Object FindS5(void* dsdtData) {
            S5Object result{};
            result.Valid = false;

            auto* header = (ACPI::CommonSDTHeader*)dsdtData;

            if (!ACPI::TestChecksum(header)) {
                KernelLogStream(ERROR, "AML") << "DSDT checksum failed";
                return result;
            }

            const uint8_t* aml = (const uint8_t*)dsdtData;
            uint32_t amlLength = header->Length;
            uint32_t dataStart = sizeof(ACPI::CommonSDTHeader);

            for (uint32_t i = dataStart; i + 4 < amlLength; i++) {
                if (aml[i] == '_' && aml[i+1] == 'S' && aml[i+2] == '5' && aml[i+3] == '_') {
                    bool validContext = false;

                    if (i >= 1 && aml[i-1] == NameOp_) {
                        validContext = true;
                    } else if (i >= 2 && aml[i-2] == NameOp_ && aml[i-1] == '\\') {
                        validContext = true;
                    }

                    if (!validContext)
                        continue;

                    KernelLogStream(OK, "AML") << "Found \\_S5_ object at offset " << base::hex << (uint64_t)i;

                    uint32_t pos = i + 4;

                    if (pos >= amlLength || aml[pos] != PackageOp_) {
                        KernelLogStream(ERROR, "AML") << "Expected PackageOp after \\_S5_, got " << base::hex << (uint64_t)aml[pos];
                        continue;
                    }
                    pos++;

                    DecodePkgLength(aml, &pos);

                    if (pos >= amlLength) continue;
                    uint8_t numElements = aml[pos];
                    pos++;

                    if (numElements < 1) {
                        KernelLogStream(ERROR, "AML") << "\\_S5_ package has no elements";
                        continue;
                    }

                    result.SLP_TYPa = (uint16_t)DecodeIntegerLegacy(aml, &pos);

                    if (numElements >= 2 && pos < amlLength) {
                        result.SLP_TYPb = (uint16_t)DecodeIntegerLegacy(aml, &pos);
                    } else {
                        result.SLP_TYPb = 0;
                    }

                    result.Valid = true;

                    KernelLogStream(OK, "AML") << "SLP_TYPa=" << base::hex << (uint64_t)result.SLP_TYPa
                        << " SLP_TYPb=" << base::hex << (uint64_t)result.SLP_TYPb;

                    return result;
                }
            }

            KernelLogStream(ERROR, "AML") << "\\_S5_ object not found in DSDT";
            return result;
        }

        // ============================================================================

        // Generalized brute-force sleep state scanner

        // ============================================================================
        SleepObject FindSleepState(void* dsdtData, int state) {
            SleepObject result{};
            result.Valid = false;

            if (state < 0 || state > 5) return result;

            auto* header = (ACPI::CommonSDTHeader*)dsdtData;

            if (!ACPI::TestChecksum(header)) {
                KernelLogStream(ERROR, "AML") << "DSDT checksum failed";
                return result;
            }

            // Build the 4-char name we're looking for: _S0_ through _S5_
            char target[4] = { '_', 'S', (char)('0' + state), '_' };

            const uint8_t* aml = (const uint8_t*)dsdtData;
            uint32_t amlLength = header->Length;
            uint32_t dataStart = sizeof(ACPI::CommonSDTHeader);

            for (uint32_t i = dataStart; i + 4 < amlLength; i++) {
                if (aml[i] == target[0] && aml[i+1] == target[1] &&
                    aml[i+2] == target[2] && aml[i+3] == target[3]) {

                    bool validContext = false;
                    if (i >= 1 && aml[i-1] == NameOp_)
                        validContext = true;
                    else if (i >= 2 && aml[i-2] == NameOp_ && aml[i-1] == '\\')
                        validContext = true;

                    if (!validContext) continue;

                    uint32_t pos = i + 4;

                    if (pos >= amlLength || aml[pos] != PackageOp_) continue;
                    pos++;

                    DecodePkgLength(aml, &pos);

                    if (pos >= amlLength) continue;
                    uint8_t numElements = aml[pos];
                    pos++;

                    if (numElements < 1) continue;

                    result.SLP_TYPa = (uint16_t)DecodeIntegerLegacy(aml, &pos);

                    if (numElements >= 2 && pos < amlLength)
                        result.SLP_TYPb = (uint16_t)DecodeIntegerLegacy(aml, &pos);
                    else
                        result.SLP_TYPb = 0;

                    result.Valid = true;

                    KernelLogStream(OK, "AML") << "\\_S" << base::dec << (uint64_t)state
                        << "_ found: SLP_TYPa=" << base::hex << (uint64_t)result.SLP_TYPa
                        << " SLP_TYPb=" << base::hex << (uint64_t)result.SLP_TYPb;

                    return result;
                }
            }

            KernelLogStream(INFO, "AML") << "\\_S" << base::dec << (uint64_t)state
                << "_ not found in DSDT";
            return result;
        }

        // ============================================================================

        // Full interpreter initialization

        // ============================================================================
        void InitializeInterpreter(void* dsdtData, ACPI::CommonSDTHeader* xsdt) {
            auto& interp = GetInterpreter();
            if (!interp.LoadTable(dsdtData)) {
                KernelLogStream(ERROR, "AML") << "Failed to load DSDT into AML interpreter";
            }

            // SSDTs extend the DSDT's namespace and are loaded after it, in
            // XSDT order. They are where most laptops keep CPU power tables,
            // and -- the reason they matter here -- the discrete GPU's power
            // control (_PR3 power resources, _PS3, NVIDIA/AMD _DSM).
            int ssdts = 0;
            if (xsdt) {
                uint32_t entryCount = (xsdt->Length - sizeof(ACPI::CommonSDTHeader)) / 8;
                auto* entries = (uint64_t*)((uint64_t)xsdt + sizeof(ACPI::CommonSDTHeader));
                for (uint32_t i = 0; i < entryCount; i++) {
                    if (entries[i] == 0) continue;
                    auto* table = (ACPI::CommonSDTHeader*)Memory::HHDM(entries[i]);
                    const char* sig = table->Signature;
                    bool isSsdt = sig[0] == 'S' && sig[1] == 'S' && sig[2] == 'D' && sig[3] == 'T';
                    bool isPsdt = sig[0] == 'P' && sig[1] == 'S' && sig[2] == 'D' && sig[3] == 'T';
                    if (!isSsdt && !isPsdt) continue;
                    if (table->Length <= sizeof(ACPI::CommonSDTHeader)) continue;
                    if (interp.LoadTable(table)) ssdts++;
                }
            }
            KernelLogStream(OK, "AML") << "Loaded DSDT + " << base::dec << (uint64_t)ssdts
                << " SSDT(s)";

            auto& ns = interp.GetNamespace();
            KernelLogStream(OK, "AML") << "Namespace: " << base::dec
                << (uint64_t)ns.NodeCount() << " / " << (uint64_t)ns.MaxCapacity() << " nodes (dynamic)";
        }


        // ============================================================================

        // Namespace initialization (_REG, _INI)

        // ============================================================================
        static bool SegIs(const char* name, const char* seg) {
            return name[0] == seg[0] && name[1] == seg[1] &&
                   name[2] == seg[2] && name[3] == seg[3];
        }

        void InitializeNamespaceObjects() {
            auto& interp = GetInterpreter();
            if (!interp.IsInitialized()) return;
            auto& ns = interp.GetNamespace();

            // Collect first, evaluate second: executing AML creates nodes,
            // which reallocates a parent's child array under a live walk.
            int32_t capacity = ns.NodeCount();
            auto* regs = (int32_t*)Memory::g_heap->Request(sizeof(int32_t) * capacity);
            auto* devices = (int32_t*)Memory::g_heap->Request(sizeof(int32_t) * capacity);
            if (!regs || !devices) {
                if (regs) Memory::g_heap->Free(regs);
                if (devices) Memory::g_heap->Free(devices);
                return;
            }
            int regCount = 0, deviceCount = 0;

            ns.WalkDescendants(ns.RootIndex(), ObjectType::None,
                [&](int32_t idx, const NamespaceNode* node) {
                    if (node->Obj.Type == ObjectType::Method && SegIs(node->Name, "_REG") &&
                        regCount < capacity)
                        regs[regCount++] = idx;
                    if (node->Obj.Type == ObjectType::Device && deviceCount < capacity)
                        devices[deviceCount++] = idx;
                });

            // _REG(PCI_Config, connect). A _REG for another space checks Arg0
            // and ignores the call, as with the EC's.
            Object args[2]{};
            args[0].Type = ObjectType::Integer;
            args[0].Integer = (uint64_t)RegionSpace::PciConfig;
            args[1].Type = ObjectType::Integer;
            args[1].Integer = 1;
            for (int i = 0; i < regCount; i++) {
                Object result{};
                interp.EvaluateNode(regs[i], args, 2, result);
            }

            // \_SB._INI runs before any device's.
            int32_t sb = ns.FindNode("\\_SB_");
            if (sb >= 0) {
                Object result{};
                if (interp.EvaluateChild(sb, "_INI", nullptr, 0, result))
                    KernelLogStream(INFO, "AML") << "Ran \\_SB._INI";
            }

            // Then each device, parents first (WalkDescendants is pre-order).
            // A device that is neither present nor functional takes its whole
            // subtree with it; one that is functional but absent skips only
            // its own _INI.
            static constexpr int MaxSkipped = 256;
            int32_t skipped[MaxSkipped];
            int skippedCount = 0;
            int ran = 0;

            for (int i = 0; i < deviceCount; i++) {
                int32_t dev = devices[i];

                bool underSkipped = false;
                for (int32_t p = ns.GetNode(dev) ? ns.GetNode(dev)->ParentIndex : -1;
                     p > 0 && !underSkipped; p = ns.GetNode(p) ? ns.GetNode(p)->ParentIndex : -1) {
                    for (int k = 0; k < skippedCount; k++)
                        if (skipped[k] == p) { underSkipped = true; break; }
                }
                if (underSkipped) continue;

                uint64_t sta = 0x0F;
                Object staObj{};
                if (interp.EvaluateChild(dev, "_STA", nullptr, 0, staObj) &&
                    staObj.Type == ObjectType::Integer)
                    sta = staObj.Integer;

                if (!(sta & 0x01)) {
                    if (!(sta & 0x08) && skippedCount < MaxSkipped)
                        skipped[skippedCount++] = dev;
                    continue;
                }

                if (interp.FindChild(dev, "_INI") < 0) continue;
                Object result{};
                bool ok = interp.EvaluateChild(dev, "_INI", nullptr, 0, result);
                ran++;

                char path[128];
                ns.GetNodePath(dev, path, sizeof(path));
                KernelLogStream(DEBUG, "AML") << "  _INI " << path << (ok ? "" : " FAILED");
            }

            KernelLogStream(OK, "AML") << "Namespace initialized: " << base::dec
                << (uint64_t)regCount << " _REG, " << (uint64_t)ran << " device _INI";

            Memory::g_heap->Free(regs);
            Memory::g_heap->Free(devices);
        }

    };
};