#include <Arch/DTB.h>
#include <OS/Panic.h>
#include <Lib/Bytes.h>
#include <Lib/Align.h>
#include <Lib/String.h>
#include <VM/PMM.h>

void DTBParse(Pointer dtb, VMBootMemoryMap* bootMap) {
    FDTHeader* header = (FDTHeader*)dtb;
    if (BytesSwap32(header->magic) != kFDTHeaderMagic) {
        OSPanic("Invalid DTB magic");
    }

    UInt32 index = bootMap->reservedCount;
    bootMap->reserved[index].base = (Address)dtb;
    bootMap->reserved[index].size = BytesSwap32(header->totalSize);
    bootMap->reservedCount++;

    UInt32 offStruct = BytesSwap32(header->offDtStruct);
    UInt32 offStrings = BytesSwap32(header->offDtStrings);

    BytePointer structs = (BytePointer)dtb + offStruct;
    ASCII* strings = (ASCII*)dtb + offStrings;

    ASCII* currentNode = "";
    UInt32 currentDepth = 0;
    UInt32 reservedMemoryDepth = 0; 
    bool inReservedMemory = false;

    while (true) {
        UInt32 token = BytesSwap32(*(UInt32*)structs);
        structs += 4;

        switch (token) {
            case FDTTokenBeginNode: {
                currentDepth++;
                currentNode = (ASCII*)structs;

                if (StringStartsWith(currentNode, "reserved-memory")) {
                    inReservedMemory = true;
                    reservedMemoryDepth = currentDepth;
                }

                UInt32 nameLength = StringGetLength(currentNode);
                structs += (nameLength + 1);
                structs = (BytePointer)AlignUp64((Address)structs, 4);
                break;
            }
            case FDTTokenProperty: {
                UInt32 propertyLength = BytesSwap32(*(UInt32*)structs);
                UInt32 nameOffset = BytesSwap32(*(UInt32*)(structs + 4));
                structs += 8;
                ASCII* propertyName = strings + nameOffset;

                if (StringCompare(propertyName, "reg") == 0) {
                    if (StringStartsWith(currentNode, "memory")) {
                        UInt32* cells = (UInt32*)structs;
                        Address base = Merge32To64(BytesSwap32(cells[1]), BytesSwap32(cells[0]));
                        Size size = Merge32To64(BytesSwap32(cells[3]), BytesSwap32(cells[2]));

                        bootMap->totalRAM.base = base;
                        bootMap->totalRAM.size = size;
                    }
                    else if (inReservedMemory && currentDepth > reservedMemoryDepth) {
                        UInt32* cells = (UInt32*)structs;
                        Address base = Merge32To64(BytesSwap32(cells[1]), BytesSwap32(cells[0]));
                        Size size = Merge32To64(BytesSwap32(cells[3]), BytesSwap32(cells[2]));

                        UInt32 index = bootMap->reservedCount;
                        bootMap->reserved[index].base = base;
                        bootMap->reserved[index].size = size;
                        bootMap->reservedCount++;
                    }
                    else if (StringStartsWith(currentNode, "pl011")) {
                        UInt32* cells = (UInt32*)structs;
                        bootMap->UART.base = Merge32To64(BytesSwap32(cells[1]), BytesSwap32(cells[0]));
                        bootMap->UART.size = Merge32To64(BytesSwap32(cells[3]), BytesSwap32(cells[2]));
                    }
                    else if (StringStartsWith(currentNode, "intc")) {
                        UInt32* cells = (UInt32*)structs;

                        bootMap->GIC.GICD.base = Merge32To64(BytesSwap32(cells[1]), BytesSwap32(cells[0]));
                        bootMap->GIC.GICD.size = Merge32To64(BytesSwap32(cells[3]), BytesSwap32(cells[2]));
                        
                        bootMap->GIC.GICC.base = Merge32To64(BytesSwap32(cells[5]), BytesSwap32(cells[4]));
                        bootMap->GIC.GICC.size = Merge32To64(BytesSwap32(cells[7]), BytesSwap32(cells[6]));
                    }
                }

                structs += propertyLength;
                structs = (BytePointer)AlignUp64((Address)structs, 4);
                break;
            }

            case FDTTokenEndNode: {
                if (inReservedMemory && currentDepth == reservedMemoryDepth) {
                    inReservedMemory = false;
                }
                currentDepth--;
                break;
            }
            
            case FDTTokenNOP: continue;
            case FDTTokenEnd: return;
            default:
                OSPanic("Invalid DTB token");
        }
    }

}