// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2026, Advanced Micro Devices, Inc.
 */

#include <drm/amdxdna_accel.h>
#include <drm/drm_print.h>
#include <linux/pci.h>

#include "aie.h"
#include "aie4_msg_priv.h"
#include "aie4_pci.h"
#include "amdxdna_mailbox.h"
#include "amdxdna_mailbox_helper.h"
#include "amdxdna_pci_drv.h"

static int aie4_destroy_vfs(struct amdxdna_dev_hdl *ndev)
{
        DECLARE_AIE_MSG(aie4_msg_destroy_vfs, AIE4_MSG_OP_DESTROY_VFS);
        int ret;

        ret = aie_send_mgmt_msg_wait(&ndev->aie, &msg);
        if (ret)
                XDNA_ERR(ndev->aie.xdna, "destroy vfs op failed: %d", ret);

        return ret;
}

static int aie4_create_vfs(struct amdxdna_dev_hdl *ndev, int num_vfs)
{
        DECLARE_AIE_MSG(aie4_msg_create_vfs, AIE4_MSG_OP_CREATE_VFS);
        int ret;

        req.vf_cnt = num_vfs;
        ret = aie_send_mgmt_msg_wait(&ndev->aie, &msg);
        if (ret)
                XDNA_ERR(ndev->aie.xdna, "create vfs op failed: %d", ret);

        return ret;
}

int aie4_sriov_stop(struct amdxdna_dev_hdl *ndev)
{
        struct amdxdna_dev *xdna = ndev->aie.xdna;
        struct pci_dev *pdev = to_pci_dev(xdna->ddev.dev);
        int ret;

        if (!pci_num_vf(pdev))
                return 0;

        ret = pci_vfs_assigned(pdev);
        if (ret) {
                XDNA_ERR(xdna, "VFs are still assigned to VMs");
                return -EPERM;
        }

        pci_disable_sriov(pdev);
        return aie4_destroy_vfs(ndev);
}

static int aie4_sriov_start(struct amdxdna_dev_hdl *ndev, int num_vfs)
{
        struct amdxdna_dev *xdna = ndev->aie.xdna;
        struct pci_dev *pdev = to_pci_dev(xdna->ddev.dev);
        int ret;

        ret = aie4_create_vfs(ndev, num_vfs);
        if (ret)
                return ret;

        ret = pci_enable_sriov(pdev, num_vfs);
        if (ret) {
                XDNA_ERR(xdna, "configure VFs failed, ret: %d", ret);
                aie4_destroy_vfs(ndev);
                return ret;
        }

        return num_vfs;
}

int aie4_sriov_configure(struct amdxdna_dev *xdna, int num_vfs)
{
        struct amdxdna_dev_hdl *ndev = xdna->dev_handle;

        drm_WARN_ON(&xdna->ddev, !mutex_is_locked(&xdna->dev_lock));

        return (num_vfs) ? aie4_sriov_start(ndev, num_vfs) : aie4_sriov_stop(ndev);
}