#include "LwIPUDP.hpp"
#include "pico/cyw43_arch.h"
#include <string.h>

LwIPUDP::LwIPUDP()
    : pcb_(nullptr), rxHead_(0), rxTail_(0),
      current_(nullptr), txPort_(0), txBuf_(nullptr)
{
    memset(rxQueue_, 0, sizeof(rxQueue_));
    ip_addr_set_zero(&txAddr_);
}

LwIPUDP::~LwIPUDP() {
    stop();
}

bool LwIPUDP::begin(uint16_t port) {
    if (pcb_) stop();

    cyw43_arch_lwip_begin();
    pcb_ = udp_new();
    if (!pcb_) {
        cyw43_arch_lwip_end();
        return false;
    }

    if (udp_bind(pcb_, IP_ADDR_ANY, port) != ERR_OK) {
        udp_remove(pcb_);
        pcb_ = nullptr;
        cyw43_arch_lwip_end();
        return false;
    }

    udp_recv(pcb_, recvCb, this);
    cyw43_arch_lwip_end();
    return true;
}

void LwIPUDP::stop() {
    if (pcb_) {
        cyw43_arch_lwip_begin();
        udp_remove(pcb_);
        cyw43_arch_lwip_end();
        pcb_ = nullptr;
    }

    // Drain rx queue
    while (rxTail_ != rxHead_) {
        auto &pkt = rxQueue_[rxTail_ % RX_QUEUE_SIZE];
        if (pkt.p) { pbuf_free(pkt.p); pkt.p = nullptr; }
        rxTail_++;
    }
    rxHead_ = rxTail_ = 0;
    current_ = nullptr;

    if (txBuf_) { pbuf_free(txBuf_); txBuf_ = nullptr; }
}

void LwIPUDP::recvCb(void *arg, struct udp_pcb *, struct pbuf *p,
                      const ip_addr_t *addr, u16_t port) {
    auto *self = static_cast<LwIPUDP *>(arg);
    uint8_t next = (self->rxHead_ + 1) % RX_QUEUE_SIZE;
    if (next == self->rxTail_ % RX_QUEUE_SIZE) {
        // Queue full — drop
        pbuf_free(p);
        return;
    }
    auto &slot = self->rxQueue_[self->rxHead_ % RX_QUEUE_SIZE];
    slot.addr = *addr;
    slot.port = port;
    slot.p = p;
    slot.offset = 0;
    self->rxHead_ = next;
}

void LwIPUDP::advanceRx() {
    if (current_) {
        if (current_->p) { pbuf_free(current_->p); current_->p = nullptr; }
        rxTail_ = (rxTail_ + 1) % RX_QUEUE_SIZE;
        current_ = nullptr;
    }
}

int LwIPUDP::available() {
    if (current_ && current_->p)
        return current_->p->tot_len - current_->offset;
    return 0;
}

int LwIPUDP::parsePacket() {
    // Drop any unfinished current packet
    advanceRx();

    if (rxTail_ == rxHead_) return 0;

    current_ = &rxQueue_[rxTail_ % RX_QUEUE_SIZE];
    return current_->p ? current_->p->tot_len : 0;
}

int LwIPUDP::read(uint8_t *buf, size_t len) {
    if (!current_ || !current_->p) return 0;

    uint16_t remaining = current_->p->tot_len - current_->offset;
    uint16_t toRead = (len < remaining) ? (uint16_t)len : remaining;
    uint16_t copied = pbuf_copy_partial(current_->p, buf, toRead, current_->offset);
    current_->offset += copied;
    return copied;
}

int LwIPUDP::read() {
    uint8_t b;
    return (read(&b, 1) == 1) ? b : -1;
}

IPAddress LwIPUDP::remoteIP() {
    if (current_) return IPAddress(current_->addr);
    return IPAddress();
}

uint16_t LwIPUDP::remotePort() {
    if (current_) return current_->port;
    return 0;
}

int LwIPUDP::beginPacket(IPAddress ip, uint16_t port) {
    if (txBuf_) { pbuf_free(txBuf_); txBuf_ = nullptr; }
    txAddr_ = ip.toLwIP();
    txPort_ = port;
    return 1;
}

size_t LwIPUDP::write(const uint8_t *buf, size_t len) {
    if (!len) return 0;

    struct pbuf *seg = pbuf_alloc(PBUF_TRANSPORT, len, PBUF_RAM);
    if (!seg) return 0;
    memcpy(seg->payload, buf, len);

    if (txBuf_) {
        pbuf_cat(txBuf_, seg);
    } else {
        txBuf_ = seg;
    }
    return len;
}

int LwIPUDP::endPacket() {
    if (!pcb_ || !txBuf_) return 0;

    cyw43_arch_lwip_begin();
    err_t err = udp_sendto(pcb_, txBuf_, &txAddr_, txPort_);
    cyw43_arch_lwip_end();
    pbuf_free(txBuf_);
    txBuf_ = nullptr;
    return (err == ERR_OK) ? 1 : 0;
}