#include "udp_client.h"

#include <arpa/inet.h>
#include <cerrno>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <map>
#include <netdb.h>
#include <poll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <unistd.h>

UDPClient::UDPClient(std::string host, uint16_t port,
                     std::string multicast_group, uint16_t data_port)
    : host_(std::move(host))
    , port_(port)
    , multicast_group_(std::move(multicast_group))
    , data_port_(data_port)
{}

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

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

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

UDPClient::Stats UDPClient::stats() const {
    Stats s;
    s.packets_rx   = rx_packets_.load(std::memory_order_relaxed);
    s.packets_lost = seq_gaps_.load(std::memory_order_relaxed);
    s.bytes_rx     = rx_bytes_.load(std::memory_order_relaxed);
    s.connected    = connected_.load(std::memory_order_relaxed);
    return s;
}

// ─── Fragmentation reassembly ────────────────────────────────────────

namespace {

struct FragSet {
    uint32_t total = 0;
    uint32_t count = 0;
    std::vector<std::vector<uint8_t>> parts;
    std::chrono::steady_clock::time_point created;
};

using ReassemblyMap = std::map<uint64_t, FragSet>;

std::vector<uint8_t> reassemble(ReassemblyMap& m,
                                const UDPSPacketHeader& hdr,
                                const uint8_t* payload, size_t plen)
{
    if (hdr.total_fragments == 1 && hdr.fragment_idx == 0)
        return std::vector<uint8_t>(payload, payload + plen);

    uint64_t key = (uint64_t(hdr.counter) << 8) | hdr.type;
    auto& fs = m[key];
    if (fs.parts.empty()) {
        fs.total   = hdr.total_fragments;
        fs.parts.resize(fs.total);
        fs.created = std::chrono::steady_clock::now();
    }
    if (hdr.fragment_idx < fs.total && fs.parts[hdr.fragment_idx].empty()) {
        fs.parts[hdr.fragment_idx].assign(payload, payload + plen);
        ++fs.count;
    }
    if (fs.count < fs.total) return {};

    size_t sz = 0;
    for (auto& p : fs.parts) sz += p.size();
    std::vector<uint8_t> out;
    out.reserve(sz);
    for (auto& p : fs.parts) out.insert(out.end(), p.begin(), p.end());
    m.erase(key);
    return out;
}

// Drop fragments older than 2 seconds to avoid unbounded map growth
void gc_fragments(ReassemblyMap& m) {
    auto now  = std::chrono::steady_clock::now();
    auto it   = m.begin();
    while (it != m.end()) {
        if (std::chrono::duration<double>(now - it->second.created).count() > 2.0)
            it = m.erase(it);
        else
            ++it;
    }
}

} // anonymous namespace

// ─── Ingest thread ───────────────────────────────────────────────────

void UDPClient::run() {
    if (!multicast_group_.empty())
        run_multicast();
    else
        run_unicast();
}

void UDPClient::run_unicast() {
    constexpr double SILENCE_SEC  = 5.0;
    constexpr double CONNECT_RETRY= 1.0;   // re-send CONNECT if no CONFIG yet
    constexpr double RETRY_SEC    = 2.0;   // wait before full reconnect

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

    while (running_) {
        // ── Resolve server address ────────────────────────────────────
        addrinfo hints{};
        hints.ai_family   = AF_INET;
        hints.ai_socktype = SOCK_DGRAM;
        char port_str[8];
        std::snprintf(port_str, sizeof(port_str), "%u", port_);

        addrinfo* res = nullptr;
        if (getaddrinfo(host_.c_str(), port_str, &hints, &res) != 0 || !res) {
            std::fprintf(stderr, "[udp_client %s:%u] cannot resolve host\n",
                         host_.c_str(), port_);
            std::this_thread::sleep_for(std::chrono::duration<double>(RETRY_SEC));
            continue;
        }
        sockaddr_in server{};
        std::memcpy(&server, res->ai_addr, sizeof(server));
        freeaddrinfo(res);

        // ── Open socket ───────────────────────────────────────────────
        int sock = socket(AF_INET, SOCK_DGRAM, 0);
        if (sock < 0) {
            std::fprintf(stderr, "[udp_client %s:%u] socket: %s\n",
                         host_.c_str(), port_, std::strerror(errno));
            std::this_thread::sleep_for(std::chrono::duration<double>(RETRY_SEC));
            continue;
        }

        // Large receive buffer
        {
            int rcv = 8 * 1024 * 1024;
            setsockopt(sock, SOL_SOCKET, SO_RCVBUF, &rcv, sizeof(rcv));
        }

        // Bind to any local port so we have a stable source address
        {
            sockaddr_in local{};
            local.sin_family      = AF_INET;
            local.sin_port        = 0;          // OS picks port
            local.sin_addr.s_addr = INADDR_ANY;
            if (bind(sock, reinterpret_cast<sockaddr*>(&local), sizeof(local)) < 0) {
                std::fprintf(stderr, "[udp_client %s:%u] bind: %s\n",
                             host_.c_str(), port_, std::strerror(errno));
                close(sock);
                std::this_thread::sleep_for(std::chrono::duration<double>(RETRY_SEC));
                continue;
            }
        }

        // ── Send CONNECT to server ────────────────────────────────────
        {
            auto pkt = build_connect();
            sendto(sock, pkt.data(), pkt.size(), 0,
                   reinterpret_cast<sockaddr*>(&server), sizeof(server));
            std::fprintf(stderr, "[udp_client %s:%u] CONNECT sent\n",
                         host_.c_str(), port_);
        }

        // ── Receive loop (poll-based, non-blocking recv) ──────────────
        ReassemblyMap fragments;
        std::vector<SignalDef> signal_defs;
        uint32_t last_counter  = 0;
        bool     seen_first    = false;
        bool     config_rx     = false;
        using clock = std::chrono::steady_clock;
        auto     last_rx       = clock::now();
        auto     last_connect  = clock::now();
        uint32_t gc_counter    = 0;

        while (running_) {
            // poll() with 200 ms timeout — keeps reconnect logic responsive
            pollfd pfd{sock, POLLIN, 0};
            int    ready = poll(&pfd, 1, 200 /*ms*/);

            if (!running_) break;

            if (ready < 0) {
                if (errno == EINTR) continue;
                std::fprintf(stderr, "[udp_client %s:%u] poll: %s\n",
                             host_.c_str(), port_, std::strerror(errno));
                break;
            }

            auto now = clock::now();
            double elapsed = std::chrono::duration<double>(now - last_rx).count();

            if (ready == 0) {
                // Timeout — check timers
                if (elapsed > SILENCE_SEC) {
                    std::fprintf(stderr,
                        "[udp_client %s:%u] silence timeout, reconnecting\n",
                        host_.c_str(), port_);
                    connected_ = false;
                    break;
                }
                // If CONFIG not yet received, re-send CONNECT periodically
                if (!config_rx) {
                    double since_connect = std::chrono::duration<double>(
                        now - last_connect).count();
                    if (since_connect >= CONNECT_RETRY) {
                        auto pkt = build_connect();
                        sendto(sock, pkt.data(), pkt.size(), 0,
                               reinterpret_cast<sockaddr*>(&server), sizeof(server));
                        last_connect = now;
                    }
                }
                continue;
            }

            // ── Read one datagram ─────────────────────────────────────
            ssize_t n = recvfrom(sock, buf, BUF_SZ, 0, nullptr, nullptr);
            if (n < 0) {
                if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
                    continue;
                std::fprintf(stderr, "[udp_client %s:%u] recv: %s\n",
                             host_.c_str(), port_, std::strerror(errno));
                break;
            }

            last_rx = clock::now();
            rx_packets_++;
            rx_bytes_ += static_cast<uint64_t>(n);

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

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

            auto complete = reassemble(fragments, hdr, payload, plen);
            if (complete.empty()) continue;

            // Periodic GC of stale fragment sets
            if (++gc_counter % 1000 == 0) gc_fragments(fragments);

            switch (hdr.type) {
            case UDPS_TYPE_CONFIG:
                signal_defs = parse_config(complete.data(), complete.size());
                config_rx   = true;
                connected_  = true;
                std::fprintf(stderr, "[udp_client %s:%u] CONFIG: %zu signals\n",
                             host_.c_str(), port_, signal_defs.size());
                if (on_config) on_config(signal_defs);
                break;

            case UDPS_TYPE_DATA: {
                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);
                if (on_data) on_data(dp, signal_defs);
                break;
            }

            case UDPS_TYPE_DISCONNECT:
                std::fprintf(stderr, "[udp_client %s:%u] server DISCONNECT\n",
                             host_.c_str(), port_);
                connected_ = false;
                goto reconnect;

            default: break;
            }
        }

        reconnect:
        {
            auto pkt = build_disconnect();
            sendto(sock, pkt.data(), pkt.size(), 0,
                   reinterpret_cast<sockaddr*>(&server), sizeof(server));
        }
        close(sock);

        if (running_)
            std::this_thread::sleep_for(std::chrono::duration<double>(RETRY_SEC));
    }

    std::fprintf(stderr, "[udp_client %s:%u] stopped\n", host_.c_str(), port_);
}

// ─── Multicast ingest thread ──────────────────────────────────────────

/** Returns true if the session ran normally (even if the server disconnected),
 *  false if a fatal socket error occurred before data could be received. */
bool UDPClient::run_multicast_session() {
    constexpr double SILENCE_SEC = 5.0;
    constexpr size_t BUF_SZ = 65536;
    uint8_t buf[BUF_SZ];

    // ── Resolve server address ────────────────────────────────────────
    addrinfo hints{};
    hints.ai_family   = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    char port_str[8];
    std::snprintf(port_str, sizeof(port_str), "%u", port_);
    addrinfo* res = nullptr;
    if (getaddrinfo(host_.c_str(), port_str, &hints, &res) != 0 || !res) {
        std::fprintf(stderr, "[udp_client %s:%u] cannot resolve host\n",
                     host_.c_str(), port_);
        return false;
    }
    sockaddr_in server{};
    std::memcpy(&server, res->ai_addr, sizeof(server));
    freeaddrinfo(res);

    // ── Open and connect TCP control socket ───────────────────────────
    int tcp_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (tcp_sock < 0) return false;
    {
        int one = 1;
        setsockopt(tcp_sock, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
    }
    if (connect(tcp_sock, reinterpret_cast<sockaddr*>(&server), sizeof(server)) < 0) {
        std::fprintf(stderr, "[udp_client %s:%u] tcp connect: %s\n",
                     host_.c_str(), port_, std::strerror(errno));
        close(tcp_sock);
        return false;
    }

    // RAII wrapper so tcp_sock is always closed on return
    struct TcpGuard {
        int& fd;
        std::function<void()> on_close;
        ~TcpGuard() {
            if (on_close) on_close();
            close(fd);
        }
    } tcp_guard{tcp_sock, [&]() {
        auto pkt = build_disconnect();
        send(tcp_sock, pkt.data(), pkt.size(), 0);
    }};

    // ── Send CONNECT via TCP ──────────────────────────────────────────
    {
        auto pkt = build_connect();
        if (send(tcp_sock, pkt.data(), pkt.size(), 0) < 0) {
            std::fprintf(stderr, "[udp_client %s:%u] tcp send CONNECT: %s\n",
                         host_.c_str(), port_, std::strerror(errno));
            return false;
        }
        std::fprintf(stderr, "[udp_client %s:%u] tcp: CONNECT sent\n",
                     host_.c_str(), port_);
    }

    // ── Read CONFIG header via TCP ────────────────────────────────────
    uint8_t hdr_buf[UDPS_HEADER_SIZE];
    {
        size_t total = 0;
        while (total < UDPS_HEADER_SIZE) {
            ssize_t n = recv(tcp_sock, hdr_buf + total, UDPS_HEADER_SIZE - total, 0);
            if (n <= 0) return false;
            total += static_cast<size_t>(n);
        }
    }
    UDPSPacketHeader cfg_hdr{};
    if (!parse_header(hdr_buf, UDPS_HEADER_SIZE, cfg_hdr) ||
        cfg_hdr.type != UDPS_TYPE_CONFIG) {
        std::fprintf(stderr, "[udp_client %s:%u] expected CONFIG header\n",
                     host_.c_str(), port_);
        return false;
    }

    // ── Read CONFIG payload via TCP ───────────────────────────────────
    std::vector<uint8_t> cfg_payload(cfg_hdr.payload_bytes);
    {
        size_t total = 0;
        while (total < cfg_hdr.payload_bytes) {
            ssize_t n = recv(tcp_sock, cfg_payload.data() + total,
                             cfg_hdr.payload_bytes - total, 0);
            if (n <= 0) return false;
            total += static_cast<size_t>(n);
        }
    }
    auto signal_defs = parse_config(cfg_payload.data(), cfg_payload.size());
    std::fprintf(stderr, "[udp_client %s:%u] tcp: CONFIG (%zu signals)\n",
                 host_.c_str(), port_, signal_defs.size());
    connected_ = true;
    if (on_config) on_config(signal_defs);

    // ── Determine multicast data port ─────────────────────────────────
    uint16_t mcast_port = (data_port_ > 0) ? data_port_
                                            : static_cast<uint16_t>(port_ + 1u);

    // ── Open UDP socket and join multicast group ──────────────────────
    int udp_sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (udp_sock < 0) return false;

    // RAII for udp_sock
    struct UdpGuard {
        int& fd;
        std::string& mgroup;
        ~UdpGuard() {
            ip_mreq mreq{};
            mreq.imr_multiaddr.s_addr = inet_addr(mgroup.c_str());
            mreq.imr_interface.s_addr = INADDR_ANY;
            setsockopt(fd, IPPROTO_IP, IP_DROP_MEMBERSHIP, &mreq, sizeof(mreq));
            close(fd);
        }
    } udp_guard{udp_sock, multicast_group_};

    {
        int reuse = 1;
        setsockopt(udp_sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse));
        int rcv = 8 * 1024 * 1024;
        setsockopt(udp_sock, SOL_SOCKET, SO_RCVBUF, &rcv, sizeof(rcv));
    }
    {
        sockaddr_in mcast_addr{};
        mcast_addr.sin_family      = AF_INET;
        mcast_addr.sin_port        = htons(mcast_port);
        mcast_addr.sin_addr.s_addr = INADDR_ANY;
        if (bind(udp_sock, reinterpret_cast<sockaddr*>(&mcast_addr),
                 sizeof(mcast_addr)) < 0) {
            std::fprintf(stderr, "[udp_client %s:%u] udp bind: %s\n",
                         host_.c_str(), port_, std::strerror(errno));
            return false;
        }
    }
    {
        ip_mreq mreq{};
        mreq.imr_multiaddr.s_addr = inet_addr(multicast_group_.c_str());
        mreq.imr_interface.s_addr = INADDR_ANY;
        if (setsockopt(udp_sock, IPPROTO_IP, IP_ADD_MEMBERSHIP,
                       &mreq, sizeof(mreq)) < 0) {
            std::fprintf(stderr, "[udp_client %s:%u] IP_ADD_MEMBERSHIP(%s): %s\n",
                         host_.c_str(), port_,
                         multicast_group_.c_str(), std::strerror(errno));
            return false;
        }
    }
    std::fprintf(stderr, "[udp_client %s:%u] joined multicast %s:%u\n",
                 host_.c_str(), port_, multicast_group_.c_str(), mcast_port);

    // ── Main DATA loop ────────────────────────────────────────────────
    ReassemblyMap fragments;
    uint32_t last_counter = 0;
    bool     seen_first   = false;
    uint32_t gc_counter   = 0;
    using clock = std::chrono::steady_clock;
    auto last_rx = clock::now();

    while (running_) {
        pollfd pfds[2];
        pfds[0] = {udp_sock, POLLIN, 0};
        pfds[1] = {tcp_sock, POLLIN, 0};
        int ready = poll(pfds, 2, 200 /*ms*/);

        if (!running_) break;
        if (ready < 0) {
            if (errno == EINTR) continue;
            break;
        }

        // TCP event: DISCONNECT or closed connection
        if (pfds[1].revents & (POLLIN | POLLHUP | POLLERR)) {
            ssize_t n = recv(tcp_sock, buf, BUF_SZ, MSG_DONTWAIT);
            if (n <= 0) {
                std::fprintf(stderr, "[udp_client %s:%u] tcp control closed\n",
                             host_.c_str(), port_);
                break;
            }
            UDPSPacketHeader ctrl_hdr{};
            if (parse_header(buf, static_cast<size_t>(n), ctrl_hdr) &&
                ctrl_hdr.type == UDPS_TYPE_DISCONNECT) {
                std::fprintf(stderr, "[udp_client %s:%u] server DISCONNECT\n",
                             host_.c_str(), port_);
                break;
            }
        }

        // Silence timeout on UDP DATA
        {
            auto now = clock::now();
            double elapsed = std::chrono::duration<double>(now - last_rx).count();
            if (!(pfds[0].revents & POLLIN)) {
                if (elapsed > SILENCE_SEC) {
                    std::fprintf(stderr,
                        "[udp_client %s:%u] multicast silence timeout\n",
                        host_.c_str(), port_);
                    break;
                }
                continue;
            }
        }

        ssize_t n = recvfrom(udp_sock, buf, BUF_SZ, 0, nullptr, nullptr);
        if (n < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) continue;
            std::fprintf(stderr, "[udp_client %s:%u] udp recv: %s\n",
                         host_.c_str(), port_, std::strerror(errno));
            break;
        }

        last_rx = clock::now();
        rx_packets_++;
        rx_bytes_ += static_cast<uint64_t>(n);

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

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

        auto complete = reassemble(fragments, hdr, payload, plen);
        if (complete.empty()) continue;

        if (++gc_counter % 1000 == 0) gc_fragments(fragments);

        if (hdr.type == UDPS_TYPE_DATA) {
            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()) {
                auto dp = parse_data(complete.data(), complete.size(), signal_defs);
                if (on_data) on_data(dp, signal_defs);
            }
        }
    }

    connected_ = false;
    return true;
}

void UDPClient::run_multicast() {
    constexpr double RETRY_SEC = 2.0;
    while (running_) {
        connected_ = false;
        run_multicast_session();
        connected_ = false;
        if (running_)
            std::this_thread::sleep_for(std::chrono::duration<double>(RETRY_SEC));
    }
    std::fprintf(stderr, "[udp_client %s:%u] multicast stopped\n",
                 host_.c_str(), port_);
}