#include <linux/pci.h>
#include <linux/module.h>
#include <linux/init.h>
#include <linux/kmod.h>
#include <linux/delay.h>

MODULE_AUTHOR("Majd Dibbiny <majd@mellanox.com>");
MODULE_DESCRIPTION("Mellanox ConnectX family PPC PCI reset module");
MODULE_LICENSE("Dual BSD/GPL");
#define MAX_NUM_OF_DEVS 8
#define MAX_DEV_NAME_LEN 128
#define MAX_POLLING_NUM 100
#define UNKNOWN_DEVICE_ID 0xffff

struct reset_info
{
    int num_of_pci_devs;
    int num_of_found_pci_devs;
    char pdev_arr[MAX_NUM_OF_DEVS][MAX_DEV_NAME_LEN];
    struct pci_dev* devs[MAX_NUM_OF_DEVS];
    int was_reset;
};

struct reset_info g_reset_info = {0};

char pci_device[MAX_DEV_NAME_LEN * MAX_NUM_OF_DEVS];
module_param_string(pci_dev, pci_device, sizeof(pci_device), 0444);

int reset_device(void)
{
    int err = 0;
    int ix;
    struct pci_dev* pdev;

    if (g_reset_info.was_reset == 1)
    {
        return 0;
    }

    for (ix = 0; ix < g_reset_info.num_of_pci_devs; ix++)
    {
        pdev = g_reset_info.devs[ix];
        err = pci_enable_device(pdev);

        if (err)
        {
            printk(KERN_ERR "%s %s %d Reset failed for device: %s - err: %d\n", dev_driver_string(&pdev->dev), __func__,
                   __LINE__, dev_name(&pdev->dev), err);
            return err;
        }
        pci_set_master(pdev);
        err = pci_save_state(pdev);
        if (err)
        {
            printk(KERN_ERR "%s %s %d Save state failed for device: %s - err: %d\n", dev_driver_string(&pdev->dev),
                   __func__, __LINE__, dev_name(&pdev->dev), err);
            return err;
        }
    }

    for (ix = 0; ix < g_reset_info.num_of_pci_devs; ix++)
    {
        pdev = g_reset_info.devs[ix];
        if (PCI_FUNC(pdev->devfn) == 0)
        {
            printk(KERN_INFO "%s %s %d Send hot reset to  device: %s\n", dev_driver_string(&pdev->dev), __func__,
                   __LINE__, dev_name(&pdev->dev));
            err = pci_set_pcie_reset_state(pdev, pcie_hot_reset);
            if (err)
            {
                printk(KERN_ERR "%s %s %d Set PCIE hot reset state failed for device: %s - err: %d\n",
                       dev_driver_string(&pdev->dev), __func__, __LINE__, dev_name(&pdev->dev), err);
                return err;
            }
        }
    }

    msleep(jiffies_to_msecs(HZ / 2));
    msleep(jiffies_to_msecs(HZ / 2));

    for (ix = 0; ix < g_reset_info.num_of_pci_devs; ix++)
    {
        pdev = g_reset_info.devs[ix];
        if (PCI_FUNC(pdev->devfn) == 0)
        {
            printk(KERN_INFO "%s %s %d Deassert device: %s\n", dev_driver_string(&pdev->dev), __func__, __LINE__,
                   dev_name(&pdev->dev));
            err = pci_set_pcie_reset_state(pdev, pcie_deassert_reset);
            if (err)
            {
                printk(KERN_ERR "%s %s %d Set PCIE deassert reset state failed for device: %s - err: %d\n",
                       dev_driver_string(&pdev->dev), __func__, __LINE__, dev_name(&pdev->dev), err);
                return err;
            }
        }
    }

    // Wait for the device to response to PCI configuration cycles.
    for (ix = 0; ix < g_reset_info.num_of_pci_devs; ix++)
    {
        unsigned short device_id = UNKNOWN_DEVICE_ID;
        int polling_counter = 0;

        // Get the current device.
        pdev = g_reset_info.devs[ix];

        // Device id still unknown ?
        while (device_id != pdev->device)
        {
            // 100ms is the minimum time that prevents error logs on
            //   dmesg (device is not ready for PCI configuration cycles).
            msleep(100);

            // Read the device id.
            // Access can fail (if device is not ready) and
            //   as a result we might get errors in dmesg.
            pci_read_config_word(pdev, PCI_DEVICE_ID, &device_id);

            // Polling counter violation.
            if (polling_counter > MAX_POLLING_NUM)
            {
                printk(KERN_ERR
                       "%s %s %d Polling on device id failed: reached max value of polling failures for device: %s\n",
                       dev_driver_string(&pdev->dev), __func__, __LINE__, dev_name(&pdev->dev));
                break;
            }

            polling_counter++;
        }
    }

    for (ix = 0; ix < g_reset_info.num_of_pci_devs; ix++)
    {
        pdev = g_reset_info.devs[ix];
        pci_restore_state(pdev);
    }

    g_reset_info.was_reset = 1;
    msleep(jiffies_to_msecs(HZ / 2));
    msleep(jiffies_to_msecs(HZ / 2));
    return 0;
}

static int init_one(struct pci_dev* pdev, const struct pci_device_id* id)
{
    int ix = 0;

    for (ix = 0; ix < g_reset_info.num_of_pci_devs; ix++)
    {
        if (!strcmp(g_reset_info.pdev_arr[ix], dev_name(&pdev->dev)))
        {
            g_reset_info.devs[g_reset_info.num_of_found_pci_devs] = pdev;
            g_reset_info.num_of_found_pci_devs++;
        }
    }

    if (g_reset_info.num_of_found_pci_devs == g_reset_info.num_of_pci_devs)
    {
        return reset_device();
    }
    return 0;
}

static void remove_one(struct pci_dev* pdev)
{
    int ix = 0;
    for (ix = 0; ix < g_reset_info.num_of_pci_devs; ix++)
    {
        if (!strcmp(g_reset_info.pdev_arr[ix], dev_name(&pdev->dev)))
        {
            pci_clear_master(pdev);
            pci_disable_device(pdev);
            return;
        }
    }
}

static const struct pci_device_id mlx5_core_pci_table[] = {{PCI_VDEVICE(MELLANOX, 0x1011)}, // Connect-IB
                                                           {PCI_VDEVICE(MELLANOX, 0x1013)}, // ConnectX-4
                                                           {PCI_VDEVICE(MELLANOX, 0x1015)}, // ConnectX-4LX
                                                           {PCI_VDEVICE(MELLANOX, 0x1017)}, // ConnectX-5
                                                           {PCI_VDEVICE(MELLANOX, 0x1019)}, // ConnectX-5EX
                                                           {PCI_VDEVICE(MELLANOX, 0x101B)}, // ConnectX-6
                                                           {PCI_VDEVICE(MELLANOX, 0x101D)}, // ConnectX-6DX
                                                           {PCI_VDEVICE(MELLANOX, 0x101F)}, // ConnectX-6LX
                                                           {PCI_VDEVICE(MELLANOX, 0x1021)}, // ConnectX-7
                                                           {PCI_VDEVICE(MELLANOX, 0x1023)}, // ConnectX-8
                                                           {PCI_VDEVICE(MELLANOX, 0xA2D2)}, // BlueField
                                                           {PCI_VDEVICE(MELLANOX, 0xA2D6)}, // BlueField2
                                                           {PCI_VDEVICE(MELLANOX, 0xA2DC)}, // BlueField3
                                                           {PCI_VDEVICE(MELLANOX, 0xA2DE)}, // BlueField4
                                                           {PCI_VDEVICE(MELLANOX, 0x2900)}, // GB-100
                                                           {
                                                             0,
                                                           }};

MODULE_DEVICE_TABLE(pci, mlx5_core_pci_table);

static struct pci_driver mst_ppc_pci_reset_driver = {
  .name = "mst_ppc_pci_reset_driver",
  .id_table = mlx5_core_pci_table,
  .probe = init_one,
  .remove = remove_one,
};

void pars_str_to_devs_arr(void)
{
    char tmp_str[MAX_NUM_OF_DEVS * MAX_DEV_NAME_LEN];
    char *end, *tok;
    strncpy(tmp_str, pci_device, MAX_NUM_OF_DEVS * MAX_DEV_NAME_LEN);
    end = tmp_str;
    while ((tok = strsep(&end, ",")) != NULL)
    {
        strncpy(g_reset_info.pdev_arr[g_reset_info.num_of_pci_devs++], tok, strlen(tok));
    }
}

static int __init init(void)
{
    pars_str_to_devs_arr();
    return pci_register_driver(&mst_ppc_pci_reset_driver);
}

static void __exit cleanup(void)
{
    pci_unregister_driver(&mst_ppc_pci_reset_driver);
}

module_init(init);
module_exit(cleanup);
