/*
    * L2cap.cpp
    * Bluetooth L2CAP implementation
    * Copyright (c) 2026 Daniel Hammer
*/

#include "L2cap.hpp"
#include "Hci.hpp"
#include "A2dp.hpp"
#include "Avrcp.hpp"
#include <Drivers/USB/Xhci.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Libraries/Memory.hpp>
#include <Timekeeping/ApicTimer.hpp>

using namespace Kt;

namespace Drivers::USB::Bluetooth::L2cap {

    // =========================================================================
    // State
    // =========================================================================

    static uint16_t g_aclHandle = 0;
    static bool g_initialized = false;
    static uint8_t g_sigIdentifier = 1;

    // Diagnostic: result of the most recent L2CAP Connection Response we got for
    // an outgoing connection (0=success, 1=pending, 2=PSM-unsupported,
    // 3=security-block, 4=no-resources).  0xFFFF = none received yet.
    static volatile uint16_t g_lastConnRspResult = 0xFFFF;

    // Channel table
    static ChannelInfo g_channels[MAX_CHANNELS] = {};
    static uint16_t g_nextCid = CID_DYNAMIC_START;

    // =========================================================================
    // Helpers
    // =========================================================================

    static uint16_t AllocCid() {
        return g_nextCid++;
    }

    static ChannelInfo* AllocChannel(uint16_t psm) {
        for (int i = 0; i < MAX_CHANNELS; i++) {
            if (!g_channels[i].Active) {
                g_channels[i].Active = true;
                g_channels[i].LocalCid = AllocCid();
                g_channels[i].RemoteCid = 0;
                g_channels[i].Psm = psm;
                g_channels[i].RemoteMtu = 672;  // Default L2CAP MTU
                g_channels[i].Configured = false;
                g_channels[i].LocalConfigDone = false;
                g_channels[i].RemoteConfigDone = false;
                return &g_channels[i];
            }
        }
        return nullptr;
    }

    // Send L2CAP signaling command
    static void SendSignal(uint8_t code, uint8_t identifier,
                           const uint8_t* payload, uint16_t payloadLen) {
        // L2CAP header + Signal header + payload
        uint16_t sigLen = sizeof(SignalHeader) + payloadLen;
        uint16_t totalPayload = sizeof(L2capHeader) + sigLen;

        uint8_t buf[128] = {};
        auto* l2hdr = (L2capHeader*)buf;
        l2hdr->Length = sigLen;
        l2hdr->ChannelId = CID_SIGNALING;

        auto* sig = (SignalHeader*)(buf + sizeof(L2capHeader));
        sig->Code = code;
        sig->Identifier = identifier;
        sig->Length = payloadLen;

        if (payload && payloadLen > 0) {
            memcpy(buf + sizeof(L2capHeader) + sizeof(SignalHeader), payload, payloadLen);
        }

        Hci::SendAcl(g_aclHandle, Hci::ACL_PB_FIRST_FLUSH,
                      buf, totalPayload);
    }

    // =========================================================================
    // Initialize
    // =========================================================================

    void Initialize(uint16_t aclHandle) {
        g_aclHandle = aclHandle;
        g_initialized = true;
        g_sigIdentifier = 1;
        g_nextCid = CID_DYNAMIC_START;

        for (int i = 0; i < MAX_CHANNELS; i++) {
            g_channels[i].Active = false;
        }

        KernelLogStream(OK, "BT-L2CAP") << "Initialized for ACL handle " << (uint64_t)aclHandle;
    }

    // =========================================================================
    // ProcessPacket
    // =========================================================================

    void ProcessPacket(uint16_t aclHandle, const uint8_t* data, uint16_t len) {
        if (len < sizeof(L2capHeader)) return;

        auto* l2hdr = (const L2capHeader*)data;
        uint16_t l2len = l2hdr->Length;
        uint16_t cid = l2hdr->ChannelId;
        const uint8_t* payload = data + sizeof(L2capHeader);

        if (l2len + sizeof(L2capHeader) > len) return;

        if (cid == CID_SIGNALING) {
            // L2CAP signaling channel.  One packet may carry MULTIPLE commands
            // back to back -- a headset commonly bundles its Config Response
            // (to our request) with its own Config Request in a single packet.
            // Loop over all of them, or we'd handle one and silently drop the
            // rest (which left our channel half-configured: localCfg=0).
            uint16_t sOff = 0;
            while (sOff + sizeof(SignalHeader) <= l2len) {
              auto* sig = (const SignalHeader*)(payload + sOff);
              const uint8_t* sigPayload = payload + sOff + sizeof(SignalHeader);
              uint16_t sigPayloadLen = sig->Length;
              // Clamp (don't bail) so a length that doesn't exactly match the
              // packet never skips this command -- the cases bounds-check their
              // own reads.  Advancing by the clamped length ends the loop after
              // the last real command.
              uint16_t avail = (uint16_t)(l2len - sOff - sizeof(SignalHeader));
              if (sigPayloadLen > avail) sigPayloadLen = avail;

              switch (sig->Code) {
                case SIG_CONN_REQ: {
                    if (sigPayloadLen >= 4) {
                        uint16_t psm = (uint16_t)sigPayload[0] | ((uint16_t)sigPayload[1] << 8);
                        uint16_t srcCid = (uint16_t)sigPayload[2] | ((uint16_t)sigPayload[3] << 8);

                        KernelLogStream(INFO, "BT-L2CAP") << "Connection Request: PSM="
                            << base::hex << (uint64_t)psm << " srcCID=" << (uint64_t)srcCid;

                        // Accept connections for AVDTP, SDP, and AVRCP control
                        if (psm == PSM_AVDTP || psm == PSM_SDP || psm == PSM_AVCTP) {
                            auto* ch = AllocChannel(psm);
                            if (ch) {
                                ch->RemoteCid = srcCid;

                                // Send Connection Response (success)
                                uint8_t rsp[8] = {};
                                rsp[0] = (uint8_t)(ch->LocalCid & 0xFF);
                                rsp[1] = (uint8_t)(ch->LocalCid >> 8);
                                rsp[2] = (uint8_t)(srcCid & 0xFF);
                                rsp[3] = (uint8_t)(srcCid >> 8);
                                rsp[4] = 0; rsp[5] = 0;  // Result: success
                                rsp[6] = 0; rsp[7] = 0;  // Status: no info
                                SendSignal(SIG_CONN_RSP, sig->Identifier, rsp, 8);

                                // As the acceptor we must ALSO send our own
                                // Configuration Request -- otherwise our side
                                // never reaches LocalConfigDone and the channel
                                // never becomes Configured (OnChannelReady never
                                // fires).  Address the remote's CID.
                                uint8_t cfgReq[4] = {};
                                cfgReq[0] = (uint8_t)(srcCid & 0xFF);
                                cfgReq[1] = (uint8_t)(srcCid >> 8);
                                cfgReq[2] = 0; cfgReq[3] = 0;  // Flags
                                SendSignal(SIG_CONFIG_REQ, g_sigIdentifier++, cfgReq, 4);
                            }
                        } else {
                            // Reject: PSM not supported
                            uint8_t rsp[8] = {};
                            rsp[0] = 0; rsp[1] = 0;  // Dest CID = 0
                            rsp[2] = (uint8_t)(srcCid & 0xFF);
                            rsp[3] = (uint8_t)(srcCid >> 8);
                            rsp[4] = 0x02; rsp[5] = 0;  // Result: PSM not supported
                            rsp[6] = 0; rsp[7] = 0;
                            SendSignal(SIG_CONN_RSP, sig->Identifier, rsp, 8);
                        }
                    }
                    break;
                }

                case SIG_CONN_RSP: {
                    if (sigPayloadLen >= 8) {
                        uint16_t dstCid = (uint16_t)sigPayload[0] | ((uint16_t)sigPayload[1] << 8);
                        uint16_t srcCid = (uint16_t)sigPayload[2] | ((uint16_t)sigPayload[3] << 8);
                        uint16_t result = (uint16_t)sigPayload[4] | ((uint16_t)sigPayload[5] << 8);
                        // Status (bytes 6-7) explains a PENDING (result=1): 0=no
                        // info, 1=authentication pending, 2=authorization pending.
                        uint16_t status = (sigPayloadLen >= 8)
                            ? ((uint16_t)sigPayload[6] | ((uint16_t)sigPayload[7] << 8)) : 0;
                        g_lastConnRspResult = result;

                        KernelLogStream(INFO, "BT-L2CAP") << "Connection Response: dstCID="
                            << base::hex << (uint64_t)dstCid << " result=" << (uint64_t)result
                            << " status=" << (uint64_t)status << base::dec;

                        if (result == CONN_SUCCESS) {
                            // Find our channel by srcCid (which is our local CID)
                            for (int i = 0; i < MAX_CHANNELS; i++) {
                                if (g_channels[i].Active && g_channels[i].LocalCid == srcCid) {
                                    g_channels[i].RemoteCid = dstCid;

                                    // Send Configuration Request
                                    uint8_t cfgReq[4] = {};
                                    cfgReq[0] = (uint8_t)(dstCid & 0xFF);
                                    cfgReq[1] = (uint8_t)(dstCid >> 8);
                                    cfgReq[2] = 0; cfgReq[3] = 0;  // Flags
                                    SendSignal(SIG_CONFIG_REQ, g_sigIdentifier++, cfgReq, 4);
                                    break;
                                }
                            }
                        } else if (result == CONN_PENDING && dstCid != 0) {
                            // A PENDING response may already carry the remote's
                            // allocated CID.  Record it (but do NOT configure --
                            // the final SUCCESS on this channel does that): the
                            // diagnostics then show the peer engaged, and if the
                            // dial is later abandoned, FreeChannel knows to send
                            // a Disconnect for the half-open remote endpoint.
                            for (int i = 0; i < MAX_CHANNELS; i++) {
                                if (g_channels[i].Active && g_channels[i].LocalCid == srcCid) {
                                    g_channels[i].RemoteCid = dstCid;
                                    break;
                                }
                            }
                        }
                    }
                    break;
                }

                case SIG_CONFIG_REQ: {
                    if (sigPayloadLen >= 4) {
                        uint16_t dstCid = (uint16_t)sigPayload[0] | ((uint16_t)sigPayload[1] << 8);

                        // Find channel
                        for (int i = 0; i < MAX_CHANNELS; i++) {
                            if (g_channels[i].Active && g_channels[i].LocalCid == dstCid) {
                                g_channels[i].RemoteConfigDone = true;

                                // Parse MTU option if present
                                uint16_t cfgOffset = 4;  // Skip dstCid + flags
                                while (cfgOffset + 2 <= sigPayloadLen) {
                                    uint8_t optType = sigPayload[cfgOffset];
                                    uint8_t optLen = sigPayload[cfgOffset + 1];
                                    if (optType == 0x01 && optLen == 2 && cfgOffset + 4 <= sigPayloadLen) {
                                        g_channels[i].RemoteMtu = (uint16_t)sigPayload[cfgOffset + 2]
                                                                | ((uint16_t)sigPayload[cfgOffset + 3] << 8);
                                    }
                                    cfgOffset += 2 + optLen;
                                }

                                // Send Config Response (success)
                                uint8_t rsp[6] = {};
                                rsp[0] = (uint8_t)(g_channels[i].RemoteCid & 0xFF);
                                rsp[1] = (uint8_t)(g_channels[i].RemoteCid >> 8);
                                rsp[2] = 0; rsp[3] = 0;  // Flags
                                rsp[4] = 0; rsp[5] = 0;  // Result: success
                                SendSignal(SIG_CONFIG_RSP, sig->Identifier, rsp, 6);

                                if (g_channels[i].LocalConfigDone && g_channels[i].RemoteConfigDone) {
                                    g_channels[i].Configured = true;
                                    KernelLogStream(OK, "BT-L2CAP") << "Channel "
                                        << (uint64_t)g_channels[i].LocalCid << " configured";

                                    // Notify A2DP if this is an AVDTP channel
                                    if (g_channels[i].Psm == PSM_AVDTP) {
                                        A2dp::OnChannelReady(g_channels[i].LocalCid);
                                    }
                                }
                                break;
                            }
                        }
                    }
                    break;
                }

                case SIG_CONFIG_RSP: {
                    if (sigPayloadLen >= 6) {
                        // The Config Response's first field is the "Source CID".
                        // In a RECEIVED Config Response it echoes back OUR local
                        // channel endpoint (the device that sent the Config
                        // Request -- us), NOT the remote's CID -- exactly like the
                        // SIG_CONN_RSP handler above matches srcCid to LocalCid.
                        // Matching RemoteCid here was THE bug that pinned every
                        // AVDTP channel at localCfg=0: a success Config Response
                        // arrived but matched no channel, so LocalConfigDone never
                        // got set and the channel never became Configured.
                        uint16_t srcCid = (uint16_t)sigPayload[0] | ((uint16_t)sigPayload[1] << 8);
                        uint16_t result = (uint16_t)sigPayload[4] | ((uint16_t)sigPayload[5] << 8);

                        if (result != CFG_SUCCESS) {
                            KernelLogStream(WARNING, "BT-L2CAP") << "Config Response result="
                                << base::hex << (uint64_t)result << base::dec
                                << " (our config rejected)";
                        }

                        if (result == CFG_SUCCESS) {
                            bool matched = false;
                            for (int i = 0; i < MAX_CHANNELS; i++) {
                                if (g_channels[i].Active && g_channels[i].LocalCid == srcCid) {
                                    matched = true;
                                    g_channels[i].LocalConfigDone = true;
                                    if (g_channels[i].LocalConfigDone && g_channels[i].RemoteConfigDone) {
                                        g_channels[i].Configured = true;
                                        KernelLogStream(OK, "BT-L2CAP") << "Channel "
                                            << (uint64_t)g_channels[i].LocalCid << " configured";

                                        if (g_channels[i].Psm == PSM_AVDTP) {
                                            A2dp::OnChannelReady(g_channels[i].LocalCid);
                                        }
                                    }
                                    break;
                                }
                            }
                            if (!matched) {
                                KernelLogStream(WARNING, "BT-L2CAP")
                                    << "Config Response (success) for unknown CID="
                                    << base::hex << (uint64_t)srcCid << base::dec;
                            }
                        }
                    }
                    break;
                }

                case SIG_DISCONN_REQ: {
                    if (sigPayloadLen >= 4) {
                        uint16_t dstCid = (uint16_t)sigPayload[0] | ((uint16_t)sigPayload[1] << 8);
                        uint16_t srcCid = (uint16_t)sigPayload[2] | ((uint16_t)sigPayload[3] << 8);

                        for (int i = 0; i < MAX_CHANNELS; i++) {
                            if (g_channels[i].Active && g_channels[i].LocalCid == dstCid) {
                                g_channels[i].Active = false;
                                break;
                            }
                        }

                        // Send Disconnect Response
                        uint8_t rsp[4] = {};
                        rsp[0] = (uint8_t)(dstCid & 0xFF);
                        rsp[1] = (uint8_t)(dstCid >> 8);
                        rsp[2] = (uint8_t)(srcCid & 0xFF);
                        rsp[3] = (uint8_t)(srcCid >> 8);
                        SendSignal(SIG_DISCONN_RSP, sig->Identifier, rsp, 4);
                    }
                    break;
                }

                case SIG_INFO_REQ: {
                    if (sigPayloadLen >= 2) {
                        uint16_t infoType = (uint16_t)sigPayload[0] | ((uint16_t)sigPayload[1] << 8);

                        if (infoType == 0x0002) {
                            // Extended features mask
                            uint8_t rsp[8] = {};
                            rsp[0] = 0x02; rsp[1] = 0x00;  // InfoType
                            rsp[2] = 0x00; rsp[3] = 0x00;  // Result: success
                            rsp[4] = 0x00; rsp[5] = 0x00;  // Features: none
                            rsp[6] = 0x00; rsp[7] = 0x00;
                            SendSignal(SIG_INFO_RSP, sig->Identifier, rsp, 8);
                        } else {
                            // Not supported
                            uint8_t rsp[4] = {};
                            rsp[0] = (uint8_t)(infoType & 0xFF);
                            rsp[1] = (uint8_t)(infoType >> 8);
                            rsp[2] = 0x01; rsp[3] = 0x00;  // Result: not supported
                            SendSignal(SIG_INFO_RSP, sig->Identifier, rsp, 4);
                        }
                    }
                    break;
                }

                case SIG_DISCONN_RSP:
                    // Response to a Disconnect Request we sent (e.g. closing the
                    // SDP client channel after a query); nothing left to do.
                    break;

                case SIG_ECHO_REQ: {
                    // Echo Request must be answered (echo the payload back) or
                    // a peer probing the link times out on us.
                    SendSignal(SIG_ECHO_RSP, sig->Identifier, sigPayload, sigPayloadLen);
                    break;
                }

                default: {
                    // A signaling command we do not handle.  Log it: a peer that
                    // sends something exotic while holding our dials at
                    // "authorization pending" is exactly what we need to see.
                    KernelLogStream(INFO, "BT-L2CAP") << "Unhandled signaling code="
                        << base::hex << (uint64_t)sig->Code << base::dec
                        << " len=" << (uint64_t)sigPayloadLen;
                    break;
                }
              }
              sOff += sizeof(SignalHeader) + sigPayloadLen;
            }
        } else {
            // Data on a dynamic channel
            for (int i = 0; i < MAX_CHANNELS; i++) {
                if (g_channels[i].Active && g_channels[i].LocalCid == cid) {
                    if (g_channels[i].Psm == PSM_AVDTP) {
                        A2dp::ProcessAvdtp(payload, l2len);
                    } else if (g_channels[i].Psm == PSM_SDP) {
                        A2dp::ProcessSdp(cid, payload, l2len);
                    } else if (g_channels[i].Psm == PSM_AVCTP) {
                        Avrcp::ProcessAvctp(cid, payload, l2len);
                    }
                    break;
                }
            }
        }
    }

    // =========================================================================
    // Connect
    // =========================================================================

    uint16_t Connect(uint16_t psm) {
        if (!g_initialized) return 0;

        g_lastConnRspResult = 0xFFFF;   // reset diagnostic for this attempt

        auto* ch = AllocChannel(psm);
        if (!ch) return 0;

        // Send Connection Request
        uint8_t req[4] = {};
        req[0] = (uint8_t)(psm & 0xFF);
        req[1] = (uint8_t)(psm >> 8);
        req[2] = (uint8_t)(ch->LocalCid & 0xFF);
        req[3] = (uint8_t)(ch->LocalCid >> 8);
        SendSignal(SIG_CONN_REQ, g_sigIdentifier++, req, 4);

        return ch->LocalCid;
    }

    // =========================================================================
    // WaitConfigured
    // =========================================================================

    bool WaitConfigured(uint16_t localCid, uint32_t timeoutMs) {
        uint64_t start = Timekeeping::GetMilliseconds();

        while (Timekeeping::GetMilliseconds() - start < timeoutMs) {
            Xhci::PollEvents();
            Hci::DrainEvents();   // process incoming ACL (the Connection Response!)

            auto* ch = GetChannel(localCid);
            if (ch && ch->Configured) return true;

            for (int j = 0; j < 100; j++) {
                asm volatile("" ::: "memory");
            }
        }
        return false;
    }

    // =========================================================================
    // SendData
    // =========================================================================

    bool SendData(uint16_t localCid, const uint8_t* data, uint16_t len) {
        if (!g_initialized) return false;

        auto* ch = GetChannel(localCid);
        if (!ch || !ch->Configured) return false;

        // Build L2CAP packet
        uint16_t totalLen = sizeof(L2capHeader) + len;
        uint8_t buf[1024] = {};
        if (totalLen > sizeof(buf)) return false;

        auto* l2hdr = (L2capHeader*)buf;
        l2hdr->Length = len;
        l2hdr->ChannelId = ch->RemoteCid;

        if (data && len > 0) {
            memcpy(buf + sizeof(L2capHeader), data, len);
        }

        return Hci::SendAcl(g_aclHandle, Hci::ACL_PB_FIRST_FLUSH, buf, totalLen);
    }

    // =========================================================================
    // Channel queries
    // =========================================================================

    ChannelInfo* GetChannel(uint16_t localCid) {
        for (int i = 0; i < MAX_CHANNELS; i++) {
            if (g_channels[i].Active && g_channels[i].LocalCid == localCid) {
                return &g_channels[i];
            }
        }
        return nullptr;
    }

    bool FreeChannel(uint16_t localCid) {
        // Reclaim a channel slot from a failed/abandoned dial during A2DP
        // bring-up retries.  Two cases:
        //  - RemoteCid == 0: the remote never acknowledged our CONN_REQ (ignored,
        //    or answered only with a stuck PENDING that never completed), so its
        //    side was never created and no Disconnect is owed -- just free it.
        //  - RemoteCid != 0: a CONN_RSP success arrived (the peer allocated its
        //    CID) but the channel never finished configuring, so the peer holds a
        //    half-open channel.  Send it a Disconnect Request before freeing, or
        //    BOTH our slot leaks (the old RemoteCid==0-only guard left it Active)
        //    AND the peer keeps a dangling channel.
        // We deliberately do NOT reset g_nextCid, so a late response for this
        // (now freed) CID matches nothing rather than cross-wiring the retry.
        for (int i = 0; i < MAX_CHANNELS; i++) {
            if (g_channels[i].Active && g_channels[i].LocalCid == localCid) {
                if (g_channels[i].RemoteCid != 0) {
                    uint8_t req[4] = {};
                    req[0] = (uint8_t)(g_channels[i].RemoteCid & 0xFF);
                    req[1] = (uint8_t)(g_channels[i].RemoteCid >> 8);
                    req[2] = (uint8_t)(g_channels[i].LocalCid & 0xFF);
                    req[3] = (uint8_t)(g_channels[i].LocalCid >> 8);
                    SendSignal(SIG_DISCONN_REQ, g_sigIdentifier++, req, 4);
                }
                g_channels[i].Active = false;
                return true;
            }
        }
        return false;
    }

    ChannelInfo* FindChannelByPsm(uint16_t psm) {
        for (int i = 0; i < MAX_CHANNELS; i++) {
            if (g_channels[i].Active && g_channels[i].Psm == psm) {
                return &g_channels[i];
            }
        }
        return nullptr;
    }

    uint16_t FindConfiguredAvdtpChannelExcept(uint16_t exceptCid) {
        for (int i = 0; i < MAX_CHANNELS; i++) {
            if (g_channels[i].Active && g_channels[i].Configured
                && g_channels[i].Psm == PSM_AVDTP
                && g_channels[i].LocalCid != exceptCid) {
                return g_channels[i].LocalCid;
            }
        }
        return 0;
    }

    uint16_t GetAclHandle() {
        return g_aclHandle;
    }

    uint16_t LastConnRspResult() {
        return g_lastConnRspResult;
    }

}
