#include "r1_voice/tts.hh"

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <vector>
#include <chrono>
#include <thread>

#include <sys/wait.h>

#include <unitree/common/time/time_tool.hpp>
#include <unitree/robot/channel/channel_factory.hpp>

namespace r1_voice {

namespace {

bool readU16(
    std::ifstream& f,
    uint16_t& value)
{
    f.read(
        reinterpret_cast<char*>(&value),
        sizeof(value)
    );

    return static_cast<bool>(f);
}

bool readU32(
    std::ifstream& f,
    uint32_t& value)
{
    f.read(
        reinterpret_cast<char*>(&value),
        sizeof(value)
    );

    return static_cast<bool>(f);
}

bool runCommand(
    const std::string& command)
{
    const int status =
        std::system(command.c_str());

    if (status == -1)
        return false;

    return WIFEXITED(status) &&
           WEXITSTATUS(status) == 0;
}

} // namespace

TextToSpeech::TextToSpeech(
    TtsConfig config)
    : config_(std::move(config))
{
}

bool TextToSpeech::init(
    const std::string& network_interface)
{
    if (network_interface.empty()) {
        std::cerr
            << "TTS: network interface is empty\n";

        return false;
    }

    try {
        /*
         * ChannelFactory MUST be initialized before
         * constructing AudioClient.
         */
        unitree::robot::ChannelFactory::
            Instance()->Init(
                0,
                network_interface
            );

        client_ =
            std::make_unique<
                unitree::robot::r1::AudioClient
            >();

        client_->Init();

        client_->SetTimeout(10.0f);

        const int32_t ret =
            client_->SetVolume(
                static_cast<uint8_t>(
                    config_.volume
                )
            );

        if (ret != 0) {
            std::cerr
                << "TTS: SetVolume failed: "
                << ret
                << '\n';
        }

        initialized_ = true;

        std::cout
            << "TTS initialized\n"
            << "  Piper model: "
            << config_.model
            << '\n'
            << "  R1 interface: "
            << network_interface
            << '\n';

        return true;
    }
    catch (const std::exception& e) {
        std::cerr
            << "TTS init failed: "
            << e.what()
            << '\n';

        client_.reset();
        initialized_ = false;

        return false;
    }
}

bool TextToSpeech::synthesize(
    const std::string& text,
    const std::string& wav_path)
{
    const std::string text_path =
        "/tmp/r1_tts_input.txt";

    {
        std::ofstream out(
            text_path,
            std::ios::binary
        );

        if (!out) {
            std::cerr
                << "Cannot create "
                << text_path
                << '\n';

            return false;
        }

        out << text;
    }

    std::ostringstream command;

    command
        << "python -m piper"
        << " --model "
        << '\''
        << config_.model
        << '\''
        << " --output_file "
        << '\''
        << wav_path
        << '\''
        << " < "
        << '\''
        << text_path
        << '\''
        << " >/dev/null 2>&1";

    std::cout
        << "Synthesizing Czech speech...\n";

    if (!runCommand(command.str())) {
        std::cerr
            << "Piper failed\n";

        return false;
    }

    return true;
}

bool TextToSpeech::readWavPcm(
    const std::string& wav_path,
    std::vector<uint8_t>& pcm,
    uint32_t& sample_rate,
    uint16_t& channels,
    uint16_t& bits_per_sample)
{
    std::ifstream f(
        wav_path,
        std::ios::binary
    );

    if (!f) {
        std::cerr
            << "Cannot open WAV: "
            << wav_path
            << '\n';

        return false;
    }

    char riff[4];
    char wave[4];

    f.read(riff, 4);

    if (!f ||
        std::memcmp(riff, "RIFF", 4) != 0) {

        std::cerr
            << "Invalid RIFF file\n";

        return false;
    }

    uint32_t riff_size = 0;

    if (!readU32(f, riff_size))
        return false;

    (void)riff_size;

    f.read(wave, 4);

    if (!f ||
        std::memcmp(wave, "WAVE", 4) != 0) {

        std::cerr
            << "Invalid WAVE file\n";

        return false;
    }

    bool found_fmt = false;
    bool found_data = false;

    uint16_t audio_format = 0;

    uint32_t data_size = 0;
    std::streampos data_position{};

    while (f && (!found_fmt || !found_data)) {

        char chunk_id[4];
        uint32_t chunk_size = 0;

        f.read(
            chunk_id,
            sizeof(chunk_id)
        );

        if (!f)
            break;

        if (!readU32(f, chunk_size))
            return false;

        if (std::memcmp(
                chunk_id,
                "fmt ",
                4) == 0) {

            if (chunk_size < 16) {
                std::cerr
                    << "Invalid fmt chunk\n";

                return false;
            }

            if (!readU16(
                    f,
                    audio_format))
                return false;

            if (!readU16(
                    f,
                    channels))
                return false;

            if (!readU32(
                    f,
                    sample_rate))
                return false;

            uint32_t byte_rate = 0;
            uint16_t block_align = 0;

            if (!readU32(
                    f,
                    byte_rate))
                return false;

            if (!readU16(
                    f,
                    block_align))
                return false;

            if (!readU16(
                    f,
                    bits_per_sample))
                return false;

            (void)byte_rate;
            (void)block_align;

            if (chunk_size > 16) {
                f.seekg(
                    chunk_size - 16,
                    std::ios::cur
                );
            }

            found_fmt = true;
        }
        else if (
            std::memcmp(
                chunk_id,
                "data",
                4) == 0) {

            data_size = chunk_size;

            data_position = f.tellg();

            f.seekg(
                chunk_size,
                std::ios::cur
            );

            found_data = true;
        }
        else {
            f.seekg(
                chunk_size,
                std::ios::cur
            );
        }

        if (chunk_size & 1)
            f.seekg(
                1,
                std::ios::cur
            );
    }

    if (!found_fmt ||
        !found_data) {

        std::cerr
            << "WAV missing fmt/data chunk\n";

        return false;
    }

    if (audio_format != 1) {
        std::cerr
            << "WAV is not PCM\n";

        return false;
    }

    if (channels != 1) {
        std::cerr
            << "WAV must be mono, got "
            << channels
            << '\n';

        return false;
    }

    if (bits_per_sample != 16) {
        std::cerr
            << "WAV must be 16-bit PCM, got "
            << bits_per_sample
            << '\n';

        return false;
    }

    f.seekg(data_position);

    pcm.resize(data_size);

    f.read(
        reinterpret_cast<char*>(pcm.data()),
        static_cast<std::streamsize>(
            data_size
        )
    );

    if (!f) {
        std::cerr
            << "Failed reading PCM data\n";

        return false;
    }

    return true;
}

bool TextToSpeech::playPcm(
    const std::vector<uint8_t>& pcm)
{
    if (!client_ || pcm.empty())
        return false;

    // 16 kHz, mono, signed 16-bit PCM:
    // 16000 samples/s * 2 bytes = 32000 bytes/s
    constexpr double BYTES_PER_SECOND = 32000.0;

    // Same chunk size as Unitree's R1 audio example:
    // 96000 bytes = 3 seconds.
    constexpr size_t CHUNK_SIZE = 96000;

    // Feed chunks faster than playback so R1 keeps
    // an audio buffer and does not underrun.
    constexpr int INTER_CHUNK_DELAY_MS = 1000;

    // Extra time after predicted playback completion before
    // enabling the microphone again.
    constexpr int PLAYBACK_MARGIN_MS = 600;

    const std::string stream_id =
        std::to_string(
            unitree::common::GetCurrentTimeMillisecond()
        );

    const auto playback_start =
        std::chrono::steady_clock::now();

    size_t offset = 0;
    size_t chunk_index = 0;

    while (offset < pcm.size()) {

        const size_t remaining =
            pcm.size() - offset;

        const size_t chunk_size =
            std::min(
                CHUNK_SIZE,
                remaining
            );

        std::vector<uint8_t> chunk(
            pcm.begin() +
                static_cast<std::ptrdiff_t>(offset),
            pcm.begin() +
                static_cast<std::ptrdiff_t>(
                    offset + chunk_size
                )
        );

        std::cout
            << "Sending TTS chunk "
            << chunk_index
            << ": "
            << chunk_size
            << " bytes\n";

        const int32_t ret =
            client_->PlayStream(
                config_.app_name,
                stream_id,
                chunk
            );

        std::cout
            << "PlayStream returned "
            << ret
            << '\n';

        if (ret != 0) {
            std::cerr
                << "PlayStream failed on chunk "
                << chunk_index
                << '\n';

            return false;
        }

        offset += chunk_size;
        ++chunk_index;

        // Keep the R1 buffer filled ahead of playback.
        if (offset < pcm.size()) {
            std::this_thread::sleep_for(
                std::chrono::milliseconds(
                    INTER_CHUNK_DELAY_MS
                )
            );
        }
    }

    std::cout
        << "All TTS chunks submitted\n";

    /*
     * Work out when the complete utterance should finish
     * relative to when we submitted its first chunk.
     */
    const auto total_audio_ms =
        std::chrono::milliseconds(
            static_cast<int64_t>(
                (static_cast<double>(pcm.size()) /
                 BYTES_PER_SECOND) *
                1000.0
            )
        );

    const auto now =
        std::chrono::steady_clock::now();

    const auto elapsed =
        std::chrono::duration_cast<
            std::chrono::milliseconds
        >(now - playback_start);

    auto remaining_playback =
        total_audio_ms - elapsed;

    if (remaining_playback.count() < 0) {
        remaining_playback =
            std::chrono::milliseconds(0);
    }

    const auto final_wait =
        remaining_playback +
        std::chrono::milliseconds(
            PLAYBACK_MARGIN_MS
        );

    std::cout
        << "Waiting "
        << final_wait.count()
        << " ms for R1 playback to finish\n";

    std::this_thread::sleep_for(
        final_wait
    );

    /*
     * IMPORTANT:
     * Don't PlayStop() here.
     *
     * We let the R1 finish the buffered stream naturally.
     */
    std::cout
        << "TTS playback finished\n";

    return true;
}

bool TextToSpeech::speak(
    const std::string& text)
{
    if (!initialized_) {
        std::cerr
            << "TTS is not initialized\n";

        return false;
    }

    if (text.empty())
        return true;

    std::cout
        << "R1: "
        << text
        << '\n';

    const std::string piper_wav =
        config_.output_wav;

    const std::string output_wav_16k =
        "/tmp/r1_tts_16k.wav";

    if (!synthesize(
            text,
            piper_wav)) {

        return false;
    }

    /*
     * Piper's Czech model produces 22050 Hz.
     *
     * R1 PlayStream requires 16000 Hz mono PCM.
     *
     * Resample with ffmpeg.
     */
    std::ostringstream resample;

    resample
        << "ffmpeg -y -loglevel error"
        << " -i "
        << '\''
        << piper_wav
        << '\''
        << " -ar 16000"
        << " -ac 1"
        << " -c:a pcm_s16le"
        << " "
        << '\''
        << output_wav_16k
        << '\''
        << " >/dev/null 2>&1";

    if (!runCommand(
            resample.str())) {

        std::cerr
            << "ffmpeg resampling failed\n";

        return false;
    }

    std::vector<uint8_t> pcm;

    uint32_t sample_rate = 0;
    uint16_t channels = 0;
    uint16_t bits = 0;

    if (!readWavPcm(
            output_wav_16k,
            pcm,
            sample_rate,
            channels,
            bits)) {

        return false;
    }

    std::cout
        << "TTS WAV: "
        << sample_rate
        << " Hz, "
        << channels
        << " channel, "
        << bits
        << " bit, "
        << pcm.size()
        << " PCM bytes\n";

    if (sample_rate != 16000 ||
        channels != 1 ||
        bits != 16) {

        std::cerr
            << "Resampled WAV has unexpected format\n";

        return false;
    }

    return playPcm(pcm);
}

} // namespace r1_voice
