//
// Copyright 2017 Ettus Research, a National Instruments Company
//
// SPDX-License-Identifier: GPL-3.0-or-later
//

#include "mpmd_impl.hpp"
#include <uhd/transport/udp_simple.hpp>
#include <uhd/utils/log.hpp>
#include <uhd/utils/safe_call.hpp>
#include <chrono>
#include <thread>

namespace {
/*************************************************************************
 * Local constants
 ************************************************************************/
//! Timeout for pings (seconds).
constexpr double MPMD_PING_TIMEOUT = 0.1;
//! Default session ID (MPM will recognize a session by this name)
const std::string MPMD_DEFAULT_SESSION_ID = "UHD";
//! Key to initialize a ping/measurement latency test
const std::string MPMD_MEAS_LATENCY_KEY = "measure_rpc_latency";
//! Duration of a latency measurement test
constexpr size_t MPMD_MEAS_LATENCY_DURATION = 1000;

using log_buf_t = std::vector<std::map<std::string, std::string>>;


/*************************************************************************
 * Helper functions
 ************************************************************************/
/*! Return true if we can MPM ping a device via discovery service.
 */
bool is_pingable(const std::string& ip_addr, const std::string& udp_port)
{
    auto udp = uhd::transport::udp_simple::make_broadcast(ip_addr, udp_port);
    const std::string send_buf(uhd::mpmd::mpmd_impl::MPM_ECHO_CMD + " ping");
    std::vector<uint8_t> recv_buf;
    recv_buf.resize(send_buf.size(), ' ');
    udp->send(boost::asio::buffer(send_buf.c_str(), send_buf.size()));
    const size_t len = udp->recv(boost::asio::buffer(recv_buf), MPMD_PING_TIMEOUT);
    if (len == 0) {
        UHD_LOG_TRACE("MPMD", "Received no MPM ping, assuming device is unreachable.");
        return false;
    }
    if (len != send_buf.size()) {
        UHD_LOG_DEBUG("MPMD",
            "Received bad return packet on MPM ping. Assuming endpoint"
            " is not a valid MPM device.");
        return false;
    }
    if (std::memcmp((void*)&recv_buf[0], (void*)&send_buf[0], send_buf.size()) != 0) {
        UHD_LOG_DEBUG("MPMD",
            "Received invalid return packet on MPM ping. Assuming endpoint"
            " is not a valid MPM device.");
        return false;
    }
    return true;
}

/*! Call init() on an MPM device.
 */
void init_device(uhd::rpc_client::sptr rpc, const uhd::device_addr_t mb_args)
{
    auto init_status = rpc->request_with_token<std::vector<std::string>>(
        MPMD_DEFAULT_INIT_TIMEOUT, "get_init_status");
    if (init_status[0] != "true") {
        throw uhd::runtime_error(
            std::string("Device is in bad state: ") + init_status[1]);
    }

    // TODO maybe put this somewhere else?
    const std::set<std::string> key_blacklist{"serial", "claimed", "type", "rev", "addr"};
    std::map<std::string, std::string> mpm_device_args;
    for (const auto& key : mb_args.keys()) {
        if (not key_blacklist.count(key)) {
            mpm_device_args[key] = mb_args[key];
        }
    }
    if (not rpc->request_with_token<bool>(
            MPMD_DEFAULT_INIT_TIMEOUT, "init", mpm_device_args)) {
        throw uhd::runtime_error("Failed to initialize device.");
    }
}

void measure_rpc_latency(
    uhd::rpc_client::sptr rpc, const size_t duration_ms = MPMD_MEAS_LATENCY_DURATION)
{
    const double alpha        = 0.99;
    const std::string payload = "1234567890";
    auto measure_once         = [payload, rpc]() {
        const auto start = std::chrono::steady_clock::now();
        rpc->request<std::string>("ping", payload);
        return (double)std::chrono::duration_cast<std::chrono::microseconds>(
            std::chrono::steady_clock::now() - start)
            .count();
    };

    double max_latency = measure_once();
    double avg_latency = max_latency;

    auto end_time =
        std::chrono::steady_clock::now() + std::chrono::milliseconds(duration_ms);
    size_t ctr = 1;
    while (std::chrono::steady_clock::now() < end_time) {
        const auto duration = measure_once();
        max_latency         = std::max(max_latency, duration);
        avg_latency         = avg_latency * alpha + (1 - alpha) * duration;
        ctr++;
        // Light throttle:
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }

    UHD_LOG_INFO("MPMD",
        "RPC latency (coarse estimate): Avg = " << avg_latency
                                                << " us, "
                                                   "Max = "
                                                << max_latency << ", n = " << ctr);
}

/*! Forward entries from a list of dictionaries to UHD's native logging
 *  system.
 */
void forward_logs(log_buf_t&& log_buf)
{
    for (const auto& log_record : log_buf) {
        if (log_record.count("levelname") == 0 or log_record.count("message") == 0) {
            UHD_LOG_ERROR("MPMD", "Invalid logging structure returned from MPM device!");
            continue;
        }
        if (log_record.at("levelname") == "TRACE") {
            UHD_LOG_TRACE(log_record.at("name"), log_record.at("message"));
        } else if (log_record.at("levelname") == "DEBUG") {
            UHD_LOG_DEBUG(log_record.at("name"), log_record.at("message"));
        } else if (log_record.at("levelname") == "INFO") {
            UHD_LOG_INFO(log_record.at("name"), log_record.at("message"));
        } else if (log_record.at("levelname") == "WARNING") {
            UHD_LOG_WARNING(log_record.at("name"), log_record.at("message"));
        } else if (log_record.at("levelname") == "ERROR") {
            UHD_LOG_ERROR(log_record.at("name"), log_record.at("message"));
        } else if (log_record.at("levelname") == "CRITICAL") {
            UHD_LOG_FATAL(log_record.at("name"), log_record.at("message"));
        } else {
            UHD_LOG_ERROR("MPMD",
                "Invalid log level returned from MPM device: "
                "`" << log_record.at("levelname")
                    << "'");
        }
    }
}

/*! Return a new rpc_client with given addr and mb args
 */
uhd::rpc_client::sptr make_mpm_rpc_client(const std::string& rpc_server_addr,
    const uhd::device_addr_t& mb_args,
    const size_t timeout_ms = MPMD_DEFAULT_RPC_TIMEOUT)
{
    return uhd::rpc_client::make(rpc_server_addr,
        mb_args.cast<size_t>(
            uhd::mpmd::mpmd_impl::MPM_RPC_PORT_KEY, uhd::mpmd::mpmd_impl::MPM_RPC_PORT),
        timeout_ms,
        uhd::mpmd::mpmd_impl::MPM_RPC_GET_LAST_ERROR_CMD);
}

} // namespace

using namespace uhd;
using namespace uhd::mpmd;

/******************************************************************************
 * Static Helpers
 *****************************************************************************/
boost::optional<device_addr_t> mpmd_mboard_impl::is_device_reachable(
    const device_addr_t& device_addr)
{
    UHD_LOG_TRACE(
        "MPMD", "Checking accessibility of device `" << device_addr.to_string() << "'");
    UHD_ASSERT_THROW(device_addr.has_key(xport::MGMT_ADDR_KEY));
    const std::string rpc_addr = device_addr.get(xport::MGMT_ADDR_KEY);
    const size_t rpc_port =
        device_addr.cast<size_t>(mpmd_impl::MPM_RPC_PORT_KEY, mpmd_impl::MPM_RPC_PORT);
    auto rpcc = uhd::rpc_client::make(rpc_addr, rpc_port);
    // 1) Read back device info
    dev_info device_info_dict;
    try {
        device_info_dict =
            rpcc->request<dev_info>(MPMD_SHORT_RPC_TIMEOUT, "get_device_info");
    } catch (const uhd::runtime_error& e) {
        UHD_LOG_ERROR("MPMD", e.what());
    } catch (...) {
        UHD_LOG_DEBUG("MPMD",
            "Unexpected exception when trying to query device info. Flagging "
            "device as unreachable.");
        return boost::optional<device_addr_t>();
    }
    // 2) Check for local device
    if (device_info_dict.count("connection")
        and device_info_dict.at("connection") == "local") {
        UHD_LOG_TRACE("MPMD", "Device is local, flagging as reachable.");
        return boost::optional<device_addr_t>(device_addr);
    }
    // 3) Check for network-reachable device
    // Note: This makes the assumption that devices will always allow RPC
    // connections on their CHDR addresses.
    const std::vector<std::string> addr_keys = {"second_addr", "addr"};
    for (const auto& addr_key : addr_keys) {
        if (not device_info_dict.count(addr_key)) {
            continue;
        }
        const std::string chdr_addr = device_info_dict.at(addr_key);
        UHD_LOG_TRACE("MPMD", "Checking reachability via network addr " << chdr_addr);
        try {
            // First do an MPM ping -- there is some issue with rpclib that can
            // lead to indefinite timeouts
            const std::string mpm_discovery_port =
                device_addr.get(mpmd_impl::MPM_DISCOVERY_PORT_KEY,
                    std::to_string(mpmd_impl::MPM_DISCOVERY_PORT));
            if (!is_pingable(chdr_addr, mpm_discovery_port)) {
                UHD_LOG_TRACE("MPMD", "Cannot MPM ping, assuming device is unreachable.");
                continue;
            }
            UHD_LOG_TRACE("MPMD", "Was able to ping device, trying RPC connection.");
            auto chdr_rpcc = uhd::rpc_client::make(chdr_addr, rpc_port);
            auto dev_info_chdr =
                chdr_rpcc->request<dev_info>(MPMD_SHORT_RPC_TIMEOUT, "get_device_info");
            if (dev_info_chdr["serial"] != device_info_dict["serial"]) {
                UHD_LOG_DEBUG("MPMD",
                    boost::format("Connected to CHDR interface, but got wrong device. "
                                  "Tried to reach serial %s, got %s")
                        % device_info_dict["serial"] % dev_info_chdr["serial"]);
                return boost::optional<device_addr_t>();
            } else {
                UHD_LOG_TRACE("MPMD",
                    boost::format("Reachable device matches expected device (serial=%s)")
                        % device_info_dict["serial"]);
            }
            device_addr_t device_addr_copy = device_addr;
            device_addr_copy["addr"]       = chdr_addr;
            return boost::optional<device_addr_t>(device_addr_copy);
        } catch (...) {
            UHD_LOG_DEBUG(
                "MPMD", "Failed to reach device on network addr " << chdr_addr << ".");
        }
    }
    // If everything fails, we probably can't talk to this chap.
    UHD_LOG_TRACE(
        "MPMD", "All reachability checks failed -- assuming device is not reachable.");
    return boost::optional<device_addr_t>();
}

/*****************************************************************************
 * Structors
 ****************************************************************************/
mpmd_mboard_impl::mpmd_mboard_impl(
    const device_addr_t& mb_args_, const std::string& rpc_server_addr)
    : mb_args(mb_args_)
    , rpc(make_mpm_rpc_client(rpc_server_addr, mb_args_))
    , num_xbars(rpc->request<size_t>("get_num_xbars"))
    , _claim_rpc(make_mpm_rpc_client(rpc_server_addr, mb_args, MPMD_CLAIMER_RPC_TIMEOUT))
    // xbar_local_addrs is not yet valid after this!
    , xbar_local_addrs(num_xbars, 0xFF)
    , _xport_mgr(xport::mpmd_xport_mgr::make(mb_args))
{
    UHD_LOGGER_TRACE("MPMD") << "Initializing mboard, connecting to RPC server address: "
                             << rpc_server_addr << " mboard args: " << mb_args.to_string()
                             << " number of crossbars: " << num_xbars;

    _claimer_task = claim_device_and_make_task();
    if (mb_args_.has_key(MPMD_MEAS_LATENCY_KEY)) {
        measure_rpc_latency(rpc, MPMD_MEAS_LATENCY_DURATION);
    }

    /// Get device info
    const auto device_info_dict = rpc->request<dev_info>("get_device_info");
    for (const auto& info_pair : device_info_dict) {
        device_info[info_pair.first] = info_pair.second;
    }
    UHD_LOGGER_TRACE("MPMD") << "MPM reports device info: " << device_info.to_string();
    /// Get dboard info
    const auto dboards_info = rpc->request<std::vector<dev_info>>("get_dboard_info");
    UHD_ASSERT_THROW(this->dboard_info.size() == 0);
    for (const auto& dboard_info_dict : dboards_info) {
        uhd::device_addr_t this_db_info;
        for (const auto& info_pair : dboard_info_dict) {
            this_db_info[info_pair.first] = info_pair.second;
        }
        UHD_LOGGER_TRACE("MPMD")
            << "MPM reports dboard info for slot " << this->dboard_info.size() << ": "
            << this_db_info.to_string();
        this->dboard_info.push_back(this_db_info);
    }
}

mpmd_mboard_impl::~mpmd_mboard_impl()
{
    try {
        dump_logs();
    } catch (...) {
        UHD_LOG_WARNING("MPMD", "Could not flush log queue on exit!");
    }
    UHD_SAFE_CALL(if (not rpc->request_with_token<bool>("unclaim")) {
        UHD_LOG_WARNING("MPMD", "Failure to ack unclaim!");
    });
}

/*****************************************************************************
 * Init
 ****************************************************************************/
void mpmd_mboard_impl::init()
{
    init_device(rpc, mb_args);
    // RFNoC block clocks are now on. Noc-IDs can be read back.
}

/*****************************************************************************
 * API
 ****************************************************************************/
void mpmd_mboard_impl::set_xbar_local_addr(
    const size_t xbar_index, const size_t local_addr)
{
    UHD_ASSERT_THROW(
        rpc->request_with_token<bool>("set_xbar_local_addr", xbar_index, local_addr));
    UHD_ASSERT_THROW(xbar_index < xbar_local_addrs.size());
    xbar_local_addrs.at(xbar_index) = local_addr;
}

uhd::both_xports_t mpmd_mboard_impl::make_transport(const sid_t& sid,
    usrp::device3_impl::xport_type_t xport_type,
    const uhd::device_addr_t& xport_args)
{
    const std::string xport_type_str = [xport_type]() {
        switch (xport_type) {
            case mpmd_impl::CTRL:
                return "CTRL";
            case mpmd_impl::ASYNC_MSG:
                return "ASYNC_MSG";
            case mpmd_impl::RX_DATA:
                return "RX_DATA";
            case mpmd_impl::TX_DATA:
                return "TX_DATA";
            default:
                UHD_THROW_INVALID_CODE_PATH();
        };
    }();

    UHD_LOGGER_TRACE("MPMD") << __func__
                             << "(): Creating new transport of type: " << xport_type_str;

    using namespace uhd::mpmd::xport;
    const auto xport_info_list =
        rpc->request_with_token<mpmd_xport_mgr::xport_info_list_t>(
            "request_xport", sid.get_dst(), sid.get_src(), xport_type_str);
    UHD_LOGGER_TRACE("MPMD") << __func__ << "(): request_xport() gave us "
                             << xport_info_list.size() << " option(s).";
    if (xport_info_list.empty()) {
        UHD_LOG_ERROR("MPMD", "No viable transport path found!");
        throw uhd::runtime_error("No viable transport path found!");
    }

    xport::mpmd_xport_mgr::xport_info_t xport_info_out;
    auto xports = _xport_mgr->make_transport(
        xport_info_list, xport_type, xport_args, xport_info_out);

    if (not rpc->request_with_token<bool>("commit_xport", xport_info_out)) {
        UHD_LOG_ERROR("MPMD", "Failed to create UDP transport!");
        throw uhd::runtime_error("commit_xport() failed!");
    }

    return xports;
}

size_t mpmd_mboard_impl::get_mtu(const uhd::direction_t dir) const
{
    return _xport_mgr->get_mtu(dir);
}

uhd::device_addr_t mpmd_mboard_impl::get_rx_hints() const
{
    // TODO: See if we need to do anything here. get_rx_stream() might care.
    device_addr_t rx_hints;
    return rx_hints;
}

uhd::device_addr_t mpmd_mboard_impl::get_tx_hints() const
{
    // TODO: See if we need to do anything here. get_tx_stream() might care.
    device_addr_t tx_hints;
    return tx_hints;
}

/*****************************************************************************
 * Private methods
 ****************************************************************************/
bool mpmd_mboard_impl::claim()
{
    try {
        return _claim_rpc->request_with_token<bool>("reclaim");
    } catch (...) {
        UHD_LOG_WARNING("MPMD", "Reclaim failed. Exiting claimer loop.");
        return false;
    }
}

uhd::task::sptr mpmd_mboard_impl::claim_device_and_make_task()
{
    auto rpc_token = _claim_rpc->request<std::string>(
        "claim", mb_args.get("session_id", MPMD_DEFAULT_SESSION_ID));
    if (rpc_token.empty()) {
        throw uhd::value_error("mpmd device claiming failed!");
    }
    UHD_LOG_TRACE("MPMD", "Received claim token " << rpc_token);
    // Save token for both RPC clients
    _claim_rpc->set_token(rpc_token);
    rpc->set_token(rpc_token);
    return uhd::task::make([this] {
        auto now = std::chrono::steady_clock::now();
        if (not this->claim()) {
            throw uhd::value_error("mpmd device reclaiming loop failed!");
        };
        this->dump_logs();
        std::this_thread::sleep_until(
            now + std::chrono::milliseconds(MPMD_RECLAIM_INTERVAL_MS));
    });
}

void mpmd_mboard_impl::dump_logs(const bool dump_to_null)
{
    // We need to use _claim_rpc instead of rpc because this currently only
    // gets called in the claimer loop.
    if (dump_to_null) {
        _claim_rpc->request_with_token<log_buf_t>("get_log_buf");
    } else {
        forward_logs(_claim_rpc->request_with_token<log_buf_t>("get_log_buf"));
    }
}


/*****************************************************************************
 * Factory
 ****************************************************************************/
mpmd_mboard_impl::uptr mpmd_mboard_impl::make(
    const uhd::device_addr_t& mb_args, const std::string& addr)
{
    mpmd_mboard_impl::uptr mb =
        mpmd_mboard_impl::uptr(new mpmd_mboard_impl(mb_args, addr));
    // implicit move
    return mb;
}