// SPDX-License-Identifier: Apache-2.0
// Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.

#include "dev_info.h"

// Test program location, all workspace paths below are relative to it
extern std::string cur_path;
// Force to use xclbin pointed to by this path
extern std::string xclbin_path;

namespace {

xclbin_info xclbin_infos[] = {
  {
    .name = "1x4.xclbin",
    .device = npu1_device_id,
    .revision_id = npu1_revision_id,
    .ip_name2idx = {
      { "DPU_PDI_0:IPUV1CNN",         {0} },
      { "DPU_PDI_1:IPUV1CNN",         {1} },
      { "DPU_PDI_2:IPUV1CNN",         {2} },
      { "DPU_PDI_3:IPUV1CNN",         {3} },
      { "DPU_PDI_4:IPUV1CNN",         {4} },
      { "DPU_PDI_5:IPUV1CNN",         {5} },
      { "DPU_PDI_6:IPUV1CNN",         {6} },
      { "DPU_PDI_7:IPUV1CNN",         {7} },
    },
    .workspace = "npu1_workspace",
    .data = "data",
  },
  {
    .name = "1x4.xclbin",
    .device = npu1_device_id1,
    .revision_id = npu1_revision_id1,
    .ip_name2idx = {
      { "DPU_PDI_0:IPUV1CNN",         {0} },
      { "DPU_PDI_1:IPUV1CNN",         {1} },
      { "DPU_PDI_2:IPUV1CNN",         {2} },
      { "DPU_PDI_3:IPUV1CNN",         {3} },
      { "DPU_PDI_4:IPUV1CNN",         {4} },
      { "DPU_PDI_5:IPUV1CNN",         {5} },
      { "DPU_PDI_6:IPUV1CNN",         {6} },
      { "DPU_PDI_7:IPUV1CNN",         {7} },
    },
    .workspace = "npu1_workspace",
    .data = "data",
  },
  {
    .name = "1x4.xclbin",
    .device = npu4_device_id,
    .revision_id = npu2_revision_id,
    .ip_name2idx = {
      { "DPU_PDI_0:IPUV1CNN",         {0} },
      { "DPU_PDI_1:IPUV1CNN",         {1} },
      { "DPU_PDI_2:IPUV1CNN",         {2} },
      { "DPU_PDI_3:IPUV1CNN",         {3} },
      { "DPU_PDI_4:IPUV1CNN",         {4} },
      { "DPU_PDI_5:IPUV1CNN",         {5} },
      { "DPU_PDI_6:IPUV1CNN",         {6} },
      { "DPU_PDI_7:IPUV1CNN",         {7} },
    },
    .workspace = "npu2_workspace",
    .data = "data",
  },
  {
    .name = "vadd.elf",
    .device = npu3_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:vadd", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "vadd.elf",
    .device = npu3_device_id1,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:vadd", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "move_memtiles.elf",
    .device = npu3_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:move_memtiles", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "move_memtiles.elf",
    .device = npu3_device_id1,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:move_memtiles", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "ddr_memtile.elf",
    .device = npu3_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:move_ddr_memtile", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "ddr_memtile.elf",
    .device = npu3_device_id1,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:move_ddr_memtile", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "remote_barrier.elf",
    .device = npu3_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:remote_barrier", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "remote_barrier.elf",
    .device = npu3_device_id1,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:remote_barrier", {0xffffffff} },
    },
    .workspace = "npu3_workspace",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "1x4.xclbin",
    .device = npu4_device_id,
    .revision_id = npu4_revision_id,
    .ip_name2idx = {
      { "DPU_PDI_0:IPUV1CNN",         {0} },
      { "DPU_PDI_1:IPUV1CNN",         {1} },
      { "DPU_PDI_2:IPUV1CNN",         {2} },
      { "DPU_PDI_3:IPUV1CNN",         {3} },
      { "DPU_PDI_4:IPUV1CNN",         {4} },
      { "DPU_PDI_5:IPUV1CNN",         {5} },
      { "DPU_PDI_6:IPUV1CNN",         {6} },
      { "DPU_PDI_7:IPUV1CNN",         {7} },
      { "DPU_PDI_8:IPUV1CNN",         {8} },
      { "DPU_PDI_9:IPUV1CNN",         {9} },
      { "DPU_PDI_10:IPUV1CNN",         {10} },
      { "DPU_PDI_11:IPUV1CNN",         {11} },
      { "DPU_PDI_12:IPUV1CNN",         {12} },
      { "DPU_PDI_13:IPUV1CNN",         {13} },
    },
    .workspace = "npu4_workspace",
    .data = "data",
  },
  {
    .name = "1x4.xclbin",
    .device = npu4_device_id,
    .revision_id = npu5_revision_id,
    .ip_name2idx = {
      { "DPU_PDI_0:IPUV1CNN",         {0} },
      { "DPU_PDI_1:IPUV1CNN",         {1} },
      { "DPU_PDI_2:IPUV1CNN",         {2} },
      { "DPU_PDI_3:IPUV1CNN",         {3} },
      { "DPU_PDI_4:IPUV1CNN",         {4} },
      { "DPU_PDI_5:IPUV1CNN",         {5} },
      { "DPU_PDI_6:IPUV1CNN",         {6} },
      { "DPU_PDI_7:IPUV1CNN",         {7} },
      { "DPU_PDI_8:IPUV1CNN",         {8} },
      { "DPU_PDI_9:IPUV1CNN",         {9} },
      { "DPU_PDI_10:IPUV1CNN",        {10} },
      { "DPU_PDI_11:IPUV1CNN",        {11} },
      { "DPU_PDI_12:IPUV1CNN",        {12} },
      { "DPU_PDI_13:IPUV1CNN",        {13} },
    },
    .workspace = "npu5_workspace",
    .data = "data",
  },
  {
    .name = "1x4.xclbin",
    .device = npu4_device_id,
    .revision_id = npu6_revision_id,
    .ip_name2idx = {
      { "DPU_PDI_0:IPUV1CNN",         {0} },
      { "DPU_PDI_1:IPUV1CNN",         {1} },
      { "DPU_PDI_2:IPUV1CNN",         {2} },
      { "DPU_PDI_3:IPUV1CNN",         {3} },
      { "DPU_PDI_4:IPUV1CNN",         {4} },
      { "DPU_PDI_5:IPUV1CNN",         {5} },
      { "DPU_PDI_6:IPUV1CNN",         {6} },
      { "DPU_PDI_7:IPUV1CNN",         {7} },
      { "DPU_PDI_8:IPUV1CNN",         {8} },
      { "DPU_PDI_9:IPUV1CNN",         {9} },
      { "DPU_PDI_10:IPUV1CNN",         {10} },
      { "DPU_PDI_11:IPUV1CNN",         {11} },
      { "DPU_PDI_12:IPUV1CNN",         {12} },
      { "DPU_PDI_13:IPUV1CNN",         {13} },
    },
    .workspace = "npu6_workspace",
    .data = "data",
  },
  {
    .name = "design.xclbin",
    .device = npu1_device_id,
    .revision_id = npu1_revision_id,
    .ip_name2idx = {
      { "DPU_ELF:IPUV1CNN", {9} },
    },
    .workspace = "local_shim_test_data/elf_txn_no_cp_npu1",
    .data = "",
    .type = KERNEL_TYPE_TXN,
  },
  {
    .name = "design.xclbin",
    .device = npu4_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:IPUV1CNN", {1} },
    },
    .workspace = "local_shim_test_data/elf_txn_no_cp_npu4",
    .data = "",
    .type = KERNEL_TYPE_TXN,
  },
  {
    .name = "pm_reload.xclbin",
    .device = npu4_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:IPUV1CNN", {1} },
    },
    .workspace = "local_shim_test_data/elf_txn_preempt_no_cp_npu4",
    .data = "",
    .type = KERNEL_TYPE_TXN_PREEMPT,
  },
  {
    .name = "yolo_fullelf_aximm.elf",
    .device = npu4_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:dpu", {0xffffffff} }, // CU index is not needed for full ELF
    },
    .workspace = "local_shim_test_data/xclbin2elf_preempt_tinyyolo_npu4",
    .data = "",
    .type = KERNEL_TYPE_TXN_FULL_ELF_PREEMPT,
  },
  {
    .name = "bad_txn.xclbin",
    .device = npu4_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:IPUV1CNN", {0} },
    },
    .workspace = "local_shim_test_data/elf_no_cp_bad_txn_npu4",
    .data = "",
    .type = KERNEL_TYPE_TXN,
  },
  {
    .name = "gemm.xclbin",
    .device = npu4_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:IPUV1CNN", {0} },
    },
    .workspace = "local_shim_test_data/elf_no_cp_gemm_npu4",
    .data = "",
    .type = KERNEL_TYPE_TXN,
  },
  {
    .name = "nop.xclbin",
    .device = npu4_device_id,
    .revision_id = npu_any_revision_id,
    .ip_name2idx = {
      { "DPU:IPUV1CNN", {1} },
    },
    .workspace = "local_shim_test_data/elf_no_op_npu4",
    .data = "",
    .type = KERNEL_TYPE_TXN,
  },
  {
    .name = "nop.xclbin",
    .device = npu1_device_id,
    .revision_id = npu1_revision_id,
    .ip_name2idx = {
      { "DPU:IPUV1CNN", {1} },
    },
    .workspace = "local_shim_test_data/elf_no_op_npu1",
    .data = "",
    .type = KERNEL_TYPE_TXN,
  },

};

}

const xclbin_info&
get_xclbin_info(device* dev, const char *xclbin_name)
{
  auto pci_dev_id = device_query<query::pcie_device>(dev);
  auto revision_id = device_query<query::pcie_id>(dev).revision_id;
  for (auto& xclbin : xclbin_infos) {
    if ((xclbin.device == pci_dev_id) &&
        ((xclbin.revision_id == revision_id) || (xclbin.revision_id == npu_any_revision_id)) &&
        (xclbin_name == nullptr || !strcmp(xclbin.name, xclbin_name)))
      return xclbin;
  }
  throw std::runtime_error("xclbin info not found");
}

std::string
get_xclbin_name(device* dev)
{
  return get_xclbin_info(dev).name;
}

std::string
get_kernel_name(device* dev, const char *xclbin)
{
  return get_xclbin_info(dev, xclbin).ip_name2idx.begin()->first;
}

kernel_type
get_kernel_type(device* dev, const char *xclbin)
{
  return get_xclbin_info(dev, xclbin).type;
}

static std::string
get_xclbin_workspace(device* dev, const char *xclbin_name)
{
  return cur_path + "/../" + get_xclbin_info(dev, xclbin_name).workspace + "/";
}

std::string
get_xclbin_data(device* dev, const char *xclbin_name)
{
  auto wrk = get_xclbin_workspace(dev, xclbin_name);
  return wrk + get_xclbin_info(dev, xclbin_name).data + "/";
}

std::string
get_xclbin_path(device* dev, const char *xclbin_name)
{
  if (!xclbin_path.empty())
    return xclbin_path;

  auto wrk = get_xclbin_workspace(dev, xclbin_name);
  if (!xclbin_name)
    return wrk + get_xclbin_name(dev);
  return wrk + std::string(xclbin_name);
}

const std::map<const char*, cuidx_type>&
get_xclbin_ip_name2index(device* dev, const char *xclbin_name)
{
  return get_xclbin_info(dev, xclbin_name).ip_name2idx;
}

std::string
get_preemption_libs_path()
{
  return cur_path + "/../" + "preemption_libs" + "/";
}
