/***************************************************************************************
* Copyright (c) 2025 Beijing Institute of Open Source Chip (BOSC)
* Copyright (c) 2020-2025 Institute of Computing Technology, Chinese Academy of Sciences
*
* DiffTest is licensed under Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
*          http://license.coscl.org.cn/MulanPSL2
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
*
* See the Mulan PSL v2 for more details.
***************************************************************************************/
#include "xdma.h"
#include "difftest-dpic.h"
#include "mpool.h"
#include "ram.h"
#include <algorithm>
#include <cstring>
#include <execinfo.h>
#include <fcntl.h>
#include <fstream>
#include <inttypes.h>
#include <iostream>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <sys/mman.h>
#include <unistd.h>

#define XDMA_USER       "/dev/xdma0_user"
#define XDMA_BYPASS     "/dev/xdma0_bypass"
#define XDMA_C2H_DEVICE "/dev/xdma0_c2h_"
#define XDMA_H2C_DEVICE "/dev/xdma0_h2c_0"
static const size_t H2C_AXIS_BYTES = 64;

void signal_handler(int sig) {
  void *array[20];
  size_t size;
  size = backtrace(array, 20);

  fprintf(stderr, "Error: signal %d:\n", sig);
  backtrace_symbols_fd(array, size, STDERR_FILENO);
  exit(1);
}

template <typename Func, typename Obj, typename... Args> void thread_wrapper(Func func, Obj obj, Args... args) {
  signal(SIGSEGV, signal_handler);
  (obj->*func)(args...);
}

void handle_sigint(int sig) {
  printf("handle sigint unlink pcie success, exit fpga-host!\n");
  exit(1);
}

FpgaXdma::FpgaXdma()
#ifdef USE_THREAD_MEMPOOL
    : xdma_mempool(sizeof(FpgaPackgeHead))
#endif // USE_THREAD_MEMPOOL
{
  signal(SIGINT, handle_sigint);
  for (int i = 0; i < CONFIG_DMA_CHANNELS; i++) {
    char c2h_device[64];
    sprintf(c2h_device, "%s%d", XDMA_C2H_DEVICE, i);
#ifdef FPGA_SIM
    xdma_sim_open(i, true);
#else
    xdma_c2h_fd[i] = open(c2h_device, O_RDONLY);
    if (xdma_c2h_fd[i] == -1) {
      std::cout << c2h_device << std::endl;
      perror("Failed to open XDMA device");
      exit(-1);
    }
    std::cout << "XDMA link " << c2h_device << std::endl;
#endif // FPGA_SIM
  }
#ifdef FPGA_SIM
  xdma_sim_axilite_open(true);
  xdma_sim_workload_open(true);
  xdma_sim_h2c_open(0, true);
#endif // FPGA_SIM
#if defined(CONFIG_USE_XDMA_H2C) && !defined(FPGA_SIM)
  xdma_h2c_fd = open(XDMA_H2C_DEVICE, O_WRONLY | O_TRUNC);
  if (xdma_h2c_fd == -1) {
    std::cout << XDMA_H2C_DEVICE << std::endl;
    perror("Failed to open XDMA device");
    exit(-1);
  }
  std::cout << "XDMA link " << XDMA_H2C_DEVICE << std::endl;
#endif
}

FpgaXdma::~FpgaXdma() {
#ifdef FPGA_SIM
  for (int i = 0; i < CONFIG_DMA_CHANNELS; i++) {
    xdma_sim_close(i);
  }
  xdma_sim_workload_close(true);
  xdma_sim_h2c_close(0);
  xdma_sim_axilite_close(true);
#endif // FPGA_SIM
}

void FpgaXdma::wait_fpga_io_done(uint64_t address, const char *tag) {
  const int max_retry = 600000; // 10 minute
  for (int retry = 0; retry < max_retry; retry++) {
    uint32_t status = fpga_io_read(address) & 0x3;
    if (status == 0x2) {
      return;
    }
    if (status == 0x3) {
      fprintf(stderr, "[fpga-host] %s failed: address range exceeds FPGA AXI address width\n", tag);
      exit(1);
    }
    usleep(1000);
  }
  fprintf(stderr, "[fpga-host] timeout waiting for %s\n", tag);
  exit(1);
}

#ifdef CONFIG_USE_XDMA_H2C
void FpgaXdma::h2c_load_workload(const char *workload, uint64_t ram_size) {
  MmapMemory memory(workload, ram_size, false, 0);
  uint64_t img_size = memory.get_img_size();
  const uint8_t *image = reinterpret_cast<const uint8_t *>(memory.as_ptr());
  if (image == nullptr) {
    fprintf(stderr, "[fpga-host] H2C load requires mmap-backed memory image\n");
    exit(-1);
  }

#ifdef FPGA_SIM
  uint64_t offset = 0;
  while (offset < img_size) {
    size_t beatBytes = std::min<uint64_t>(H2C_AXIS_BYTES, img_size - offset);
    char beat[H2C_AXIS_BYTES] = {};
    memcpy(beat, image + offset, beatBytes);
    uint64_t tkeep = beatBytes == H2C_AXIS_BYTES ? UINT64_MAX : ((1ULL << beatBytes) - 1);
    if (xdma_sim_h2c_write(0, beat, tkeep, offset + beatBytes >= img_size, sizeof(beat)) != (int)sizeof(beat)) {
      fprintf(stderr, "[fpga-host] FPGA_SIM H2C shared-memory write failed\n");
      exit(-1);
    }
    offset += beatBytes;
  }
  printf("[fpga-host] FPGA_SIM H2C queued %" PRIu64 " bytes\n", img_size);
#else
  const char *payload = reinterpret_cast<const char *>(image);
  ssize_t written = write(xdma_h2c_fd, payload, img_size);
  if (written < 0) {
    perror("[fpga-host] XDMA H2C write failed");
    exit(-1);
  }
  if ((size_t)written != img_size) {
    fprintf(stderr, "[fpga-host] XDMA H2C partial write: %zd/%" PRIu64 " bytes\n", written, img_size);
    exit(-1);
  }
  printf("[fpga-host] XDMA H2C queued %" PRIu64 " bytes\n", img_size);
#endif // FPGA_SIM
}
#endif // CONFIG_USE_XDMA_H2C

// write xdma_bypass memory or xdma_user
void FpgaXdma::device_write(bool is_bypass, const char *workload, uint64_t addr, uint64_t value) {
  (void)workload;
#ifdef FPGA_SIM
  if (is_bypass) {
    fprintf(stderr, "[fpga-host] FPGA_SIM XDMA bypass write is unsupported\n");
    exit(-1);
  }
  if (xdma_sim_axilite_write(static_cast<uint32_t>(addr), static_cast<uint32_t>(value), 0xf) != 0) {
    fprintf(stderr, "[fpga-host] FPGA_SIM AXI-Lite command queue is full, addr=0x%lx value=0x%lx\n", addr, value);
    exit(-1);
  }
  return;
#endif // FPGA_SIM

  uint64_t pg_size = sysconf(_SC_PAGE_SIZE);
  uint64_t size = !is_bypass ? 0x1000 : 0x100000;
  uint64_t aligned_size = (size + 0xffful) & ~0xffful;
  uint64_t base = addr & ~0xffful;
  uint32_t offset = addr & 0xfffu;
  int fd = -1;

  if (base % pg_size != 0) {
    printf("base must be a multiple of system page size\n");
    exit(-1);
  }

  if (is_bypass)
    fd = open(XDMA_BYPASS, O_RDWR | O_SYNC);
  else
    fd = open(XDMA_USER, O_RDWR | O_SYNC);
  if (fd < 0) {
    printf("Failed to open %s\n", is_bypass ? XDMA_BYPASS : XDMA_USER);
    exit(-1);
  }

  void *m_ptr = mmap(nullptr, aligned_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, base);
  if (m_ptr == MAP_FAILED) {
    close(fd);
    printf("failed to mmap\n");
    exit(-1);
  }

  if (is_bypass) {
    if (simMemory->get_img_size() > aligned_size) {
      printf("The loaded workload size exceeds the xdma bypass size");
      exit(-1);
    }
    memcpy(static_cast<char *>(m_ptr) + offset, static_cast<const void *>(simMemory->as_ptr()),
           simMemory->get_img_size());
  } else {
    ((volatile uint32_t *)m_ptr)[offset >> 2] = value;
  }

  munmap(m_ptr, aligned_size);
  close(fd);
}

uint32_t FpgaXdma::device_read(bool is_bypass, uint64_t addr) {
#ifdef FPGA_SIM
  if (is_bypass) {
    fprintf(stderr, "[fpga-host] FPGA_SIM XDMA bypass read is unsupported\n");
    exit(-1);
  }
  uint32_t data = 0;
  if (xdma_sim_axilite_read(static_cast<uint32_t>(addr), &data) != 0) {
    fprintf(stderr, "[fpga-host] FPGA_SIM AXI-Lite read failed, addr=0x%lx\n", addr);
    exit(-1);
  }
  return data;
#endif // FPGA_SIM

  uint64_t pg_size = sysconf(_SC_PAGE_SIZE);
  uint64_t size = !is_bypass ? 0x1000 : 0x100000;
  uint64_t aligned_size = (size + 0xffful) & ~0xffful;
  uint64_t base = addr & ~0xffful;
  uint32_t offset = addr & 0xfffu;

  if (base % pg_size != 0) {
    printf("base must be a multiple of system page size\n");
    exit(-1);
  }

  int fd = open(is_bypass ? XDMA_BYPASS : XDMA_USER, O_RDWR | O_SYNC);
  if (fd < 0) {
    printf("Failed to open %s\n", is_bypass ? XDMA_BYPASS : XDMA_USER);
    exit(-1);
  }

  void *m_ptr = mmap(nullptr, aligned_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, base);
  if (m_ptr == MAP_FAILED) {
    close(fd);
    printf("failed to mmap\n");
    exit(-1);
  }

  uint32_t value = ((volatile uint32_t *)m_ptr)[offset >> 2];
  munmap(m_ptr, aligned_size);
  close(fd);
  return value;
}

#ifdef USE_THREAD_MEMPOOL
void FpgaXdma::start_transmit_thread() {
  for (int i = 0; i < CONFIG_DMA_CHANNELS; i++) {
    printf("start channel %d \n", i);
    receive_thread[i] = std::thread(thread_wrapper<decltype(&FpgaXdma::read_xdma_thread), FpgaXdma *, int>,
                                    &FpgaXdma::read_xdma_thread, this, i);
  }
  process_thread = std::thread(thread_wrapper<decltype(&FpgaXdma::write_difftest_thread), FpgaXdma *>,
                               &FpgaXdma::write_difftest_thread, this);
}

void FpgaXdma::stop_thansmit_thread() {
  for (int i = 0; i < CONFIG_DMA_CHANNELS; i++) {
    if (receive_thread[i].joinable())
      receive_thread[i].join();
#ifdef FPGA_SIM
    xdma_sim_close(i);
#else
    close(xdma_c2h_fd[i]);
#endif // FPGA_SIM
  }

  if (process_thread.joinable())
    process_thread.join();
#if defined(CONFIG_USE_XDMA_H2C) && !defined(FPGA_SIM)
  close(xdma_h2c_fd);
#endif
}

void FpgaXdma::read_xdma_thread(int channel) {
  size_t mem_get_idx = 0;
  while (running) {
    char *mem = xdma_mempool.get_free_chunk(&mem_get_idx);
#ifdef FPGA_SIM
    size_t size = xdma_sim_read(channel, mem, sizeof(FpgaPackgeHead));
#else
    size_t size = read(xdma_c2h_fd[channel], mem, sizeof(FpgaPackgeHead));
#endif // FPGA_SIM
    if (xdma_mempool.write_free_chunk(mem[0], mem_get_idx) == false) {
      printf("It should not be the case that no available block can be found\n");
      assert(0);
    }
  }
}

void FpgaXdma::write_difftest_thread() {
  FpgaPackgeHead *packge;
  uint8_t recv_count = 0;
  xdma_mempool.wait_mempool_start();
  while (running) {
    packge = reinterpret_cast<FpgaPackgeHead *>(xdma_mempool.read_busy_chunk());
    if (packge == nullptr) {
      printf("Failed to read data from the XDMA memory pool\n");
      assert(0);
    }
    if (packge->diff_packge[0].packge_idx != recv_count) {
      printf("read mempool idx failed, packge_idx %d need_idx %d\n", packge->diff_packge[0].packge_idx, recv_count);
      assert(0);
    }
    recv_count++;
    // packge unpack
    for (size_t i = 0; i < DMA_PACKGE_NUM; i++) {
      v_difftest_Batch(packge->diff_packge[i].diff_packge);
    }
    xdma_mempool.set_free_chunk();
  }
}

#else // !USE_THREAD_MEMPOOL

void *posix_memalignd_malloc(size_t size) {
  void *ptr = nullptr;
  int ret = posix_memalign(&ptr, 4096, size);
  if (ret != 0) {
    perror("posix_memalign failed");
    return nullptr;
  }
  return ptr;
}
void FpgaXdma::read_and_process() {
  printf("start channel 0\n");
  FpgaPackgeHead *packge = (FpgaPackgeHead *)posix_memalignd_malloc(sizeof(FpgaPackgeHead));
  memset(packge, 0, sizeof(FpgaPackgeHead));
  while (running) {
#ifdef FPGA_SIM
    size_t size = xdma_sim_read(0, (char *)packge, sizeof(FpgaPackgeHead));
#else
    size_t size = read(xdma_c2h_fd[0], packge, sizeof(FpgaPackgeHead));
#endif // FPGA_SIM
    for (size_t i = 0; i < DMA_PACKGE_NUM; i++) {
      v_difftest_Batch(packge->diff_packge[i].diff_packge);
    }
  }
  free(packge);
}
#endif // USE_THREAD_MEMPOOL
