#include "r1_voice/stt.hh"

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

namespace r1_voice {

namespace {

constexpr const char* GROUP_IP =
    "239.168.123.161";

constexpr uint16_t PORT = 5555;

constexpr int SAMPLE_RATE = 16000;

std::string findR1InterfaceAddress(
    const std::string& interface_name)
{
    struct ifaddrs* ifaddr = nullptr;

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

    std::string result;

    for (auto* ifa = ifaddr;
         ifa != nullptr;
         ifa = ifa->ifa_next) {

        if (!ifa->ifa_addr)
            continue;

        if (interface_name != 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;
}

float rms(
    const int16_t* samples,
    size_t count)
{
    if (count == 0)
        return 0.0f;

    double sum = 0.0;

    for (size_t i = 0; i < count; ++i) {
        const float sample =
            static_cast<float>(samples[i]) /
            32768.0f;

        sum +=
            static_cast<double>(sample) *
            static_cast<double>(sample);
    }

    return static_cast<float>(
        std::sqrt(sum / count)
    );
}

void writeWav(
    const std::string& filename,
    const std::vector<int16_t>& pcm)
{
    std::ofstream file(
        filename,
        std::ios::binary
    );

    if (!file)
        throw std::runtime_error(
            "Cannot open WAV file: " + filename
        );

    const uint16_t audio_format = 1;
    const uint16_t channels = 1;
    const uint16_t bits_per_sample = 16;

    const uint32_t sample_rate = SAMPLE_RATE;

    const uint16_t block_align =
        channels * bits_per_sample / 8;

    const uint32_t byte_rate =
        sample_rate * block_align;

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

    const uint32_t riff_size =
        36 + data_size;

    file.write(
        "RIFF",
        4
    );

    file.write(
        reinterpret_cast<const char*>(&riff_size),
        sizeof(riff_size)
    );

    file.write(
        "WAVE",
        4
    );

    file.write(
        "fmt ",
        4
    );

    const uint32_t fmt_size = 16;

    file.write(
        reinterpret_cast<const char*>(&fmt_size),
        sizeof(fmt_size)
    );

    file.write(
        reinterpret_cast<const char*>(&audio_format),
        sizeof(audio_format)
    );

    file.write(
        reinterpret_cast<const char*>(&channels),
        sizeof(channels)
    );

    file.write(
        reinterpret_cast<const char*>(&sample_rate),
        sizeof(sample_rate)
    );

    file.write(
        reinterpret_cast<const char*>(&byte_rate),
        sizeof(byte_rate)
    );

    file.write(
        reinterpret_cast<const char*>(&block_align),
        sizeof(block_align)
    );

    file.write(
        reinterpret_cast<const char*>(&bits_per_sample),
        sizeof(bits_per_sample)
    );

    file.write(
        "data",
        4
    );

    file.write(
        reinterpret_cast<const char*>(&data_size),
        sizeof(data_size)
    );

    if (!pcm.empty()) {
        file.write(
            reinterpret_cast<const char*>(pcm.data()),
            static_cast<std::streamsize>(
                data_size
            )
        );
    }
}

} // namespace

SpeechToText::SpeechToText(
    SttConfig config)
    : config_(std::move(config))
{
}

std::string SpeechToText::listen()
{
    constexpr const char* WAV_PATH =
        "/tmp/r1_voice_record.wav";

    if (!recordToWav(WAV_PATH))
        return {};

    return transcribe(WAV_PATH);
}

bool SpeechToText::recordToWav(
    const std::string& path)
{
    const int sock =
        socket(
            AF_INET,
            SOCK_DGRAM,
            0
        );

    if (sock < 0) {
        std::cerr
            << "socket() failed: "
            << std::strerror(errno)
            << '\n';

        return false;
    }

    int reuse = 1;

    if (setsockopt(
            sock,
            SOL_SOCKET,
            SO_REUSEADDR,
            &reuse,
            sizeof(reuse)) != 0) {

        std::cerr
            << "SO_REUSEADDR failed: "
            << std::strerror(errno)
            << '\n';

        close(sock);
        return false;
    }

    sockaddr_in local_addr{};

    local_addr.sin_family =
        AF_INET;

    local_addr.sin_port =
        htons(PORT);

    local_addr.sin_addr.s_addr =
        htonl(INADDR_ANY);

    if (bind(
            sock,
            reinterpret_cast<sockaddr*>(&local_addr),
            sizeof(local_addr)) != 0) {

        std::cerr
            << "bind() failed: "
            << std::strerror(errno)
            << '\n';

        close(sock);
        return false;
    }

    const char* env_interface =
        std::getenv(
            "R1_NETWORK_INTERFACE"
        );

    if (!env_interface) {
        std::cerr
            << "R1_NETWORK_INTERFACE is not set\n"
            << "Example:\n"
            << "  export R1_NETWORK_INTERFACE=enp3s0\n";

        close(sock);
        return false;
    }

    const std::string interface =
        env_interface;

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

    if (local_ip.empty()) {
        std::cerr
            << "Could not find IPv4 address for "
            << interface
            << '\n';

        close(sock);
        return false;
    }

    ip_mreq mreq{};

    if (inet_pton(
            AF_INET,
            GROUP_IP,
            &mreq.imr_multiaddr) != 1) {

        close(sock);
        return false;
    }

    if (inet_pton(
            AF_INET,
            local_ip.c_str(),
            &mreq.imr_interface) != 1) {

        close(sock);
        return false;
    }

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

        std::cerr
            << "IP_ADD_MEMBERSHIP failed: "
            << std::strerror(errno)
            << '\n';

        close(sock);
        return false;
    }

    /*
     * R1 sends 5120 byte packets.
     */
    constexpr size_t PACKET_SIZE = 5120;

    std::vector<int16_t> audio;

    const size_t reserve_samples =
        static_cast<size_t>(
            SAMPLE_RATE *
            config_.max_record_ms /
            1000
        );

    audio.reserve(reserve_samples);

    std::vector<int16_t> preroll;

    const size_t preroll_samples =
        static_cast<size_t>(
            SAMPLE_RATE *
            config_.preroll_ms /
            1000
        );

    preroll.reserve(preroll_samples);

    char buffer[PACKET_SIZE];

    bool speaking = false;

    int speech_ms = 0;
    int silence_ms = 0;
    int total_ms = 0;

    std::cout
        << "\nListening...\n";

    while (total_ms <
           config_.max_record_ms) {

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

        if (len <= 0) {
            continue;
        }

        if ((len % 2) != 0)
            continue;

        const auto* samples =
            reinterpret_cast<
                const int16_t*
            >(buffer);

        const size_t sample_count =
            static_cast<size_t>(len) / 2;

        const int packet_ms =
            static_cast<int>(
                sample_count * 1000 /
                SAMPLE_RATE
            );

        total_ms += packet_ms;

        const float level =
            rms(
                samples,
                sample_count
            );

        if (!speaking) {

            preroll.insert(
                preroll.end(),
                samples,
                samples + sample_count
            );

            if (preroll.size() >
                preroll_samples) {

                const size_t remove =
                    preroll.size() -
                    preroll_samples;

                preroll.erase(
                    preroll.begin(),
                    preroll.begin() +
                        static_cast<
                            std::ptrdiff_t
                        >(remove)
                );
            }

            if (level >=
                config_.speech_threshold) {

                speaking = true;

                audio.insert(
                    audio.end(),
                    preroll.begin(),
                    preroll.end()
                );

                preroll.clear();

                speech_ms += packet_ms;

                std::cout
                    << "Speech detected\n";
            }
        }
        else {

            audio.insert(
                audio.end(),
                samples,
                samples + sample_count
            );

            speech_ms += packet_ms;

            if (level <
                config_.speech_threshold) {

                silence_ms += packet_ms;
            }
            else {
                silence_ms = 0;
            }

            if (
                speech_ms >=
                    config_.min_speech_ms &&
                silence_ms >=
                    config_.silence_ms
            ) {
                break;
            }
        }
    }

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

    close(sock);

    if (!speaking ||
        audio.empty()) {

        std::cout
            << "No speech detected.\n";

        return false;
    }

    try {
        writeWav(
            path,
            audio
        );
    }
    catch (const std::exception& e) {
        std::cerr
            << "WAV write failed: "
            << e.what()
            << '\n';

        return false;
    }

    std::cout
        << "Recorded "
        << audio.size() /
            static_cast<double>(
                SAMPLE_RATE
            )
        << " seconds\n";

    return true;
}

std::string SpeechToText::transcribe(
    const std::string& wav_path)
{
    std::ostringstream command;

    command
        << '"'
        << config_.whisper_cli
        << '"'
        << " -m "
        << '"'
        << config_.whisper_model
        << '"'
        << " -f "
        << '"'
        << wav_path
        << '"'
        << " -l "
        << config_.language
        << " -t "
        << config_.threads
        << " --no-timestamps"
        << " -nt"
        << " 2>/dev/null";

    std::cout
        << "Transcribing...\n";

    FILE* pipe =
        popen(
            command.str().c_str(),
            "r"
        );

    if (!pipe) {
        std::cerr
            << "Failed to start whisper-cli\n";

        return {};
    }

    char buffer[4096];

    std::string output;

    while (fgets(
        buffer,
        sizeof(buffer),
        pipe)) {

        output += buffer;
    }

    const int status =
        pclose(pipe);

    if (status != 0) {
        std::cerr
            << "whisper-cli failed: "
            << status
            << '\n';

        return {};
    }

    const std::string text =
        trim(output);

    if (!text.empty()) {
        std::cout
            << "You: "
            << text
            << '\n';
    }

    return text;
}

std::string SpeechToText::trim(
    const std::string& text)
{
    const auto first =
        text.find_first_not_of(
            " \t\r\n"
        );

    if (first == std::string::npos)
        return {};

    const auto last =
        text.find_last_not_of(
            " \t\r\n"
        );

    return text.substr(
        first,
        last - first + 1
    );
}

} // namespace r1_voice
