//
// Copyright 2010-2013 Ettus Research LLC
// Copyright 2018 Ettus Research, a National Instruments Company
// Copyright 2019 Ettus Research, a National Instruments Brand
//
// SPDX-License-Identifier: GPL-3.0-or-later
//

#include <uhd/config.hpp>
#include <uhd/types/stream_cmd.hpp>
#include <uhd/usrp/multi_usrp.hpp>
#include <uhd/utils/safe_main.hpp>
#include <uhd/utils/static.hpp>
#include <uhd/utils/thread.hpp>
#include <boost/format.hpp>
#include <boost/program_options.hpp>
#include <complex>
#include <cstdlib>
#include <ctime>
#include <functional>
#include <iostream>

namespace po = boost::program_options;

/*!
 * Test the late command message:
 *    Issue a stream command with a time that is in the past.
 *    We expect to get an inline late command message.
 */
bool test_late_command_message(uhd::usrp::multi_usrp::sptr usrp,
    uhd::rx_streamer::sptr rx_stream,
    uhd::tx_streamer::sptr)
{
    std::cout << "Test late command message... " << std::flush;

    usrp->set_time_now(uhd::time_spec_t(200.0)); // set time

    uhd::stream_cmd_t stream_cmd(uhd::stream_cmd_t::STREAM_MODE_NUM_SAMPS_AND_DONE);
    stream_cmd.num_samps  = rx_stream->get_max_num_samps();
    stream_cmd.stream_now = false;
    stream_cmd.time_spec  = uhd::time_spec_t(100.0); // time in the past
    rx_stream->issue_stream_cmd(stream_cmd);

    std::vector<std::complex<float>> buff(rx_stream->get_max_num_samps());
    uhd::rx_metadata_t md;

    const size_t nsamps = rx_stream->recv(&buff.front(), buff.size(), md);

    switch (md.error_code) {
        case uhd::rx_metadata_t::ERROR_CODE_LATE_COMMAND:
            std::cout << "success:\n"
                      << "    Got error code late command message.\n"
                      << std::endl;
            return true;

        case uhd::rx_metadata_t::ERROR_CODE_TIMEOUT:
            std::cout << "failed:\n"
                      << "    Inline message recv timed out.\n"
                      << std::endl;
            return false;

        default:
            std::cout << boost::format(
                             "failed:\n"
                             "    Got unexpected error code 0x%x (%s), nsamps %u.\n")
                             % md.error_code % md.strerror() % nsamps
                      << std::endl;
            return false;
    }
}

/*!
 * Test the broken chain message:
 *    Issue a stream command with num samps and more.
 *    We expect to get an inline broken chain message.
 */
bool test_broken_chain_message(UHD_UNUSED(uhd::usrp::multi_usrp::sptr usrp),
    uhd::rx_streamer::sptr rx_stream,
    uhd::tx_streamer::sptr)
{
    std::cout << "Test broken chain message... " << std::flush;

    uhd::stream_cmd_t stream_cmd(uhd::stream_cmd_t::STREAM_MODE_NUM_SAMPS_AND_MORE);
    stream_cmd.stream_now = true;
    stream_cmd.num_samps  = rx_stream->get_max_num_samps();
    rx_stream->issue_stream_cmd(stream_cmd);

    std::vector<std::complex<float>> buff(rx_stream->get_max_num_samps());
    uhd::rx_metadata_t md;

    rx_stream->recv( // once for the requested samples
        &buff.front(),
        buff.size(),
        md);

    rx_stream->recv( // again for the inline message
        &buff.front(),
        buff.size(),
        md);

    switch (md.error_code) {
        case uhd::rx_metadata_t::ERROR_CODE_BROKEN_CHAIN:
            std::cout << boost::format("success:\n"
                                       "    Got error code broken chain message.\n")
                      << std::endl;
            return true;

        case uhd::rx_metadata_t::ERROR_CODE_TIMEOUT:
            std::cout << boost::format("failed:\n"
                                       "    Inline message recv timed out.\n")
                      << std::endl;
            return false;

        default:
            std::cout << boost::format("failed:\n"
                                       "    Got unexpected error code 0x%x (%s).\n")
                             % md.error_code % md.strerror()
                      << std::endl;
            return false;
    }
}

/*!
 * Test the burst ack message:
 *    Send a burst of many samples that will fragment internally.
 *    We expect to get an burst ack async message.
 */
bool test_burst_ack_message(
    uhd::usrp::multi_usrp::sptr, uhd::rx_streamer::sptr, uhd::tx_streamer::sptr tx_stream)
{
    std::cout << "Test burst ack message... " << std::flush;

    uhd::tx_metadata_t md;
    md.start_of_burst = true;
    md.end_of_burst   = true;
    md.has_time_spec  = false;

    std::vector<std::complex<float>> buff(tx_stream->get_max_num_samps());
    tx_stream->send(&buff.front(), buff.size(), md);

    uhd::async_metadata_t async_md;
    if (not tx_stream->recv_async_msg(async_md)) {
        std::cout << "failed:\n"
                  << "    Async message recv timed out.\n"
                  << std::endl;
        return false;
    }

    switch (async_md.event_code) {
        case uhd::async_metadata_t::EVENT_CODE_BURST_ACK:
            std::cout << "success:\n"
                      << "    Got event code burst ack message.\n"
                      << std::endl;
            return true;

        default:
            std::cout << boost::format("failed:\n"
                                       "    Got unexpected event code 0x%x.\n")
                             % async_md.event_code
                      << std::endl;
            return false;
    }
}

/*!
 * Test the underflow message:
 *    Send a start of burst packet with no following end of burst.
 *    We expect to get an underflow(within a burst) async message.
 */
bool test_underflow_message(
    uhd::usrp::multi_usrp::sptr, uhd::rx_streamer::sptr, uhd::tx_streamer::sptr tx_stream)
{
    std::cout << "Test underflow message... " << std::flush;

    uhd::tx_metadata_t md;
    md.start_of_burst = true;
    md.end_of_burst   = false;
    md.has_time_spec  = false;

    std::vector<std::complex<float>> buff(tx_stream->get_max_num_samps());
    tx_stream->send(&buff.front(), buff.size(), md);

    uhd::async_metadata_t async_md;
    bool result = false;
    if (not tx_stream->recv_async_msg(async_md, 1)) {
        std::cout << boost::format("failed:\n"
                                   "    Async message recv timed out.\n")
                  << std::endl;
    } else {
        switch (async_md.event_code) {
            case uhd::async_metadata_t::EVENT_CODE_UNDERFLOW:
                std::cout << boost::format("success:\n"
                                           "    Got event code underflow message.\n")
                          << std::endl;
                result = true;
                break;

            default:
                std::cout << boost::format("failed:\n"
                                           "    Got unexpected event code 0x%x.\n")
                                 % async_md.event_code
                          << std::endl;
                break;
        }
    }

    // Finish the burst to avoid affecting other tests. The DUC only considers
    // time specs from the beginning of a burst, if other tests need a time
    // spec they would be affected by an unfinished burst.
    md.end_of_burst = true;
    tx_stream->send(&buff.front(), buff.size(), md);
    return result;
}

/*!
 * Test the time error message:
 *    Send a burst packet that occurs at a time in the past.
 *    We expect to get a time error async message.
 */
bool test_time_error_message(uhd::usrp::multi_usrp::sptr usrp,
    uhd::rx_streamer::sptr,
    uhd::tx_streamer::sptr tx_stream)
{
    std::cout << "Test time error message... " << std::flush;

    uhd::tx_metadata_t md;
    md.start_of_burst = true;
    md.end_of_burst   = true;
    md.has_time_spec  = true;
    md.time_spec      = uhd::time_spec_t(100.0); // send at 100s

    usrp->set_time_now(uhd::time_spec_t(200.0)); // time at 200s

    std::vector<std::complex<float>> buff(tx_stream->get_max_num_samps());
    tx_stream->send(&buff.front(), buff.size(), md);

    uhd::async_metadata_t async_md;
    if (not tx_stream->recv_async_msg(async_md)) {
        std::cout << boost::format("failed:\n"
                                   "    Async message recv timed out.\n")
                  << std::endl;
        return false;
    }

    switch (async_md.event_code) {
        case uhd::async_metadata_t::EVENT_CODE_TIME_ERROR:
            std::cout << boost::format("success:\n"
                                       "    Got event code time error message.\n")
                      << std::endl;
            return true;

        default:
            std::cout << boost::format("failed:\n"
                                       "    Got unexpected event code 0x%x.\n")
                             % async_md.event_code
                      << std::endl;
            return false;
    }
}

void flush_async(uhd::tx_streamer::sptr tx_stream)
{
    uhd::async_metadata_t async_md;
    while (tx_stream->recv_async_msg(async_md)) {
    }
}

void flush_recv(uhd::rx_streamer::sptr rx_stream)
{
    uhd::stream_cmd_t stream_cmd(uhd::stream_cmd_t::STREAM_MODE_NUM_SAMPS_AND_DONE);
    stream_cmd.num_samps  = rx_stream->get_max_num_samps() * 3;
    stream_cmd.stream_now = true;
    rx_stream->issue_stream_cmd(stream_cmd);

    std::vector<std::complex<float>> buff(stream_cmd.num_samps);
    uhd::rx_metadata_t md;

    do {
        rx_stream->recv(&buff.front(), buff.size(), md);
    } while (md.error_code != uhd::rx_metadata_t::ERROR_CODE_TIMEOUT);
}

int UHD_SAFE_MAIN(int argc, char* argv[])
{
    const std::string program_doc =
        "usage: test_messages [-h] [--args ARGS] [--ntests NTESTS] [--test-chain]"
        "\n\n"
        "This example program demonstrates how to trigger, receive, and inspect\n"
        "standard UHD (USRP Hardware Driver) error and event messages using the\n"
        "UHD API.\n"
        "It runs a series of tests that issue specific stream commands and\n"
        "transmit operations to intentionally provoke built-in UHD responses such\n"
        "as late command, broken chain, burst acknowledgment, underflow, and time\n"
        "error messages. The program prints the outcome of each test, helping\n"
        "users understand how UHD reports errors and how to implement robust\n"
        "error handling in message-based communication with USRP devices.\n";
    // variables to be set by po
    std::string args;
    size_t ntests;

    // setup the program options
    po::options_description desc("Allowed options");
    // clang-format off
    desc.add_options()
        ("help", "Show this help message and exit.")
        ("args",   po::value<std::string>(&args)->default_value(""), "Single USRP device selection and "
            "configuration arguments."
            "\nSpecify key-value pairs (e.g., addr, serial, type, master_clock_rate) separated by commas."
            "\nSee the UHD manual for model-specific options."
            "\nExamples:"
            "\n  --args \"addr=192.168.10.2\""
            "\n  --args \"addr=192.168.10.2,master_clock_rate=200e6\""
            "\nIf not specified, UHD connects to the first available device.")
        ("ntests", po::value<size_t>(&ntests)->default_value(50),    "Number of tests to run.")
        ("test-chain", "Include tests that trigger a \"broken chain\" error, which happens when the host "
            "stops receiving data too early after starting a special streaming mode."
            "\nThe test does this by issuing a stream command with STREAM_MODE_NUM_SAMPS_AND_MORE but not continuing to "
            "call recv as required.")
    ;
    // clang-format on
    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);

    // print the help message
    if (vm.count("help")) {
        std::cout << program_doc << std::endl;
        std::cout << desc << std::endl;
        return ~0;
    }

    // create a usrp device
    std::cout << std::endl;
    std::cout << boost::format("Creating the usrp device with: %s...") % args
              << std::endl;
    uhd::usrp::multi_usrp::sptr usrp = uhd::usrp::multi_usrp::make(args);
    std::cout << boost::format("Using Device: %s") % usrp->get_pp_string() << std::endl;

    // create RX and TX streamers
    uhd::stream_args_t stream_args("fc32"); // complex floats
    uhd::rx_streamer::sptr rx_stream;
    uhd::tx_streamer::sptr tx_stream;
    if (usrp->get_rx_num_channels()) {
        rx_stream = usrp->get_rx_stream(stream_args);
    }
    if (usrp->get_tx_num_channels()) {
        tx_stream = usrp->get_tx_stream(stream_args);
    }

    //------------------------------------------------------------------
    // begin messages test
    //------------------------------------------------------------------
    using test_executor_fn = std::function<bool(
        uhd::usrp::multi_usrp::sptr, uhd::rx_streamer::sptr, uhd::tx_streamer::sptr)>;
    uhd::dict<std::string, test_executor_fn> tests;
    if (tx_stream) {
        tests["Test Burst ACK "] = &test_burst_ack_message;
        tests["Test Underflow "] = &test_underflow_message;
        tests["Test Time Error"] = &test_time_error_message;
    }
    if (rx_stream) {
        tests["Test Late Command"] = &test_late_command_message;
    }

    if (vm.count("test-chain")) {
        tests["Test Broken Chain"] = &test_broken_chain_message;
    }

    // init result counts
    uhd::dict<std::string, size_t> failures, successes;
    for (const std::string& key : tests.keys()) {
        failures[key]  = 0;
        successes[key] = 0;
    }

    // run the tests, pick at random
    std::srand((unsigned int)time(NULL));
    for (size_t n = 0; n < ntests; n++) {
        std::string key = tests.keys()[std::rand() % tests.size()];
        bool pass       = tests[key](usrp, rx_stream, tx_stream);

        if (rx_stream) {
            flush_recv(rx_stream);
        }
        if (tx_stream) {
            flush_async(tx_stream);
        }

        // store result
        if (pass)
            successes[key]++;
        else
            failures[key]++;
    }

    // print the result summary
    bool any_failure = false;
    std::cout << std::endl << "Summary:" << std::endl << std::endl;
    for (const std::string& key : tests.keys()) {
        std::cout << boost::format("%s   ->   %3u successes, %3u failures") % key
                         % successes[key] % failures[key]
                  << std::endl;
        any_failure = any_failure or (failures[key] > 0);
    }

    // finished
    std::cout << std::endl << "Done!" << std::endl << std::endl;

    return any_failure ? EXIT_FAILURE : EXIT_SUCCESS;
}
