#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");

char pci_device[128];
module_param_string(pci_dev, pci_device, sizeof(pci_device), 0444);

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

	printk(KERN_INFO "%s %s %d Going to reset device: %s\n",
	       dev_driver_string(&pdev->dev), __func__, __LINE__, dev_name(&pdev->dev));

	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;
        }

	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));

	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;
        }

	pci_restore_state(pdev);

	printk(KERN_INFO "%s %s %d Reset succeeded for device: %s\n",
	       dev_driver_string(&pdev->dev), __func__, __LINE__, dev_name(&pdev->dev));

	return 0;
}

static int init_one(struct pci_dev *pdev,
		    const struct pci_device_id *id)
{
	if (!strcmp(pci_device, dev_name(&pdev->dev)))
		return reset_device(pdev);
	else
		return 0;

}

static void remove_one(struct pci_dev *pdev)
{
	if (!strcmp(pci_device, dev_name(&pdev->dev))) {
		pci_clear_master(pdev);
		pci_disable_device(pdev);
	}
}

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 */
	{ 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,
};

static int __init init(void)
{
	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);
