/*
 * Copyright (C) Intel Corporation, 2023-2026. ALL RIGHTS RESERVED.
 * See file LICENSE for terms.
 */

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

#include "ze_base.h"

#include <ucs/sys/module.h>
#include <ucs/sys/string.h>
#include <ucs/debug/log.h>
#include <ucs/sys/math.h>


#define UCT_ZE_MAX_DEVICES 32 /* Max root devices (GPUs) */


/* Global state */
typedef struct {
    ze_driver_handle_t driver;
    uct_ze_device_t    devices[UCT_ZE_MAX_DEVICES];
    int                num_devices;
    int                num_subdevices; /* Total sub-devices across all devices */
    ucs_init_once_t    init_once;
    ze_result_t        init_status; /* Store init result for later calls */
} uct_ze_base_state_t;


static uct_ze_base_state_t uct_ze_base = {
    .init_once   = UCS_INIT_ONCE_INITIALIZER,
    .init_status = ZE_RESULT_ERROR_UNINITIALIZED
};


typedef struct {
    ucs_sys_bus_id_t bus_id;
    ucs_sys_device_t sys_dev;
    int              valid;
} uct_ze_base_pci_info_t;


/* Static sub-device array - populated during init, read-only after */
static uct_ze_subdevice_t
        ze_subdevices[UCT_ZE_MAX_DEVICES * UCT_ZE_MAX_SUBDEVICES];
static int ze_num_subdevices = 0;


static void uct_ze_base_pci_info_init(uct_ze_base_pci_info_t *pci_info)
{
    int i;

    for (i = 0; i < UCT_ZE_MAX_DEVICES; i++) {
        pci_info[i].valid = 0;
    }
}

ze_driver_handle_t uct_ze_base_get_driver(void)
{
    if (uct_ze_base_init() != ZE_RESULT_SUCCESS) {
        return NULL;
    }

    return uct_ze_base.driver;
}

/**
 * Get PCI properties from device using separate API call
 * Compatible with older Level Zero versions
 */
static ze_result_t
uct_ze_get_pci_properties(ze_device_handle_t device, ucs_sys_bus_id_t *bus_id)
{
    ze_pci_ext_properties_t pci_props = {
        .stype = ZE_STRUCTURE_TYPE_PCI_EXT_PROPERTIES,
        .pNext = NULL
    };
    ze_result_t ret;

    ret = zeDevicePciGetPropertiesExt(device, &pci_props);
    if (ret != ZE_RESULT_SUCCESS) {
        return ret;
    }

    bus_id->domain   = (uint16_t)pci_props.address.domain;
    bus_id->bus      = (uint8_t)pci_props.address.bus;
    bus_id->slot     = (uint8_t)pci_props.address.device;
    bus_id->function = (uint8_t)pci_props.address.function;
    return ZE_RESULT_SUCCESS;
}

static ze_result_t
uct_ze_base_get_root_devices(ze_device_handle_t *root_devices,
                             uint32_t *root_dev_count_p)
{
    uint32_t driver_count = 0;
    ze_driver_handle_t *drivers;
    ze_result_t ret;
    uint32_t dev_count;
    uint32_t i;

    ret = zeInit(ZE_INIT_FLAG_GPU_ONLY);
    if (ret != ZE_RESULT_SUCCESS) {
        ucs_debug("failure to initialize ze library: 0x%x", ret);
        return ret;
    }

    ret = zeDriverGet(&driver_count, NULL);
    if ((ret != ZE_RESULT_SUCCESS) || (driver_count == 0)) {
        ucs_debug("failure to get ze driver count: 0x%x, count=%u",
                  ret, driver_count);
        return (ret != ZE_RESULT_SUCCESS) ? ret : ZE_RESULT_ERROR_UNKNOWN;
    }

    drivers = ucs_calloc(driver_count, sizeof(*drivers), "ze_drivers");
    if (drivers == NULL) {
        ucs_debug("failed to allocate ze drivers array, count=%u",
                  driver_count);
        return ZE_RESULT_ERROR_UNKNOWN;
    }

    ret = zeDriverGet(&driver_count, drivers);
    if (ret != ZE_RESULT_SUCCESS) {
        ucs_debug("failure to enumerate ze drivers: 0x%x", ret);
        ucs_free(drivers);
        return ret;
    }

    for (i = 0; i < driver_count; i++) {
        dev_count = 0;
        ret       = zeDeviceGet(drivers[i], &dev_count, NULL);
        if ((ret != ZE_RESULT_SUCCESS) || (dev_count == 0)) {
            continue;
        }

        if (dev_count > UCT_ZE_MAX_DEVICES) {
            ucs_warn("ze returned %u devices, limiting to %u", dev_count,
                     UCT_ZE_MAX_DEVICES);
            dev_count = UCT_ZE_MAX_DEVICES;
        }

        *root_dev_count_p = dev_count;
        ret               = zeDeviceGet(drivers[i], root_dev_count_p,
                                        root_devices);
        if (ret == ZE_RESULT_SUCCESS) {
            uct_ze_base.driver = drivers[i];
            ucs_free(drivers);
            return ZE_RESULT_SUCCESS;
        }

        ucs_debug("failure to get ze devices from driver %u: 0x%x", i, ret);
    }

    ucs_debug("no ze driver exposed any devices (tried %u drivers)",
              driver_count);
    ucs_free(drivers);
    return ZE_RESULT_ERROR_UNKNOWN;
}

static ze_result_t
uct_ze_base_init_device_props(uct_ze_device_t *device,
                              ze_device_handle_t root_device,
                              int device_idx)
{
    ze_result_t ret;

    device->root_device        = root_device;
    device->device_index       = device_idx;
    device->device_props.stype = ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES;
    device->device_props.pNext = NULL;

    ret = zeDeviceGetProperties(device->root_device, &device->device_props);
    if (ret != ZE_RESULT_SUCCESS) {
        ucs_debug("failure to get ze device properties for device %d: 0x%x",
                  device_idx, ret);
        return ret;
    }

    return ZE_RESULT_SUCCESS;
}

static void
uct_ze_base_set_device_sys_dev(uct_ze_device_t *device, int device_idx,
                               const ucs_sys_bus_id_t *bus_id,
                               uct_ze_base_pci_info_t *pci_info)
{
    int found_idx = -1;
    ucs_status_t status;
    int i;
    char name[16];

    for (i = 0; i < device_idx; i++) {
        if ((pci_info[i].valid != 0) &&
            (pci_info[i].bus_id.domain == bus_id->domain) &&
            (pci_info[i].bus_id.bus == bus_id->bus) &&
            (pci_info[i].bus_id.slot == bus_id->slot) &&
            (pci_info[i].bus_id.function == bus_id->function)) {
            found_idx = i;
            break;
        }
    }

    if (found_idx >= 0) {
        device->sys_dev = pci_info[found_idx].sys_dev;
        ucs_debug("device %d shares sys_dev %u with device %d (same pci)",
                  device_idx, device->sys_dev, found_idx);
    } else {
        status = ucs_topo_find_device_by_bus_id_and_user_value(
                bus_id, (uintptr_t)device_idx, &device->sys_dev);
        if (status != UCS_OK) {
            ucs_debug("ucs_topo_find_device_by_bus_id_and_user_value failed "
                      "for device %d",
                      device_idx);
            device->sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
        } else {
            ucs_snprintf_safe(name, sizeof(name), "GPU%d", device_idx);
            ucs_topo_sys_device_set_name(device->sys_dev, name, 10);

            status = ucs_topo_sys_device_enable_aux_path(device->sys_dev);
            if (status != UCS_OK) {
                ucs_debug("ucs_topo_sys_device_enable_aux_path failed "
                          "for device %d",
                          device_idx);
            }
        }

        pci_info[device_idx].bus_id  = *bus_id;
        pci_info[device_idx].sys_dev = device->sys_dev;
        pci_info[device_idx].valid   = 1;
    }
}

static void
uct_ze_base_init_device_subdevices(uct_ze_device_t *device, int device_idx)
{
    uint32_t subdev_count = 0;
    ze_result_t ret;
    ze_result_t subret;

    subret = zeDeviceGetSubDevices(device->root_device, &subdev_count, NULL);
    if ((subret == ZE_RESULT_SUCCESS) && (subdev_count > 0)) {
        if (subdev_count > UCT_ZE_MAX_SUBDEVICES) {
            ucs_warn("device %d has %u sub-devices, limiting to %d",
                     device_idx, subdev_count, UCT_ZE_MAX_SUBDEVICES);
            subdev_count = UCT_ZE_MAX_SUBDEVICES;
        }

        ret = zeDeviceGetSubDevices(device->root_device, &subdev_count,
                                    device->subdevices);
        if (ret != ZE_RESULT_SUCCESS) {
            ucs_debug("failure to get ze sub-devices for device %d: 0x%x",
                      device_idx, ret);
            device->num_subdevices = 1;
            device->subdevices[0]  = device->root_device;
        } else {
            device->num_subdevices = (int)subdev_count;
            ucs_debug("device %d has %d sub-devices (hierarchical)",
                      device_idx, device->num_subdevices);
        }
    } else {
        device->num_subdevices = 1;
        device->subdevices[0]  = device->root_device;
        ucs_debug("device %d is single sub-device (flat)", device_idx);
    }
}

static void
uct_ze_base_add_global_subdevices(uct_ze_device_t *device,
                                  int *global_subdevice_id_p)
{
    uct_ze_subdevice_t *subdevice;
    int subdevice_idx;

    for (subdevice_idx = 0; subdevice_idx < device->num_subdevices;
         subdevice_idx++) {
        if (*global_subdevice_id_p >= ucs_static_array_size(ze_subdevices)) {
            ucs_error("too many sub-devices! max %zu",
                      ucs_static_array_size(ze_subdevices));
            device->num_subdevices = subdevice_idx;
            break;
        }

        subdevice                = &ze_subdevices[*global_subdevice_id_p];
        subdevice->device        = device;
        subdevice->subdevice_idx = subdevice_idx;
        subdevice->global_id     = *global_subdevice_id_p;

        (*global_subdevice_id_p)++;
    }
}

/**
 * Initialize ZE driver and enumerate all devices and sub-devices
 * Thread-safe via UCS_INIT_ONCE
 */
ze_result_t uct_ze_base_init(void)
{
    ze_result_t ret         = ZE_RESULT_SUCCESS;
    uint32_t root_dev_count = 0;
    int global_subdevice_id = 0;
    ze_device_handle_t root_devices[UCT_ZE_MAX_DEVICES];
    uct_ze_base_pci_info_t pci_info[UCT_ZE_MAX_DEVICES];
    ucs_sys_bus_id_t bus_id;
    uct_ze_device_t *device;
    int device_idx;

    UCS_INIT_ONCE(&uct_ze_base.init_once) {
        uct_ze_base.num_devices    = 0;
        uct_ze_base.num_subdevices = 0;
        ze_num_subdevices          = 0;

        ret = uct_ze_base_get_root_devices(root_devices, &root_dev_count);
        if (ret != ZE_RESULT_SUCCESS) {
            uct_ze_base.init_status = ret;
            continue;
        }

        uct_ze_base.num_devices = (int)root_dev_count;
        ucs_debug("found %d ze root devices", uct_ze_base.num_devices);

        uct_ze_base_pci_info_init(pci_info);

        for (device_idx = 0; device_idx < uct_ze_base.num_devices;
             device_idx++) {
            device = &uct_ze_base.devices[device_idx];

            ret = uct_ze_base_init_device_props(device,
                                                root_devices[device_idx],
                                                device_idx);
            if (ret != ZE_RESULT_SUCCESS) {
                continue;
            }

            ret = uct_ze_get_pci_properties(device->root_device, &bus_id);
            if (ret != ZE_RESULT_SUCCESS) {
                ucs_debug("failure to get pci properties for device %d: 0x%x"
                          ", using sys_dev=unknown",
                          device_idx, ret);
                device->sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
            } else {
                uct_ze_base_set_device_sys_dev(device, device_idx, &bus_id,
                                               pci_info);
            }
            uct_ze_base_init_device_subdevices(device, device_idx);
            uct_ze_base_add_global_subdevices(device, &global_subdevice_id);

            uct_ze_base.num_subdevices += device->num_subdevices;
        }

        ze_num_subdevices       = global_subdevice_id;
        uct_ze_base.init_status = ZE_RESULT_SUCCESS;

        ucs_debug("ze init complete: %d devices, %d total sub-devices",
                  uct_ze_base.num_devices, ze_num_subdevices);
    }

    return uct_ze_base.init_status;
}

const uct_ze_subdevice_t *uct_ze_base_get_subdevice_by_global_id(int global_id)
{
    if (uct_ze_base_init() != ZE_RESULT_SUCCESS) {
        return NULL;
    }

    if ((global_id < 0) || (global_id >= ze_num_subdevices)) {
        return NULL;
    }

    return &ze_subdevices[global_id];
}

ze_device_handle_t
uct_ze_base_get_device_handle_from_subdevice(const uct_ze_subdevice_t *subdevice)
{
    if ((subdevice == NULL) || (subdevice->device == NULL)) {
        return NULL;
    }

    if ((subdevice->subdevice_idx < 0) ||
        (subdevice->subdevice_idx >= subdevice->device->num_subdevices)) {
        return NULL;
    }

    return subdevice->device->subdevices[subdevice->subdevice_idx];
}

ucs_sys_device_t uct_ze_base_get_sys_dev_from_handle(ze_device_handle_t device)
{
    const uct_ze_device_t *dev;
    int i, j;

    if ((device == NULL) || (uct_ze_base_init() != ZE_RESULT_SUCCESS)) {
        return UCS_SYS_DEVICE_ID_UNKNOWN;
    }

    for (i = 0; i < uct_ze_base.num_devices; i++) {
        dev = &uct_ze_base.devices[i];

        /* Skip entries not fully initialized during discovery. */
        if (dev->num_subdevices <= 0) {
            continue;
        }

        if (dev->root_device == device) {
            return dev->sys_dev;
        }

        for (j = 0; j < dev->num_subdevices; j++) {
            if (dev->subdevices[j] == device) {
                return dev->sys_dev;
            }
        }
    }

    ucs_trace("cannot map ze device handle %p to sys_dev, "
              "falling back to unknown",
              (void*)device);

    return UCS_SYS_DEVICE_ID_UNKNOWN;
}

/**
 * Query MD resources - returns one MD per sub-device
 * This is correct because each sub-device has separate memory
 */
ucs_status_t
uct_ze_base_query_md_resources(uct_component_h component,
                               uct_md_resource_desc_t **resources_p,
                               unsigned *num_resources_p)
{
    if ((uct_ze_base_init() != ZE_RESULT_SUCCESS) ||
        (ze_num_subdevices == 0)) {
        ucs_debug("ze initialization failed or no sub-devices, "
                  "returning empty resources");
        return uct_md_query_empty_md_resource(resources_p, num_resources_p);
    }

    /* Return single MD resource - actual sub-device selection happens in md_open */
    return uct_md_query_single_md_resource(component, resources_p,
                                           num_resources_p);
}

/**
 * Query devices for transport layer - returns one device per sub-device
 * Each sub-device is a separate memory domain but shares sys_dev with siblings
 */
ucs_status_t uct_ze_base_query_devices(uct_md_h md,
                                       uct_tl_device_resource_t **tl_devices_p,
                                       unsigned *num_tl_devices_p)
{
    uct_tl_device_resource_t *resources;
    const uct_ze_subdevice_t *subdevice;
    const uct_ze_device_t *device;
    int i;

    if ((uct_ze_base_init() != ZE_RESULT_SUCCESS) ||
        (ze_num_subdevices == 0)) {
        *tl_devices_p     = NULL;
        *num_tl_devices_p = 0;
        return UCS_OK;
    }

    resources = ucs_calloc(ze_num_subdevices, sizeof(*resources),
                           "ze_tl_devices");
    if (resources == NULL) {
        return UCS_ERR_NO_MEMORY;
    }

    /* Return one entry per sub-device */
    for (i = 0; i < ze_num_subdevices; i++) {
        subdevice = &ze_subdevices[i];
        device    = subdevice->device;

        /* Name format: "GPU0" for single sub-device, "GPU0.0" for multi sub-device */
        if (device->num_subdevices == 1) {
            ucs_snprintf_safe(resources[i].name, UCT_DEVICE_NAME_MAX, "GPU%d",
                              device->device_index);
        } else {
            ucs_snprintf_safe(resources[i].name, UCT_DEVICE_NAME_MAX,
                              "GPU%d.%d", device->device_index,
                              subdevice->subdevice_idx);
        }

        resources[i].type = UCT_DEVICE_TYPE_ACC;

        /* CRITICAL: All sub-devices on same device share the same sys_dev */
        /* This enables correct IB affinity for all sub-devices */
        resources[i].sys_device = device->sys_dev;

        ucs_debug("sub-device %d: name=%s sys_dev=%u (device %d, sub-device "
                  "%d/%d)",
                  i, resources[i].name, resources[i].sys_device,
                  device->device_index, subdevice->subdevice_idx,
                  device->num_subdevices);
    }

    *tl_devices_p     = resources;
    *num_tl_devices_p = ze_num_subdevices;

    return UCS_OK;
}

/* Cleanup function - called on module unload */
static void UCS_F_DTOR uct_ze_base_cleanup(void)
{
    /* Note: ze_subdevices array is static, no cleanup needed */
    /* ZE driver cleanup happens automatically on process exit */
}

UCS_MODULE_INIT()
{
    return UCS_OK;
}
