/*
 *
 * Copyright (c) 2021-2022 The Khronos Group Inc.
 * Copyright (c) 2021-2022 Valve Corporation
 * Copyright (c) 2021-2022 LunarG, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Author: Mark Young <marky@lunarg.com>
 *
 */

// Non-windows and non-apple only header file, guard it so that accidental
// inclusion doesn't cause unknown header include errors
#if defined(LOADER_ENABLE_LINUX_SORT)

#include <stdio.h>
#include <stdlib.h>

#include "loader_linux.h"

#include "allocation.h"
#include "loader_environment.h"
#include "loader.h"
#include "log.h"

// Determine a priority based on device type with the higher value being higher priority.
uint32_t determine_priority_type_value(VkPhysicalDeviceType type) {
    switch (type) {
        case VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
            return 10;
        case VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
            return 5;
        case VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU:
            return 3;
        case VK_PHYSICAL_DEVICE_TYPE_OTHER:
            return 2;
        case VK_PHYSICAL_DEVICE_TYPE_CPU:
            return 1;
        case VK_PHYSICAL_DEVICE_TYPE_MAX_ENUM:  // Not really an enum, but throws warning if it's not here
            break;
    }
    return 0;
}

// Compare the two device types.
// This behaves similar to a qsort compare.
int32_t device_type_compare(VkPhysicalDeviceType a, VkPhysicalDeviceType b) {
    uint32_t a_value = determine_priority_type_value(a);
    uint32_t b_value = determine_priority_type_value(b);
    if (a_value > b_value) {
        return -1;
    } else if (b_value > a_value) {
        return 1;
    }
    return 0;
}

// Used to compare two devices and determine which one should have priority.  The criteria is
// simple:
//   1) Default device ALWAYS wins
//   2) Sort by type
//   3) Sort by PCI bus ID
//   4) Ties broken by device_ID XOR vendor_ID comparison
int32_t compare_devices(const void *a, const void *b) {
    struct LinuxSortedDeviceInfo *left = (struct LinuxSortedDeviceInfo *)a;
    struct LinuxSortedDeviceInfo *right = (struct LinuxSortedDeviceInfo *)b;

    // Default device always gets priority
    if (left->default_device) {
        return -1;
    } else if (right->default_device) {
        return 1;
    }

    // Order by device type next
    int32_t dev_type_comp = device_type_compare(left->device_type, right->device_type);
    if (0 != dev_type_comp) {
        return dev_type_comp;
    }

    // Sort by PCI info (prioritize devices that have info over those that don't)
    if (left->has_pci_bus_info && !right->has_pci_bus_info) {
        return -1;
    } else if (!left->has_pci_bus_info && right->has_pci_bus_info) {
        return 1;
    } else if (left->has_pci_bus_info && right->has_pci_bus_info) {
        // Sort low to high PCI domain
        if (left->pci_domain < right->pci_domain) {
            return -1;
        } else if (left->pci_domain > right->pci_domain) {
            return 1;
        }
        // Sort low to high PCI bus
        if (left->pci_bus < right->pci_bus) {
            return -1;
        } else if (left->pci_bus > right->pci_bus) {
            return 1;
        }
        // Sort low to high PCI device
        if (left->pci_device < right->pci_device) {
            return -1;
        } else if (left->pci_device > right->pci_device) {
            return 1;
        }
        // Sort low to high PCI function
        if (left->pci_function < right->pci_function) {
            return -1;
        } else if (left->pci_function > right->pci_function) {
            return 1;
        }
    }

    // Somehow we have a tie above, so XOR vendorID and deviceID and compare
    uint32_t left_xord_dev_vend = left->device_id ^ left->vendor_id;
    uint32_t right_xord_dev_vend = right->device_id ^ right->vendor_id;
    if (left_xord_dev_vend < right_xord_dev_vend) {
        return -1;
    } else if (right_xord_dev_vend < left_xord_dev_vend) {
        return 1;
    }
    return 0;
}

// Used to compare two device groups and determine which one should have priority.
// NOTE: This assumes that devices in each group have already been sorted.
// The group sort criteria is simple:
//   1) Group with the default device ALWAYS wins
//   2) Group with the best device type for device 0 wins
//   3) Group with best PCI bus ID for device 0 wins
//   4) Ties broken by group device 0 device_ID XOR vendor_ID comparison
int32_t compare_device_groups(const void *a, const void *b) {
    struct loader_physical_device_group_term *grp_a = (struct loader_physical_device_group_term *)a;
    struct loader_physical_device_group_term *grp_b = (struct loader_physical_device_group_term *)b;

    // Use the first GPU's info from each group to sort the groups by
    struct LinuxSortedDeviceInfo *left = &grp_a->internal_device_info[0];
    struct LinuxSortedDeviceInfo *right = &grp_b->internal_device_info[0];

    // Default device always gets priority
    if (left->default_device) {
        return -1;
    } else if (right->default_device) {
        return 1;
    }

    // Order by device type next
    int32_t dev_type_comp = device_type_compare(left->device_type, right->device_type);
    if (0 != dev_type_comp) {
        return dev_type_comp;
    }

    // Sort by PCI info (prioritize devices that have info over those that don't)
    if (left->has_pci_bus_info && !right->has_pci_bus_info) {
        return -1;
    } else if (!left->has_pci_bus_info && right->has_pci_bus_info) {
        return 1;
    } else if (left->has_pci_bus_info && right->has_pci_bus_info) {
        // Sort low to high PCI domain
        if (left->pci_domain < right->pci_domain) {
            return -1;
        } else if (left->pci_domain > right->pci_domain) {
            return 1;
        }
        // Sort low to high PCI bus
        if (left->pci_bus < right->pci_bus) {
            return -1;
        } else if (left->pci_bus > right->pci_bus) {
            return 1;
        }
        // Sort low to high PCI device
        if (left->pci_device < right->pci_device) {
            return -1;
        } else if (left->pci_device > right->pci_device) {
            return 1;
        }
        // Sort low to high PCI function
        if (left->pci_function < right->pci_function) {
            return -1;
        } else if (left->pci_function > right->pci_function) {
            return 1;
        }
    }

    // Somehow we have a tie above, so XOR vendorID and deviceID and compare
    uint32_t left_xord_dev_vend = left->device_id ^ left->vendor_id;
    uint32_t right_xord_dev_vend = right->device_id ^ right->vendor_id;
    if (left_xord_dev_vend < right_xord_dev_vend) {
        return -1;
    } else if (right_xord_dev_vend < left_xord_dev_vend) {
        return 1;
    }
    return 0;
}

// Search for the default device using the loader environment variable.
void linux_env_var_default_device(struct loader_instance *inst, uint32_t device_count,
                                  struct LinuxSortedDeviceInfo *sorted_device_info) {
    char *selection = loader_getenv("VK_LOADER_DEVICE_SELECT", inst);
    if (NULL != selection) {
        loader_log(inst, VULKAN_LOADER_DEBUG_BIT | VULKAN_LOADER_DRIVER_BIT, 0,
                   "linux_env_var_default_device:  Found \'VK_LOADER_DEVICE_SELECT\' set to %s", selection);

        // The environment variable exists, so grab the vendor ID and device ID of the
        // selected default device
        unsigned vendor_id, device_id;
        int32_t matched = sscanf(selection, "%x:%x", &vendor_id, &device_id);
        if (matched == 2) {
            for (int32_t i = 0; i < (int32_t)device_count; ++i) {
                if (sorted_device_info[i].vendor_id == vendor_id && sorted_device_info[i].device_id == device_id) {
                    loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0,
                               "linux_env_var_default_device:  Found default at index %u \'%s\'", i,
                               sorted_device_info[i].device_name);
                    sorted_device_info[i].default_device = true;
                    break;
                }
            }
        }

        loader_free_getenv(selection, inst);
    }
}

// This function allocates an array in sorted_devices which must be freed by the caller if not null
VkResult linux_read_sorted_physical_devices(struct loader_instance *inst, uint32_t icd_count,
                                            struct loader_icd_physical_devices *icd_devices, uint32_t phys_dev_count,
                                            struct loader_physical_device_term **sorted_device_term) {
    VkResult res = VK_SUCCESS;
    bool app_is_vulkan_1_1 = loader_check_version_meets_required(LOADER_VERSION_1_1_0, inst->app_api_version);

    struct LinuxSortedDeviceInfo *sorted_device_info = loader_instance_heap_calloc(
        inst, phys_dev_count * sizeof(struct LinuxSortedDeviceInfo), VK_SYSTEM_ALLOCATION_SCOPE_COMMAND);
    if (NULL == sorted_device_info) {
        res = VK_ERROR_OUT_OF_HOST_MEMORY;
        goto out;
    }

    loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "linux_read_sorted_physical_devices:");
    loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "     Original order:");

    // Grab all the necessary info we can about each device
    uint32_t index = 0;
    for (uint32_t icd_idx = 0; icd_idx < icd_count; ++icd_idx) {
        for (uint32_t phys_dev = 0; phys_dev < icd_devices[icd_idx].device_count; ++phys_dev) {
            struct loader_icd_term *icd_term = icd_devices[icd_idx].icd_term;
            VkPhysicalDeviceProperties dev_props = {};

            sorted_device_info[index].physical_device = icd_devices[icd_idx].physical_devices[phys_dev];
            sorted_device_info[index].icd_index = icd_idx;
            sorted_device_info[index].icd_term = icd_term;
            sorted_device_info[index].has_pci_bus_info = false;

            icd_term->dispatch.GetPhysicalDeviceProperties(sorted_device_info[index].physical_device, &dev_props);
            sorted_device_info[index].device_type = dev_props.deviceType;
            strncpy(sorted_device_info[index].device_name, dev_props.deviceName, VK_MAX_PHYSICAL_DEVICE_NAME_SIZE);
            sorted_device_info[index].vendor_id = dev_props.vendorID;
            sorted_device_info[index].device_id = dev_props.deviceID;

            bool device_is_1_1_capable =
                loader_check_version_meets_required(LOADER_VERSION_1_1_0, loader_make_version(dev_props.apiVersion));
            if (!sorted_device_info[index].has_pci_bus_info) {
                uint32_t ext_count = 0;
                icd_term->dispatch.EnumerateDeviceExtensionProperties(sorted_device_info[index].physical_device, NULL, &ext_count,
                                                                      NULL);
                if (ext_count > 0) {
                    VkExtensionProperties *ext_props =
                        (VkExtensionProperties *)loader_stack_alloc(sizeof(VkExtensionProperties) * ext_count);
                    if (NULL == ext_props) {
                        res = VK_ERROR_OUT_OF_HOST_MEMORY;
                        goto out;
                    }
                    icd_term->dispatch.EnumerateDeviceExtensionProperties(sorted_device_info[index].physical_device, NULL,
                                                                          &ext_count, ext_props);
                    for (uint32_t ext = 0; ext < ext_count; ++ext) {
                        if (!strcmp(ext_props[ext].extensionName, VK_EXT_PCI_BUS_INFO_EXTENSION_NAME)) {
                            sorted_device_info[index].has_pci_bus_info = true;
                            break;
                        }
                    }
                }
            }

            if (sorted_device_info[index].has_pci_bus_info) {
                VkPhysicalDevicePCIBusInfoPropertiesEXT pci_props = (VkPhysicalDevicePCIBusInfoPropertiesEXT){
                    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PCI_BUS_INFO_PROPERTIES_EXT};
                VkPhysicalDeviceProperties2 dev_props2 = (VkPhysicalDeviceProperties2){
                    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2, .pNext = (VkBaseInStructure *)&pci_props};

                PFN_vkGetPhysicalDeviceProperties2 GetPhysDevProps2 = NULL;
                if (app_is_vulkan_1_1 && device_is_1_1_capable) {
                    GetPhysDevProps2 = icd_term->dispatch.GetPhysicalDeviceProperties2;
                } else {
                    GetPhysDevProps2 = (PFN_vkGetPhysicalDeviceProperties2)icd_term->dispatch.GetPhysicalDeviceProperties2KHR;
                }
                if (NULL != GetPhysDevProps2) {
                    GetPhysDevProps2(sorted_device_info[index].physical_device, &dev_props2);
                    sorted_device_info[index].pci_domain = pci_props.pciDomain;
                    sorted_device_info[index].pci_bus = pci_props.pciBus;
                    sorted_device_info[index].pci_device = pci_props.pciDevice;
                    sorted_device_info[index].pci_function = pci_props.pciFunction;
                } else {
                    sorted_device_info[index].has_pci_bus_info = false;
                }
            }
            loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "           [%u] %s", index,
                       sorted_device_info[index].device_name);
            index++;
        }
    }

    // Select default device if set in the environment variable
    linux_env_var_default_device(inst, phys_dev_count, sorted_device_info);

    // Sort devices by PCI info
    qsort(sorted_device_info, phys_dev_count, sizeof(struct LinuxSortedDeviceInfo), compare_devices);

    // If we have a selected index, add that first.
    loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "     Sorted order:");

    // Add all others after (they've already been sorted)
    for (uint32_t dev = 0; dev < phys_dev_count; ++dev) {
        sorted_device_term[dev]->this_icd_term = sorted_device_info[dev].icd_term;
        sorted_device_term[dev]->icd_index = sorted_device_info[dev].icd_index;
        sorted_device_term[dev]->phys_dev = sorted_device_info[dev].physical_device;
        loader_set_dispatch((void *)sorted_device_term[dev], inst->disp);
        loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "           [%u] %s  %s", dev,
                   sorted_device_info[dev].device_name, (sorted_device_info[dev].default_device ? "[default]" : ""));
    }

out:
    loader_instance_heap_free(inst, sorted_device_info);

    return res;
}

// This function sorts an array of physical device groups
VkResult linux_sort_physical_device_groups(struct loader_instance *inst, uint32_t group_count,
                                           struct loader_physical_device_group_term *sorted_group_term) {
    VkResult res = VK_SUCCESS;
    bool app_is_vulkan_1_1 = loader_check_version_meets_required(LOADER_VERSION_1_1_0, inst->app_api_version);

    loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "linux_sort_physical_device_groups:  Original order:");

    for (uint32_t group = 0; group < group_count; ++group) {
        loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "           Group %u", group);

        struct loader_icd_term *icd_term = sorted_group_term[group].this_icd_term;
        for (uint32_t gpu = 0; gpu < sorted_group_term[group].group_props.physicalDeviceCount; ++gpu) {
            VkPhysicalDeviceProperties dev_props = {};

            sorted_group_term[group].internal_device_info[gpu].physical_device =
                sorted_group_term[group].group_props.physicalDevices[gpu];
            sorted_group_term[group].internal_device_info[gpu].has_pci_bus_info = false;

            icd_term->dispatch.GetPhysicalDeviceProperties(sorted_group_term[group].internal_device_info[gpu].physical_device,
                                                           &dev_props);
            sorted_group_term[group].internal_device_info[gpu].device_type = dev_props.deviceType;
            strncpy(sorted_group_term[group].internal_device_info[gpu].device_name, dev_props.deviceName,
                    VK_MAX_PHYSICAL_DEVICE_NAME_SIZE);
            sorted_group_term[group].internal_device_info[gpu].vendor_id = dev_props.vendorID;
            sorted_group_term[group].internal_device_info[gpu].device_id = dev_props.deviceID;

            bool device_is_1_1_capable =
                loader_check_version_meets_required(LOADER_VERSION_1_1_0, loader_make_version(dev_props.apiVersion));
            if (!sorted_group_term[group].internal_device_info[gpu].has_pci_bus_info) {
                uint32_t ext_count;
                icd_term->dispatch.EnumerateDeviceExtensionProperties(
                    sorted_group_term[group].internal_device_info[gpu].physical_device, NULL, &ext_count, NULL);
                if (ext_count > 0) {
                    VkExtensionProperties *ext_props =
                        (VkExtensionProperties *)loader_stack_alloc(sizeof(VkExtensionProperties) * ext_count);
                    if (NULL == ext_props) {
                        return VK_ERROR_OUT_OF_HOST_MEMORY;
                    }
                    icd_term->dispatch.EnumerateDeviceExtensionProperties(
                        sorted_group_term[group].internal_device_info[gpu].physical_device, NULL, &ext_count, ext_props);
                    for (uint32_t ext = 0; ext < ext_count; ++ext) {
                        if (!strcmp(ext_props[ext].extensionName, VK_EXT_PCI_BUS_INFO_EXTENSION_NAME)) {
                            sorted_group_term[group].internal_device_info[gpu].has_pci_bus_info = true;
                            break;
                        }
                    }
                }
            }

            if (sorted_group_term[group].internal_device_info[gpu].has_pci_bus_info) {
                VkPhysicalDevicePCIBusInfoPropertiesEXT pci_props = (VkPhysicalDevicePCIBusInfoPropertiesEXT){
                    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PCI_BUS_INFO_PROPERTIES_EXT};
                VkPhysicalDeviceProperties2 dev_props2 = (VkPhysicalDeviceProperties2){
                    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2, .pNext = (VkBaseInStructure *)&pci_props};

                PFN_vkGetPhysicalDeviceProperties2 GetPhysDevProps2 = NULL;
                if (app_is_vulkan_1_1 && device_is_1_1_capable) {
                    GetPhysDevProps2 = icd_term->dispatch.GetPhysicalDeviceProperties2;
                } else {
                    GetPhysDevProps2 = (PFN_vkGetPhysicalDeviceProperties2)icd_term->dispatch.GetPhysicalDeviceProperties2KHR;
                }
                if (NULL != GetPhysDevProps2) {
                    GetPhysDevProps2(sorted_group_term[group].internal_device_info[gpu].physical_device, &dev_props2);
                    sorted_group_term[group].internal_device_info[gpu].pci_domain = pci_props.pciDomain;
                    sorted_group_term[group].internal_device_info[gpu].pci_bus = pci_props.pciBus;
                    sorted_group_term[group].internal_device_info[gpu].pci_device = pci_props.pciDevice;
                    sorted_group_term[group].internal_device_info[gpu].pci_function = pci_props.pciFunction;
                } else {
                    sorted_group_term[group].internal_device_info[gpu].has_pci_bus_info = false;
                }
            }
            loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "               [%u] %s", gpu,
                       sorted_group_term[group].internal_device_info[gpu].device_name);
        }

        // Select default device if set in the environment variable
        linux_env_var_default_device(inst, sorted_group_term[group].group_props.physicalDeviceCount,
                                     sorted_group_term[group].internal_device_info);

        // Sort GPUs in each group
        qsort(sorted_group_term[group].internal_device_info, sorted_group_term[group].group_props.physicalDeviceCount,
              sizeof(struct LinuxSortedDeviceInfo), compare_devices);

        // Match the externally used physical device list with the sorted physical device list for this group.
        for (uint32_t dev = 0; dev < sorted_group_term[group].group_props.physicalDeviceCount; ++dev) {
            sorted_group_term[group].group_props.physicalDevices[dev] =
                sorted_group_term[group].internal_device_info[dev].physical_device;
        }
    }

    // Sort device groups by PCI info
    qsort(sorted_group_term, group_count, sizeof(struct loader_physical_device_group_term), compare_device_groups);

    loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "linux_sort_physical_device_groups:  Sorted order:");
    for (uint32_t group = 0; group < group_count; ++group) {
        loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "           Group %u", group);
        for (uint32_t gpu = 0; gpu < sorted_group_term[group].group_props.physicalDeviceCount; ++gpu) {
            loader_log(inst, VULKAN_LOADER_INFO_BIT | VULKAN_LOADER_DRIVER_BIT, 0, "               [%u] %s %p %s", gpu,
                       sorted_group_term[group].internal_device_info[gpu].device_name,
                       sorted_group_term[group].internal_device_info[gpu].physical_device,
                       (sorted_group_term[group].internal_device_info[gpu].default_device ? "[default]" : ""));
        }
    }

    return res;
}

#endif  // LOADER_ENABLE_LINUX_SORT