/*
    * Hub.cpp
    * USB hub class driver (USB 2.0 and SuperSpeed hubs)
    * Copyright (c) 2026 Daniel Hammer
*/

#include "Hub.hpp"
#include "Xhci.hpp"
#include "UsbDevice.hpp"
#include <Arch/Cpu.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Timekeeping/Timer.hpp>
#include <atomic>

using namespace Kt;

static void BusyWaitMs(uint64_t ms) {
    if (Arch::InterruptsEnabled()) {
        uint64_t start = Timekeeping::GetMilliseconds();
        while (Timekeeping::GetMilliseconds() - start < ms) {
            Arch::Pause();
        }
    } else {
        // The tick is not advancing: count ~1 us delays instead
        for (uint64_t i = 0; i < ms * 1000; i++) {
            Arch::IoDelay();
        }
    }
}

namespace Drivers::USB::Hub {

    // -------------------------------------------------------------------------
    // Hub class requests and port status (USB 2.0 11.24, USB 3.2 10.16)
    // -------------------------------------------------------------------------

    constexpr uint8_t REQTYPE_HUB_IN     = 0xA0;  // Dev-to-host, class, device
    constexpr uint8_t REQTYPE_HUB_OUT    = 0x20;  // Host-to-dev, class, device
    constexpr uint8_t REQTYPE_PORT_IN    = 0xA3;  // Dev-to-host, class, other (port)
    constexpr uint8_t REQTYPE_PORT_OUT   = 0x23;  // Host-to-dev, class, other (port)

    constexpr uint8_t REQ_GET_STATUS     = 0x00;
    constexpr uint8_t REQ_CLEAR_FEATURE  = 0x01;
    constexpr uint8_t REQ_SET_FEATURE    = 0x03;
    constexpr uint8_t REQ_GET_DESCRIPTOR = 0x06;
    constexpr uint8_t REQ_SET_HUB_DEPTH  = 0x0C;  // SuperSpeed hubs only

    constexpr uint8_t DESC_HUB           = 0x29;
    constexpr uint8_t DESC_SS_HUB        = 0x2A;

    // Port feature selectors
    constexpr uint16_t FEAT_PORT_RESET   = 4;
    constexpr uint16_t FEAT_PORT_POWER   = 8;

    // wPortStatus
    constexpr uint16_t PS_CONNECTION     = (1 << 0);
    constexpr uint16_t PS_ENABLE         = (1 << 1);
    constexpr uint16_t PS_RESET          = (1 << 4);
    constexpr uint16_t PS_LOW_SPEED      = (1 << 9);   // USB 2 hubs
    constexpr uint16_t PS_HIGH_SPEED     = (1 << 10);  // USB 2 hubs

    // wPortChange
    constexpr uint16_t PC_CONNECTION     = (1 << 0);
    constexpr uint16_t PC_RESET          = (1 << 4);

    // The route string holds a 4-bit port number per tier, five tiers.
    constexpr uint8_t MAX_TIERS          = 5;
    constexpr uint8_t MAX_ROUTABLE_PORT  = 15;

    // -------------------------------------------------------------------------
    // State
    // -------------------------------------------------------------------------

    struct HubState {
        bool              Active;
        HubDescriptorInfo Info;
    };

    static HubState g_hubs[Xhci::MAX_SLOTS + 1] = {};

    // Written from the event-ring drain, consumed by ProcessPendingWork.
    static std::atomic<uint32_t> g_changedPorts[Xhci::MAX_SLOTS + 1] = {};  // bit n = port n
    static std::atomic<bool>     g_rearm[Xhci::MAX_SLOTS + 1] = {};
    static std::atomic<bool>     g_endpointFailed[Xhci::MAX_SLOTS + 1] = {};
    static std::atomic<bool>     g_pending{false};

    // -------------------------------------------------------------------------
    // Requests
    // -------------------------------------------------------------------------

    static bool Succeeded(uint32_t cc) {
        return cc == Xhci::CC_SUCCESS || cc == Xhci::CC_SHORT_PACKET;
    }

    static bool GetPortStatus(uint8_t hubSlot, uint8_t port, uint16_t& status, uint16_t& change) {
        uint8_t buf[4] = {};
        uint32_t cc = Xhci::ControlTransfer(hubSlot, REQTYPE_PORT_IN, REQ_GET_STATUS,
                                            0, port, sizeof(buf), buf, true);
        if (!Succeeded(cc)) return false;
        status = (uint16_t)(buf[0] | (buf[1] << 8));
        change = (uint16_t)(buf[2] | (buf[3] << 8));
        return true;
    }

    static uint32_t SetPortFeature(uint8_t hubSlot, uint8_t port, uint16_t feature) {
        return Xhci::ControlTransfer(hubSlot, REQTYPE_PORT_OUT, REQ_SET_FEATURE,
                                     feature, port, 0, nullptr, false);
    }

    static uint32_t ClearPortFeature(uint8_t hubSlot, uint8_t port, uint16_t feature) {
        return Xhci::ControlTransfer(hubSlot, REQTYPE_PORT_OUT, REQ_CLEAR_FEATURE,
                                     feature, port, 0, nullptr, false);
    }

    // Acknowledge every change bit in `change`. Bits 0-4 clear with
    // C_PORT_CONNECTION..C_PORT_RESET (16-20); SuperSpeed hubs add
    // C_BH_PORT_RESET, C_PORT_LINK_STATE and C_PORT_CONFIG_ERROR.
    static void ClearPortChanges(uint8_t hubSlot, uint8_t port, uint16_t change, bool superSpeed) {
        for (uint16_t bit = 0; bit <= 4; bit++) {
            if (change & (1u << bit)) ClearPortFeature(hubSlot, port, (uint16_t)(16 + bit));
        }
        if (superSpeed) {
            if (change & (1u << 5)) ClearPortFeature(hubSlot, port, 29);
            if (change & (1u << 6)) ClearPortFeature(hubSlot, port, 25);
            if (change & (1u << 7)) ClearPortFeature(hubSlot, port, 26);
        }
    }

    // Acknowledge hub-wide changes (local power, over-current): until they
    // are cleared the hub keeps reporting bit 0 of its change bitmap.
    static void ClearHubChanges(uint8_t hubSlot) {
        uint8_t buf[4] = {};
        uint32_t cc = Xhci::ControlTransfer(hubSlot, REQTYPE_HUB_IN, REQ_GET_STATUS,
                                            0, 0, sizeof(buf), buf, true);
        if (!Succeeded(cc)) return;
        uint16_t change = (uint16_t)(buf[2] | (buf[3] << 8));
        for (uint16_t bit = 0; bit <= 1; bit++) {  // C_HUB_LOCAL_POWER, C_HUB_OVER_CURRENT
            if (change & (1u << bit)) {
                Xhci::ControlTransfer(hubSlot, REQTYPE_HUB_OUT, REQ_CLEAR_FEATURE,
                                      bit, 0, 0, nullptr, false);
            }
        }
    }

    bool ReadDescriptor(uint8_t slotId, uint32_t speed, HubDescriptorInfo& out) {
        out = {};
        out.SuperSpeed = speed >= Xhci::SPEED_SUPER;
        uint8_t type = out.SuperSpeed ? DESC_SS_HUB : DESC_HUB;

        uint8_t buf[16] = {};
        uint32_t cc = Xhci::ControlTransfer(slotId, REQTYPE_HUB_IN, REQ_GET_DESCRIPTOR,
                                            (uint16_t)(type << 8), 0, sizeof(buf), buf, true);
        if (!Succeeded(cc) || buf[0] < 7 || buf[1] != type || buf[2] == 0) {
            KernelLogStream(ERROR, "USB/Hub") << "Slot " << base::dec << (uint64_t)slotId
                << ": hub descriptor unreadable, cc=" << (uint64_t)cc;
            return false;
        }

        uint16_t characteristics = (uint16_t)(buf[3] | (buf[4] << 8));
        out.Ports = buf[2];
        out.ThinkTime = out.SuperSpeed ? 0 : (uint8_t)((characteristics >> 5) & 3);
        out.PowerOnMs = (uint16_t)(buf[5] * 2);
        return true;
    }

    // -------------------------------------------------------------------------
    // Ports
    // -------------------------------------------------------------------------

    static uint8_t FindChild(uint8_t hubSlot, uint8_t port) {
        for (uint8_t s = 1; s <= Xhci::MAX_SLOTS; s++) {
            auto* dev = Xhci::GetDevice(s);
            if (dev->Active && dev->ParentSlot == hubSlot && dev->ParentPort == port) return s;
        }
        return 0;
    }

    // Reset a connected port and report the speed the device came up at.
    static bool ResetPort(uint8_t hubSlot, uint8_t port, bool superSpeed, uint32_t& speed) {
        if (SetPortFeature(hubSlot, port, FEAT_PORT_RESET) != Xhci::CC_SUCCESS) return false;

        uint16_t status = 0, change = 0;
        bool done = false;
        for (int attempt = 0; attempt < 50; attempt++) {  // ~500 ms
            BusyWaitMs(10);
            if (!GetPortStatus(hubSlot, port, status, change)) return false;
            if ((change & PC_RESET) && !(status & PS_RESET)) { done = true; break; }
            if (!(status & PS_CONNECTION)) return false;
        }
        ClearPortChanges(hubSlot, port, change, superSpeed);
        if (!done || !(status & PS_ENABLE)) return false;

        if (superSpeed)                   speed = Xhci::SPEED_SUPER;
        else if (status & PS_LOW_SPEED)   speed = Xhci::SPEED_LOW;
        else if (status & PS_HIGH_SPEED)  speed = Xhci::SPEED_HIGH;
        else                              speed = Xhci::SPEED_FULL;
        return true;
    }

    // Bring the port's state in line with what is attached to it: remove a
    // device that left (or was replaced), enumerate one that arrived.
    static void HandlePort(uint8_t hubSlot, uint8_t port, bool settled) {
        const HubDescriptorInfo& info = g_hubs[hubSlot].Info;

        uint16_t status = 0, change = 0;
        if (!GetPortStatus(hubSlot, port, status, change)) return;
        if (change != 0) ClearPortChanges(hubSlot, port, change, info.SuperSpeed);

        bool connected = (status & PS_CONNECTION) != 0;
        uint8_t child = FindChild(hubSlot, port);
        if (child != 0 && (!connected || (change & PC_CONNECTION))) {
            KernelLogStream(INFO, "USB/Hub") << "Slot " << base::dec << (uint64_t)hubSlot
                << " port " << (uint64_t)port << ": device disconnected";
            Xhci::RemoveDevice(child);
            child = 0;
        }
        if (!connected || child != 0) return;

        if (!settled) {
            // Connect debounce (USB 2.0 7.1.7.3): let contacts settle first.
            BusyWaitMs(100);
            if (!GetPortStatus(hubSlot, port, status, change)) return;
            if (!(status & PS_CONNECTION)) return;
        }

        auto* hub = Xhci::GetDevice(hubSlot);
        if (hub->Depth + 1 > MAX_TIERS || port > MAX_ROUTABLE_PORT) {
            KernelLogStream(WARNING, "USB/Hub") << "Slot " << base::dec << (uint64_t)hubSlot
                << " port " << (uint64_t)port << ": too deep to address, ignored";
            return;
        }

        uint32_t speed = 0;
        if (!ResetPort(hubSlot, port, info.SuperSpeed, speed)) {
            KernelLogStream(WARNING, "USB/Hub") << "Slot " << base::dec << (uint64_t)hubSlot
                << " port " << (uint64_t)port << ": reset failed";
            return;
        }
        BusyWaitMs(10);  // reset recovery (USB 2.0 7.1.7.5)

        UsbDevice::DeviceLocation location = {};
        location.RootPort    = hub->PortId;
        location.ParentSlot  = hubSlot;
        location.ParentPort  = port;
        location.Depth       = (uint8_t)(hub->Depth + 1);
        location.RouteString = hub->RouteString | ((uint32_t)port << (4 * hub->Depth));
        if (hub->Speed == Xhci::SPEED_HIGH &&
            (speed == Xhci::SPEED_LOW || speed == Xhci::SPEED_FULL)) {
            // This hub's transaction translator turns the high-speed
            // traffic into low/full-speed transactions. Single-TT mode:
            // the configuration's default alternate setting.
            location.TtHubSlot = hubSlot;
            location.TtPort    = port;
            location.TtMulti   = false;
        } else {
            location.TtHubSlot = hub->TtHubSlot;
            location.TtPort    = hub->TtPort;
            location.TtMulti   = hub->TtMulti;
        }

        UsbDevice::EnumerateDevice(location, speed);
    }

    // -------------------------------------------------------------------------
    // Public API
    // -------------------------------------------------------------------------

    void RegisterHub(uint8_t slotId, const HubDescriptorInfo& info) {
        auto* dev = Xhci::GetDevice(slotId);
        if (dev == nullptr || !dev->Active) return;

        g_hubs[slotId].Info = info;
        g_hubs[slotId].Active = true;
        g_changedPorts[slotId].store(0, std::memory_order_relaxed);
        g_rearm[slotId].store(false, std::memory_order_relaxed);
        g_endpointFailed[slotId].store(false, std::memory_order_relaxed);

        uint8_t ports = info.Ports;
        if (ports > 31) ports = 31;  // status-change bitmap tracked in 32 bits

        KernelLogStream(OK, "USB/Hub") << "Slot " << base::dec << (uint64_t)slotId << ": "
            << (info.SuperSpeed ? "SuperSpeed" : "USB 2.0") << " hub, "
            << (uint64_t)info.Ports << " ports, tier " << (uint64_t)(dev->Depth + 1);

        if (info.SuperSpeed) {
            Xhci::ControlTransfer(slotId, REQTYPE_HUB_OUT, REQ_SET_HUB_DEPTH,
                                  dev->Depth, 0, 0, nullptr, false);
        }

        for (uint8_t port = 1; port <= ports; port++) {
            SetPortFeature(slotId, port, FEAT_PORT_POWER);
        }
        // Power-on-to-power-good, then the connect debounce interval.
        uint64_t wait = info.PowerOnMs < 100 ? 100 : info.PowerOnMs;
        BusyWaitMs(wait + 100);

        for (uint8_t port = 1; port <= ports && g_hubs[slotId].Active; port++) {
            HandlePort(slotId, port, true);
        }

        if (g_hubs[slotId].Active) Xhci::QueueInterruptTransfer(slotId);
    }

    void UnregisterHub(uint8_t slotId) {
        if (slotId == 0 || slotId > Xhci::MAX_SLOTS) return;
        g_hubs[slotId] = {};
        g_changedPorts[slotId].store(0, std::memory_order_relaxed);
        g_rearm[slotId].store(false, std::memory_order_relaxed);
        g_endpointFailed[slotId].store(false, std::memory_order_relaxed);
    }

    void OnStatusChange(uint8_t slotId, const uint8_t* data, uint16_t length,
                        uint32_t completionCode) {
        if (slotId == 0 || slotId > Xhci::MAX_SLOTS) return;

        if (Succeeded(completionCode)) {
            uint32_t bits = 0;
            for (uint16_t i = 0; i < length && i < 4; i++) {
                bits |= (uint32_t)data[i] << (8 * i);
            }
            g_changedPorts[slotId].fetch_or(bits, std::memory_order_relaxed);
            g_rearm[slotId].store(true, std::memory_order_relaxed);
        } else {
            // The endpoint halted; it needs a reset before it reports again,
            // and changes may have been missed meanwhile.
            g_endpointFailed[slotId].store(true, std::memory_order_relaxed);
        }
        g_pending.store(true, std::memory_order_release);
    }

    void RequestPortCheck(uint8_t hubSlot, uint8_t port) {
        if (hubSlot == 0 || hubSlot > Xhci::MAX_SLOTS || port == 0 || port > 31) return;
        g_changedPorts[hubSlot].fetch_or(1u << port, std::memory_order_relaxed);
        g_pending.store(true, std::memory_order_release);
    }

    bool HasPendingWork() {
        return g_pending.load(std::memory_order_acquire);
    }

    void ProcessPendingWork() {
        if (!g_pending.exchange(false, std::memory_order_acq_rel)) return;

        for (uint8_t slot = 1; slot <= Xhci::MAX_SLOTS; slot++) {
            if (!g_hubs[slot].Active) continue;

            bool failed = g_endpointFailed[slot].exchange(false, std::memory_order_acq_rel);
            bool rearm = g_rearm[slot].exchange(false, std::memory_order_acq_rel);
            uint32_t changed = g_changedPorts[slot].exchange(0, std::memory_order_acq_rel);
            if (!failed && !rearm && changed == 0) continue;

            uint8_t ports = g_hubs[slot].Info.Ports;
            if (ports > 31) ports = 31;

            if (failed) {
                // Reset Endpoint re-arms the endpoint itself.
                Xhci::ResetInterruptEndpoint(slot);
                changed = ((ports == 31) ? 0xFFFFFFFFu : ((1u << (ports + 1)) - 1)) & ~1u;
            }

            if (changed & 1u) ClearHubChanges(slot);
            for (uint8_t port = 1; port <= ports && g_hubs[slot].Active; port++) {
                if (changed & (1u << port)) HandlePort(slot, port, false);
            }

            if (rearm && !failed && g_hubs[slot].Active) Xhci::QueueInterruptTransfer(slot);
        }
    }

}