/*
    * Uart.cpp
    * Polled serial console from the ACPI SPCR table
    * Copyright (c) 2026 Daniel Hammer
*/

#include "Uart.hpp"
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Memory/HHDM.hpp>
#include <Memory/Paging.hpp>
#include <Io/IoPort.hpp>
#include <Arch/Cpu.hpp>

using namespace Kt;

namespace Drivers::Serial {

    struct GenericAddress {
        uint8_t  AddressSpaceId;  // 0 = system memory, 1 = system I/O
        uint8_t  RegisterBitWidth;
        uint8_t  RegisterBitOffset;
        uint8_t  AccessSize;      // 1 = byte, 2 = word, 3 = dword, 4 = qword
        uint64_t Address;
    } __attribute__((packed));

    struct Spcr {
        Hal::ACPI::CommonSDTHeader Header;
        uint8_t        InterfaceType;
        uint8_t        Reserved[3];
        GenericAddress BaseAddress;
        uint8_t        InterruptType;
        uint8_t        Irq;
        uint32_t       GlobalSystemInterrupt;
        uint8_t        BaudRate;
        uint8_t        Parity;
        uint8_t        StopBits;
        uint8_t        FlowControl;
        uint8_t        TerminalType;
    } __attribute__((packed));

    // SPCR / DBG2 serial interface subtypes
    static constexpr uint8_t Type16550        = 0x00;
    static constexpr uint8_t Type16450        = 0x01;
    static constexpr uint8_t TypePl011        = 0x03;
    static constexpr uint8_t TypeSbsa32       = 0x0D;
    static constexpr uint8_t TypeSbsaGeneric  = 0x0E;
    static constexpr uint8_t TypeBcm2835      = 0x10;
    static constexpr uint8_t Type16550Generic = 0x12;

    // PL011: data register, flag register (TX FIFO full)
    static constexpr uint64_t Pl011Dr = 0x00;
    static constexpr uint64_t Pl011Fr = 0x18;
    static constexpr uint32_t Pl011FrTxff = 1u << 5;

    // 16550: transmit holding register, line status register (THR empty)
    static constexpr uint64_t Uart16550Thr = 0;
    static constexpr uint64_t Uart16550Lsr = 5;
    static constexpr uint8_t Uart16550LsrThre = 1u << 5;

    // Bounded so a wedged or absent UART cannot hang the console.
    static constexpr int TxSpinLimit = 100000;

    static volatile uint8_t* g_mmio = nullptr;
    static uint16_t g_port = 0;
    static uint32_t g_stride = 1;

    static inline uint32_t Read32(uint64_t offset) {
        return *(volatile uint32_t*)(g_mmio + offset);
    }

    static inline void Write32(uint64_t offset, uint32_t value) {
        *(volatile uint32_t*)(g_mmio + offset) = value;
    }

    static void Pl011Putchar(char c) {
        for (int i = 0; i < TxSpinLimit && (Read32(Pl011Fr) & Pl011FrTxff); i++) {
            Arch::Pause();
        }
        Write32(Pl011Dr, (uint8_t)c);
    }

    static uint8_t Mmio16550Read(uint64_t reg) {
        if (g_stride == 4) return (uint8_t)Read32(reg * 4);
        return g_mmio[reg];
    }

    static void Mmio16550Write(uint64_t reg, uint8_t value) {
        if (g_stride == 4) Write32(reg * 4, value);
        else g_mmio[reg] = value;
    }

    static void Mmio16550Putchar(char c) {
        for (int i = 0; i < TxSpinLimit && !(Mmio16550Read(Uart16550Lsr) & Uart16550LsrThre); i++) {
            Arch::Pause();
        }
        Mmio16550Write(Uart16550Thr, (uint8_t)c);
    }

    static void Port16550Putchar(char c) {
        for (int i = 0; i < TxSpinLimit && !(Io::In8(g_port + Uart16550Lsr) & Uart16550LsrThre); i++) {
            Arch::Pause();
        }
        Io::Out8((uint8_t)c, g_port + Uart16550Thr);
    }

    static Spcr* FindSpcr(Hal::ACPI::CommonSDTHeader* xsdt) {
        uint32_t entryCount = (xsdt->Length - sizeof(Hal::ACPI::CommonSDTHeader)) / 8;
        uint64_t* entries = (uint64_t*)((uint64_t)xsdt + sizeof(Hal::ACPI::CommonSDTHeader));

        for (uint32_t i = 0; i < entryCount; i++) {
            auto* header = (Hal::ACPI::CommonSDTHeader*)Memory::HHDM(entries[i]);
            const char* sig = header->Signature;
            if (sig[0] == 'S' && sig[1] == 'P' && sig[2] == 'C' && sig[3] == 'R') {
                return (Spcr*)header;
            }
        }
        return nullptr;
    }

    static volatile uint8_t* MapRegisters(uint64_t phys) {
        uint64_t page = phys & ~0xFFFULL;
        Memory::VMM::g_paging->MapMMIO(page, Memory::HHDM(page));
        return (volatile uint8_t*)Memory::HHDM(phys);
    }

    void InitializeConsole(Hal::ACPI::CommonSDTHeader* xsdt) {
        if (xsdt == nullptr) return;

        Spcr* spcr = FindSpcr(xsdt);
        if (spcr == nullptr || spcr->Header.Length < sizeof(Spcr)) {
            return;
        }

        const GenericAddress& reg = spcr->BaseAddress;
        if (reg.Address == 0) return;

        SerialSink sink = nullptr;
        const char* name = "";
        bool isIo = reg.AddressSpaceId == 1;

        switch (spcr->InterfaceType) {
            case TypePl011:
            case TypeSbsa32:
            case TypeSbsaGeneric:
                if (isIo) return;
                g_mmio = MapRegisters(reg.Address);
                sink = Pl011Putchar;
                name = "PL011";
                break;

            case Type16550:
            case Type16450:
            case Type16550Generic:
            case TypeBcm2835:
                if (isIo) {
                    if (!Io::HasPortIo) return;
                    g_port = (uint16_t)reg.Address;
                    sink = Port16550Putchar;
                } else {
                    // The BCM2835 mini UART is a 16550 with 32-bit registers.
                    g_stride = (reg.RegisterBitWidth == 32 || reg.AccessSize == 3
                                || spcr->InterfaceType == TypeBcm2835) ? 4 : 1;
                    g_mmio = MapRegisters(reg.Address);
                    sink = Mmio16550Putchar;
                }
                name = spcr->InterfaceType == TypeBcm2835 ? "BCM2835 mini UART" : "16550";
                break;

            default:
                KernelLogStream(WARNING, "Serial") << "SPCR interface type 0x" << base::hex
                    << (uint64_t)spcr->InterfaceType << " not supported";
                return;
        }

        AttachSerialConsole(sink);
        KernelLogStream(OK, "Serial") << "Console on " << name << " at " << (isIo ? "port " : "")
            << "0x" << base::hex << reg.Address;
    }

}