#include <arpa/inet.h>
#include <cstring>
#include <fstream>
#include <ifaddrs.h>
#include <iostream>
#include <netdb.h>
#include <netinet/in.h>
#include <string>
#include <sys/socket.h>
#include <unistd.h>
#include <vector>

constexpr const char* GROUP_IP = "239.168.123.161";
constexpr int PORT = 5555;

constexpr int SAMPLE_RATE = 16000;
constexpr int SECONDS = 5;
constexpr int PACKET_BYTES = 5120;
constexpr int TOTAL_BYTES =
    SAMPLE_RATE * 2 * SECONDS;

static std::string get_local_ip(
    const std::string& interface)
{
    ifaddrs* ifaddr = nullptr;

    if (getifaddrs(&ifaddr) != 0)
        return {};

    std::string result;

    for (ifaddrs* ifa = ifaddr;
         ifa;
         ifa = ifa->ifa_next) {

        if (!ifa->ifa_addr)
            continue;

        if (interface != ifa->ifa_name)
            continue;

        if (ifa->ifa_addr->sa_family != AF_INET)
            continue;

        char host[NI_MAXHOST]{};

        if (getnameinfo(
                ifa->ifa_addr,
                sizeof(sockaddr_in),
                host,
                sizeof(host),
                nullptr,
                0,
                NI_NUMERICHOST) == 0) {

            result = host;
            break;
        }
    }

    freeifaddrs(ifaddr);
    return result;
}

static void write_wav(
    const std::string& path,
    const std::vector<int16_t>& pcm)
{
    std::ofstream f(
        path,
        std::ios::binary
    );

    const uint32_t data_size =
        pcm.size() * sizeof(int16_t);

    const uint32_t riff_size =
        36 + data_size;

    const uint16_t format = 1;
    const uint16_t channels = 1;
    const uint16_t bits = 16;

    const uint32_t sample_rate = 16000;

    const uint16_t block_align =
        channels * bits / 8;

    const uint32_t byte_rate =
        sample_rate * block_align;

    f.write("RIFF", 4);
    f.write(
        reinterpret_cast<const char*>(&riff_size),
        4
    );

    f.write("WAVE", 4);
    f.write("fmt ", 4);

    const uint32_t fmt_size = 16;

    f.write(
        reinterpret_cast<const char*>(&fmt_size),
        4
    );

    f.write(
        reinterpret_cast<const char*>(&format),
        2
    );

    f.write(
        reinterpret_cast<const char*>(&channels),
        2
    );

    f.write(
        reinterpret_cast<const char*>(&sample_rate),
        4
    );

    f.write(
        reinterpret_cast<const char*>(&byte_rate),
        4
    );

    f.write(
        reinterpret_cast<const char*>(&block_align),
        2
    );

    f.write(
        reinterpret_cast<const char*>(&bits),
        2
    );

    f.write("data", 4);

    f.write(
        reinterpret_cast<const char*>(&data_size),
        4
    );

    f.write(
        reinterpret_cast<const char*>(pcm.data()),
        data_size
    );
}

int main(int argc, char** argv)
{
    if (argc != 2) {
        std::cerr
            << "Usage: "
            << argv[0]
            << " <interface>\n";

        return 1;
    }

    const std::string interface = argv[1];

    const std::string local_ip =
        get_local_ip(interface);

    std::cout
        << "Interface: "
        << interface
        << '\n'
        << "Local IP: "
        << local_ip
        << '\n';

    int sock =
        socket(AF_INET, SOCK_DGRAM, 0);

    if (sock < 0) {
        perror("socket");
        return 1;
    }

    int reuse = 1;

    setsockopt(
        sock,
        SOL_SOCKET,
        SO_REUSEADDR,
        &reuse,
        sizeof(reuse)
    );

    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) {

        perror("bind");
        close(sock);
        return 1;
    }

    ip_mreq mreq{};

    inet_pton(
        AF_INET,
        GROUP_IP,
        &mreq.imr_multiaddr
    );

    inet_pton(
        AF_INET,
        local_ip.c_str(),
        &mreq.imr_interface
    );

    if (setsockopt(
            sock,
            IPPROTO_IP,
            IP_ADD_MEMBERSHIP,
            &mreq,
            sizeof(mreq)) != 0) {

        perror("IP_ADD_MEMBERSHIP");
        close(sock);
        return 1;
    }

    std::vector<int16_t> pcm;
    pcm.reserve(SAMPLE_RATE * SECONDS);

    int total_bytes = 0;

    std::cout
        << "\nSpeak for 5 seconds...\n";

    while (total_bytes < TOTAL_BYTES) {

        char buffer[PACKET_BYTES];

        const ssize_t len =
            recvfrom(
                sock,
                buffer,
                sizeof(buffer),
                0,
                nullptr,
                nullptr
            );

        if (len <= 0)
            continue;

        const size_t samples =
            len / sizeof(int16_t);

        const int16_t* p =
            reinterpret_cast<
                const int16_t*
            >(buffer);

        pcm.insert(
            pcm.end(),
            p,
            p + samples
        );

        total_bytes += len;

        std::cout
            << "\r"
            << total_bytes
            << " / "
            << TOTAL_BYTES
            << " bytes"
            << std::flush;
    }

    std::cout << '\n';

    write_wav(
        "/tmp/r1_raw_test.wav",
        pcm
    );

    std::cout
        << "Written /tmp/r1_raw_test.wav\n";

    setsockopt(
        sock,
        IPPROTO_IP,
        IP_DROP_MEMBERSHIP,
        &mreq,
        sizeof(mreq)
    );

    close(sock);

    return 0;
}
