#include "udp_receiver.h"
#include <arpa/inet.h>
#include <cerrno>
#include <cstdio>
#include <cstring>
#include <map>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>

UDPReceiver::UDPReceiver(uint16_t port, const std::string& bind_addr,
                         std::shared_ptr<RingBuf> fullrate,
                         std::shared_ptr<RingBuf> overview)
    : port_(port), bind_addr_(bind_addr),
      fullrate_(std::move(fullrate)), overview_(std::move(overview))
{}

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

void UDPReceiver::start() {
    running_ = true;
    thread_ = std::thread(&UDPReceiver::run, this);
}

void UDPReceiver::stop() {
    running_ = false;
    if (thread_.joinable()) thread_.join();
}

UDPReceiver::Stats UDPReceiver::stats() const {
    Stats s;
    s.datagrams_rx = rx_datagrams_;
    s.bytes_rx     = rx_bytes_;
    s.seq_gaps     = seq_gaps_;
    s.overflows    = fullrate_ ? fullrate_->overflow_count() : 0;
    s.connected    = connected_;
    return s;
}

void UDPReceiver::run() {
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        std::fprintf(stderr, "[udp_rx] socket: %s\n", std::strerror(errno));
        return;
    }

    // Large receive buffer
    int rcvbuf = 8 * 1024 * 1024;
    setsockopt(sock, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf));
    socklen_t optlen = sizeof(rcvbuf);
    getsockopt(sock, SOL_SOCKET, SO_RCVBUF, &rcvbuf, &optlen);
    std::fprintf(stderr, "[udp_rx:%u] SO_RCVBUF = %d\n", port_, rcvbuf);

    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(port_);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(sock, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) < 0) {
        std::fprintf(stderr, "[udp_rx:%u] bind: %s\n", port_, std::strerror(errno));
        close(sock);
        return;
    }

    // Send CONNECT — this triggers the MARTe2 streamer to start sending
    auto connect_pkt = build_connect();
    sendto(sock, connect_pkt.data(), connect_pkt.size(), 0,
           reinterpret_cast<sockaddr*>(&addr), sizeof(addr));

    std::fprintf(stderr, "[udp_rx:%u] listening, CONNECT sent\n", port_);

    // Fragmentation reassembly state
    struct FragSet {
        uint32_t        total = 0;
        uint32_t        count = 0;
        std::vector<std::vector<uint8_t>> parts;
    };
    std::map<uint64_t, FragSet> reassembler;  // key = counter<<8 | type
    uint32_t last_counter = 0;
    bool     seen_first   = false;
    std::vector<SignalDef> signal_defs;

    constexpr size_t BUF_SZ = 65536;
    uint8_t buf[BUF_SZ];

    while (running_) {
        ssize_t n = recvfrom(sock, buf, BUF_SZ, 0, nullptr, nullptr);
        if (n < 0) {
            if (errno == EAGAIN || errno == EINTR) continue;
            std::fprintf(stderr, "[udp_rx:%u] recv: %s\n", port_, std::strerror(errno));
            break;
        }

        rx_datagrams_++;
        rx_bytes_ += static_cast<uint64_t>(n);

        UDPSPacketHeader hdr;
        if (!parse_header(buf, static_cast<size_t>(n), hdr)) continue;

        size_t payload_len = static_cast<size_t>(n) - UDPS_HEADER_SIZE;
        const uint8_t* payload = buf + UDPS_HEADER_SIZE;

        // Reassembly
        std::vector<uint8_t> complete;
        if (hdr.total_fragments == 1 && hdr.fragment_idx == 0) {
            complete.assign(payload, payload + payload_len);
        } else {
            uint64_t key = (uint64_t(hdr.counter) << 8) | hdr.type;
            auto& fs = reassembler[key];
            if (fs.parts.empty()) {
                fs.total = hdr.total_fragments;
                fs.parts.resize(fs.total);
            }
            if (hdr.fragment_idx < fs.total && fs.parts[hdr.fragment_idx].empty()) {
                fs.parts[hdr.fragment_idx].assign(payload, payload + payload_len);
                fs.count++;
            }
            if (fs.count < fs.total) continue;

            // Concatenate
            size_t total_sz = 0;
            for (auto& p : fs.parts) total_sz += p.size();
            complete.reserve(total_sz);
            for (auto& p : fs.parts) complete.insert(complete.end(), p.begin(), p.end());
            reassembler.erase(key);
        }

        // Process complete packet
        switch (hdr.type) {
        case UDPS_TYPE_CONFIG: {
            signal_defs = parse_config(complete.data(), complete.size());
            connected_ = true;
            std::fprintf(stderr, "[udp_rx:%u] CONFIG: %zu signals\n",
                         port_, signal_defs.size());
            break;
        }
        case UDPS_TYPE_DATA: {
            // Sequence gap detection
            if (seen_first && hdr.counter != last_counter + 1) {
                seq_gaps_ += (hdr.counter - last_counter - 1);
            }
            seen_first = true;
            last_counter = hdr.counter;

            if (signal_defs.empty()) break;

            auto dp = parse_data(complete.data(), complete.size(), signal_defs);

            // Write only the first element of each signal to the full-rate buffer
            // (scalar / PacketTime mode). Temporal arrays with per-sample timestamps
            // require the full-array time mode which writes all elements.
            for (size_t si = 0; si < signal_defs.size(); si++) {
                const auto& sig = signal_defs[si];
                const auto& vals = dp.values[si];

                if (sig.time_mode == TM_FULL_ARRAY && sig.num_elements > 1) {
                    // Full array mode — write all elements
                    for (uint32_t i = 0; i < sig.num_elements && i < vals.size(); i++) {
                        fullrate_->write(static_cast<float>(vals[i]));
                    }
                } else if (!vals.empty()) {
                    fullrate_->write(static_cast<float>(vals[0]));
                }
            }
            break;
        }
        case UDPS_TYPE_DISCONNECT:
            std::fprintf(stderr, "[udp_rx:%u] server sent DISCONNECT\n", port_);
            connected_ = false;
            break;
        default:
            break;
        }
    }

    auto disc = build_disconnect();
    sendto(sock, disc.data(), disc.size(), 0,
           reinterpret_cast<sockaddr*>(&addr), sizeof(addr));
    close(sock);
    std::fprintf(stderr, "[udp_rx:%u] stopped\n", port_);
}
