/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "ucp_tl_info.h"

#include <ucs/algorithm/qsort_r.h>
#include <ucs/datastruct/string_buffer.h>
#include <ucs/debug/log.h>
#include <ucs/debug/table.h>
#include <ucs/sys/string.h>
#include <ucs/sys/topo/base/topo.h>
#include <string.h>


#define UCP_TL_INFO_DEVS_PER_LINE 3
#define UCP_TL_INFO_MARK_ENABLED  "+"
#define UCP_TL_INFO_MARK_DISABLED "-"
#define UCP_TL_INFO_HDR_TYPE      "Type"
#define UCP_TL_INFO_HDR_TRANSPORT "Transport"
#define UCP_TL_INFO_HDR_DEVICE    "Device (System device)"
#define UCP_TL_INFO_HDR_COMPONENT "Component"
#define UCP_TL_INFO_UNAVAILABLE   "<unavailable>"

#define UCP_TL_INFO_NUM_COLS 4

typedef struct {
    int is_first_type;
    int is_first_cmpt;
    int is_first_tl;
} ucp_tl_info_tl_group_state_t;


static const char *ucp_tl_info_legend_rows[] = {
    "Legend: + = enabled, - = disabled",
    "All of the available transports are listed, some may be disabled or unsupported on your system.",
    "All of the visible devices are listed per transport, some may be disabled.",
};


static int ucp_tl_info_is_same_tl_group(const ucp_tl_info_entry_t *a,
                                        const ucp_tl_info_entry_t *b)
{
    return (a->cmpt_index == b->cmpt_index) &&
           (strcmp(a->rsc.tl_name, b->rsc.tl_name) == 0);
}

static int
ucp_tl_info_compare(const void *a, const void *b, void *UCS_V_UNUSED arg)
{
    const ucp_tl_info_entry_t *ea = a;
    const ucp_tl_info_entry_t *eb = b;
    int diff;

    diff = (int)ea->rsc.dev_type - (int)eb->rsc.dev_type;
    if (diff != 0) {
        return diff;
    }

    diff = (int)ea->cmpt_index - (int)eb->cmpt_index;
    if (diff != 0) {
        return diff;
    }

    return strcmp(ea->rsc.tl_name, eb->rsc.tl_name);
}

static int ucp_tl_info_rscs_has_cmpt(const ucp_tl_info_array_t *all_rscs,
                                     ucp_rsc_index_t cmpt_idx)
{
    const ucp_tl_info_entry_t *entry;

    ucs_array_for_each(entry, all_rscs) {
        if (entry->cmpt_index == cmpt_idx) {
            return 1;
        }
    }

    return 0;
}

static void ucp_tl_info_emit_row(ucs_table_t *table, const char *type_str,
                                 const char *cmpt_str, const char *tl_str,
                                 const char *dev_str,
                                 ucp_tl_info_tl_group_state_t *group_state)
{
    ucs_table_row_h row;

    ucs_table_add_row(table, &row);

    if (group_state->is_first_type) {
        ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                                   type_str);
        group_state->is_first_type = 0;
    } else {
        ucs_table_row_add_cell_empty(table, row, 1);
    }

    if (group_state->is_first_cmpt) {
        ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                                   cmpt_str);
        group_state->is_first_cmpt = 0;
    } else {
        ucs_table_row_add_cell_empty(table, row, 1);
    }

    if (group_state->is_first_tl) {
        ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                                   tl_str);
        group_state->is_first_tl = 0;
    } else {
        ucs_table_row_add_cell_empty(table, row, 1);
    }

    ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                               dev_str);
}

static void ucp_tl_info_add_title(ucs_table_t *table, ucp_context_h context)
{
    UCS_STRING_BUFFER_ONSTACK(title_strb, 128);
    ucs_table_row_h row;

    ucs_string_buffer_appendf(&title_strb, "Available Transports and Devices");
    if (!ucs_string_is_empty(context->name)) {
        ucs_string_buffer_appendf(&title_strb, " (ctx: %s)", context->name);
    }

    /* Title spans all body columns. */
    ucs_table_add_row(table, &row);
    ucs_table_row_add_cell_fmt(table, row, UCP_TL_INFO_NUM_COLS,
                               UCS_TABLE_ALIGN_LEFT, "%s",
                               ucs_string_buffer_cstr(&title_strb));
    ucs_table_add_separator(table);
}

static void ucp_tl_info_add_headers(ucs_table_t *table)
{
    ucs_table_row_h row;

    ucs_table_add_row(table, &row);
    ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                               UCP_TL_INFO_HDR_TYPE);
    ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                               UCP_TL_INFO_HDR_COMPONENT);
    ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                               UCP_TL_INFO_HDR_TRANSPORT);
    ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT, "%s",
                               UCP_TL_INFO_HDR_DEVICE);
    ucs_table_add_separator(table);
}

static unsigned
ucp_tl_info_tl_group_end(const ucp_tl_info_array_t *all_rscs, unsigned start)
{
    const ucp_tl_info_entry_t *start_entry = &ucs_array_elem(all_rscs, start);
    unsigned end;

    for (end = start + 1; end < ucs_array_length(all_rscs); ++end) {
        const ucp_tl_info_entry_t *end_entry = &ucs_array_elem(all_rscs, end);

        if (!ucp_tl_info_is_same_tl_group(start_entry, end_entry)) {
            break;
        }
    }

    return end;
}

static void ucp_tl_info_add_tl_group(ucs_table_t *table,
                                        ucp_context_h context,
                                        const ucp_tl_info_array_t *all_rscs,
                                        unsigned group_start,
                                        unsigned group_end, int *printed_any)
{
    UCS_STRING_BUFFER_ONSTACK(tl_strb, UCT_TL_NAME_MAX + 8);
    UCS_STRING_BUFFER_ONSTACK(dev_strb, 512);
    const ucp_tl_info_entry_t *first_entry = &ucs_array_elem(all_rscs,
                                                             group_start);
    uct_device_type_t dev_type             = first_entry->rsc.dev_type;
    ucp_rsc_index_t cmpt_idx               = first_entry->cmpt_index;
    int tl_enabled                         = 0;
    ucp_tl_info_tl_group_state_t group_state;
    const ucp_tl_info_entry_t *entry;
    unsigned i, dev_count;
    const char *sysdev_name;
    const char *mark;

    group_state.is_first_type =
            (group_start == 0) ||
            (ucs_array_elem(all_rscs, group_start - 1).rsc.dev_type !=
             dev_type);
    group_state.is_first_cmpt =
            group_state.is_first_type ||
            (ucs_array_elem(all_rscs, group_start - 1).cmpt_index != cmpt_idx);
    group_state.is_first_tl   = 1;

    if (*printed_any) {
        /* Carry-over cols on the separator above this row:
         * 2 for a new TL in the same component, 1 for a
         * new component in the same dev_type, 0 otherwise. */
        ucs_table_add_separator_with_merged_cols(
                table, !group_state.is_first_cmpt ?
                               2 :
                               (group_state.is_first_type ? 0 : 1));
    }

    /* The transport is enabled if any of the devices in the group is enabled */
    for (i = group_start; i < group_end; ++i) {
        if (ucs_array_elem(all_rscs, i).enabled) {
            tl_enabled = 1;
            break;
        }
    }

    ucs_string_buffer_appendf(&tl_strb, "%s %s",
                              tl_enabled ? UCP_TL_INFO_MARK_ENABLED :
                                           UCP_TL_INFO_MARK_DISABLED,
                              first_entry->rsc.tl_name);

    for (i = group_start; i < group_end; ++i) {
        entry     = &ucs_array_elem(all_rscs, i);
        dev_count = i - group_start + 1;

        /* Add space between devices on the same line. */
        if (ucs_string_buffer_length(&dev_strb) > 0) {
            ucs_string_buffer_appendf(&dev_strb, "  ");
        }

        mark = entry->enabled ? UCP_TL_INFO_MARK_ENABLED :
                                UCP_TL_INFO_MARK_DISABLED;

        sysdev_name = (entry->rsc.sys_device != UCS_SYS_DEVICE_ID_UNKNOWN) ?
                      ucs_topo_sys_device_get_name(entry->rsc.sys_device) :
                      NULL;

        /* Show the system device only when it adds information: skip it
         * when the device name already starts with the system device name. */
        if ((sysdev_name != NULL) && (strncmp(entry->rsc.dev_name, sysdev_name,
                                              strlen(sysdev_name)) != 0)) {
            ucs_string_buffer_appendf(&dev_strb, "%s %s (%s)", mark,
                                      entry->rsc.dev_name, sysdev_name);
        } else {
            ucs_string_buffer_appendf(&dev_strb, "%s %s", mark,
                                      entry->rsc.dev_name);
        }

        /* Emit once a line is full or this is the last device in the group. */
        if ((dev_count % UCP_TL_INFO_DEVS_PER_LINE == 0) ||
            (i == (group_end - 1))) {
            ucp_tl_info_emit_row(table, uct_device_type_names[dev_type],
                                 context->tl_cmpts[cmpt_idx].attr.name,
                                 ucs_string_buffer_cstr(&tl_strb),
                                 ucs_string_buffer_cstr(&dev_strb),
                                 &group_state);
            ucs_string_buffer_reset(&dev_strb);
            *printed_any = 1;
        }
    }
}

static void
ucp_tl_info_add_unavailable_cmpts(ucs_table_t *table, ucp_context_h context,
                                     const ucp_tl_info_array_t *all_rscs,
                                     int *printed_any)
{
    int first_unavail = 1;
    ucp_rsc_index_t cmpt_idx;
    ucs_table_row_h row;

    for (cmpt_idx = 0; cmpt_idx < context->num_cmpts; ++cmpt_idx) {
        if (ucp_tl_info_rscs_has_cmpt(all_rscs, cmpt_idx)) {
            continue;
        }

        /* Don't list components with no memory domains. (e.g. rdmacm) */
        if (context->tl_cmpts[cmpt_idx].attr.md_resource_count == 0) {
            continue;
        }

        if (first_unavail) {
            if (*printed_any) {
                ucs_table_add_separator(table);
            }

            ucs_table_add_row(table, &row);
            ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT,
                                       "%s", UCP_TL_INFO_UNAVAILABLE);
            ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT,
                                       "%s",
                                       context->tl_cmpts[cmpt_idx].attr.name);
            ucs_table_row_add_cell_empty(table, row, 1);
            ucs_table_row_add_cell_empty(table, row, 1);
            first_unavail = 0;
        } else {
            /* Carry over the empty "type" column with merged_cols=1
             * so the separator stays blank above it. */
            ucs_table_add_separator_with_merged_cols(table, 1);
            ucs_table_add_row(table, &row);
            ucs_table_row_add_cell_empty(table, row, 1);
            ucs_table_row_add_cell_fmt(table, row, 1, UCS_TABLE_ALIGN_LEFT,
                                       "%s",
                                       context->tl_cmpts[cmpt_idx].attr.name);
            ucs_table_row_add_cell_empty(table, row, 1);
            ucs_table_row_add_cell_empty(table, row, 1);
        }

        *printed_any = 1;
    }
}

static void ucp_tl_info_add_legend(ucs_table_t *table)
{
    ucs_table_row_h row;
    unsigned i;

    ucs_table_add_separator(table);
    for (i = 0; i < ucs_static_array_size(ucp_tl_info_legend_rows); ++i) {
        ucs_table_add_row(table, &row);
        ucs_table_row_add_cell_fmt(table, row, UCP_TL_INFO_NUM_COLS,
                                   UCS_TABLE_ALIGN_LEFT, "%s",
                                   ucp_tl_info_legend_rows[i]);
    }
}

ucs_status_t ucp_context_render_tl_info(ucp_context_h context,
                                        ucp_tl_info_array_t *all_rscs,
                                        ucs_string_buffer_t *strb)
{
    const ucs_table_config_t tcfg = {
        .n_cols = UCP_TL_INFO_NUM_COLS
    };
    unsigned group_start, group_end;
    ucs_status_t status;
    ucs_table_t table;
    int printed_any;

    ucs_assertv_always(!ucs_array_is_empty(all_rscs),
                       "no resource info captured");

    /* Sort by (dev_type, component, transport) so that the rows of each group
     * are contiguous and group boundaries can be found by comparing adjacent
     * entries, instead of rescanning the whole array per group. */
    ucs_qsort_r(ucs_array_begin(all_rscs), ucs_array_length(all_rscs),
                sizeof(*ucs_array_begin(all_rscs)), ucp_tl_info_compare, NULL);

    ucs_table_init(&table, &tcfg);

    ucp_tl_info_add_title(&table, context);
    ucp_tl_info_add_headers(&table);

    printed_any = 0;
    group_start = 0;
    while (group_start < ucs_array_length(all_rscs)) {
        group_end = ucp_tl_info_tl_group_end(all_rscs, group_start);
        ucp_tl_info_add_tl_group(&table, context, all_rscs, group_start,
                                    group_end, &printed_any);
        group_start = group_end;
    }

    ucp_tl_info_add_unavailable_cmpts(&table, context, all_rscs,
                                         &printed_any);
    ucp_tl_info_add_legend(&table);

    ucs_table_render(&table, strb);

    status = ucs_table_get_status(&table);
    ucs_table_cleanup(&table);
    return status;
}

void ucp_context_log_tl_info(ucp_context_h context,
                             ucp_tl_info_array_t *all_rscs)
{
    ucs_string_buffer_t strb = UCS_STRING_BUFFER_INITIALIZER;
    ucs_status_t status;

    if (!ucp_context_print_transport_tables_enabled(context)) {
        return;
    }

    if (ucs_array_is_empty(all_rscs)) {
        ucs_warn("skipping transport table: no resource info captured");
        return;
    }

    status = ucp_context_render_tl_info(context, all_rscs, &strb);
    if (status != UCS_OK) {
        ucs_warn("transport table render incomplete: %s",
                 ucs_status_string(status));
    }

    ucs_log_print_compact(ucs_string_buffer_cstr(&strb));
    ucs_string_buffer_cleanup(&strb);
}
