#ifndef TCP_TEST_UTIL_H_
#define TCP_TEST_UTIL_H_

/**
 * @file TCPTestUtil.h
 * @brief Raw-socket helpers shared by the TCPS* GTests (C++98, STL allowed).
 */

#include <arpa/inet.h>
#include <errno.h>
#include <netinet/in.h>
#include <poll.h>
#include <string.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#include <vector>

#include "CompilerTypes.h"
#include "UDPSProtocol.h"

namespace tcptest {

using namespace MARTe;

inline uint64 NowMs() {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return static_cast<uint64>(ts.tv_sec) * 1000u + static_cast<uint64>(ts.tv_nsec) / 1000000u;
}

/** Connect a blocking socket to 127.0.0.1:port. @return fd, or -1. */
inline int Connect(uint16 port, int rcvBuf = 0) {
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        return -1;
    }
    if (rcvBuf > 0) {
        (void)setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &rcvBuf, sizeof(rcvBuf));
    }
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    if (connect(fd, reinterpret_cast<struct sockaddr *>(&addr), sizeof(addr)) != 0) {
        close(fd);
        return -1;
    }
    return fd;
}

/** Read exactly @p n bytes within @p timeoutMs. */
inline bool ReadExact(int fd, uint8 *dst, size_t n, int timeoutMs) {
    const uint64 deadline = NowMs() + static_cast<uint64>(timeoutMs);
    size_t got = 0u;
    while (got < n) {
        int64 left = static_cast<int64>(deadline) - static_cast<int64>(NowMs());
        if (left < 0) {
            return false;
        }
        struct pollfd p;
        p.fd = fd;
        p.events = POLLIN;
        p.revents = 0;
        int r = poll(&p, 1, static_cast<int>(left));
        if (r <= 0) {
            return false;
        }
        ssize_t k = recv(fd, dst + got, n - got, 0);
        if (k <= 0) {
            return false;
        }
        got += static_cast<size_t>(k);
    }
    return true;
}

struct Frame {
    uint8 type;
    uint32 counter;
    uint16 fragmentIdx;
    uint16 totalFragments;
    std::vector<uint8> payload;
};

/** Read one UDPS frame. @return false on timeout, EOF or a bad header. */
inline bool ReadFrame(int fd, Frame &f, int timeoutMs) {
    uint8 hb[UDPS_HEADER_SIZE];
    if (!ReadExact(fd, hb, UDPS_HEADER_SIZE, timeoutMs)) {
        return false;
    }
    UDPSPacketHeader h;
    memcpy(&h, hb, UDPS_HEADER_SIZE);
    if (h.magic != UDPS_MAGIC) {
        return false;
    }
    f.type = h.type;
    f.counter = h.counter;
    f.fragmentIdx = h.fragmentIdx;
    f.totalFragments = h.totalFragments;
    f.payload.resize(h.payloadBytes);
    if (h.payloadBytes > 0u) {
        return ReadExact(fd, &f.payload[0], h.payloadBytes, timeoutMs);
    }
    return true;
}

/** @return true if the peer closed the connection within @p timeoutMs. */
inline bool PeerClosed(int fd, int timeoutMs) {
    struct pollfd p;
    p.fd = fd;
    p.events = POLLIN;
    p.revents = 0;
    if (poll(&p, 1, timeoutMs) <= 0) {
        return false;
    }
    uint8 b;
    ssize_t k = recv(fd, &b, 1, MSG_DONTWAIT);
    return (k == 0) || ((k < 0) && (errno != EAGAIN));
}

/** Deterministic payload pattern. */
inline std::vector<uint8> Pattern(uint32 n, uint8 seed) {
    std::vector<uint8> v(n);
    for (uint32 i = 0u; i < n; i++) {
        v[i] = static_cast<uint8>(seed + i);
    }
    return v;
}


/** Listen on 127.0.0.1:@p port (0 = ephemeral). @return fd, or -1; @p port is updated. */
inline int Listen(uint16 &port) {
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        return -1;
    }
    int one = 1;
    (void)setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    socklen_t len = sizeof(addr);
    if ((bind(fd, reinterpret_cast<struct sockaddr *>(&addr), sizeof(addr)) != 0) ||
        (listen(fd, 8) != 0) ||
        (getsockname(fd, reinterpret_cast<struct sockaddr *>(&addr), &len) != 0)) {
        close(fd);
        return -1;
    }
    port = ntohs(addr.sin_port);
    return fd;
}

/** Accept one connection within @p timeoutMs. @return fd, or -1. */
inline int AcceptOne(int listenFd, int timeoutMs) {
    struct pollfd p;
    p.fd = listenFd;
    p.events = POLLIN;
    p.revents = 0;
    if (poll(&p, 1, timeoutMs) <= 0) {
        return -1;
    }
    return accept(listenFd, NULL, NULL);
}

inline bool SendAll(int fd, const uint8 *data, size_t n) {
    size_t sent = 0u;
    while (sent < n) {
        ssize_t k = send(fd, data + sent, n - sent, MSG_NOSIGNAL);
        if (k <= 0) {
            return false;
        }
        sent += static_cast<size_t>(k);
    }
    return true;
}

/** Build a complete UDPS frame (header + payload). */
inline std::vector<uint8> MakeFrame(uint8 type, uint32 counter, const std::vector<uint8> &payload,
                                    uint16 totalFragments = 1u, uint16 fragmentIdx = 0u) {
    std::vector<uint8> f(UDPS_HEADER_SIZE + payload.size());
    UDPSBuildHeader(&f[0], type, counter, fragmentIdx, totalFragments, static_cast<uint32>(payload.size()));
    if (!payload.empty()) {
        memcpy(&f[UDPS_HEADER_SIZE], &payload[0], payload.size());
    }
    return f;
}

} // namespace tcptest

#endif // TCP_TEST_UTIL_H_