// Mostly copied from https://github.com/rivosinc/hammer

#include <algorithm>
#include <cerrno>
#include <cstdint>
#include <cstdlib>
#include <fesvr/memif.h>
#include <gdbstub.h>
#include <memory>
#include <optional>
#include <riscv/cfg.h>
#include <riscv/debug_module.h>
#include <riscv/devices.h>
#include <riscv/mmu.h>
#include <riscv/processor.h>
#include <riscv/sim.h>
#include <riscv/trap.h>
#include <spdlog/cfg/env.h>
#include <spdlog/logger.h>
#include <spdlog/sinks/stdout_color_sinks.h>
#include <spdlog/spdlog.h>
#include <utility>

class Difftest {
public:
  std::shared_ptr<spdlog::logger> logger =
      spdlog::stdout_color_mt("spike-diff");

  Difftest();

  template <typename T>
  std::optional<std::vector<T>>
  get_mem(addr_t virtual_address, size_t bytes_to_read, uint8_t hart_id = 0);

  template <typename T>
  int set_mem(addr_t virtual_address, const std::vector<T> &data,
              uint8_t hart_id = 0);

  uint32_t set_breakpoint(addr_t addr, uint8_t hart_id = 0);

  reg_t get_gpr(uint8_t gpr_id, uint8_t hart_id = 0);
  void set_gpr(uint8_t gpr_id, reg_t new_gpr_value, uint8_t hart_id = 0);

  reg_t get_pc(uint8_t hart_id = 0);
  void set_pc(reg_t new_pc_value, uint8_t hart_id = 0);

  void single_step(uint8_t hart_id = 0);
  std::optional<bp_type_t> cont(void);
  void get_state(uint8_t hart_id = 0);

private:
  sim_t *simulator;
};

Difftest::Difftest() {
  debug_module_config_t dmconfig = {.progbufsize = 2,
                                    .max_sba_data_width = 0,
                                    .require_authentication = false,
                                    .abstract_rti = 0,
                                    .support_hasel = true,
                                    .support_abstract_csr_access = true,
                                    .support_abstract_fpr_access = true,
                                    .support_haltgroups = true,
                                    .support_impebreak = true};
  std::vector<std::pair<reg_t, abstract_mem_t *>> mem(
      1, std::make_pair(reg_t(DRAM_BASE), new mem_t(0x8000000)));

  std::vector<device_factory_sargs_t> plugin_devices_factories;

  cfg_t cfg;
  cfg.initrd_bounds = std::make_pair((reg_t)0, (reg_t)0),
  cfg.bootargs = nullptr;
  cfg.isa = "RV32I";
  cfg.priv = "M";
  cfg.misaligned = false;
  cfg.endianness = endianness_little;
  cfg.pmpregions = 0;
  cfg.pmpgranularity = 4;
  cfg.mem_layout = std::vector<mem_cfg_t>();
  cfg.hartids = std::vector<size_t>(1, 0);
  cfg.explicit_hartids = false;
  cfg.real_time_clint = false;
  cfg.trigger_count = 4;

  simulator = new sim_t{&cfg,
                        false,
                        mem,
                        plugin_devices_factories,
                        std::vector<std::string>(1, ""),
                        dmconfig,
                        nullptr,
                        false,
                        nullptr,
                        false,
                        nullptr};
  simulator->set_debug(false);
  processor_t *hart = simulator->get_core(0);
  hart->reset();
  // spdlog::trace("Sim: {} {}", simulator->get_core(0)))
}

template <typename T>
std::optional<std::vector<T>> Difftest::get_mem(addr_t virtual_address,
                                                size_t bytes_to_read,
                                                uint8_t hart_id) {
  mmu_t *mmu = simulator->get_core(hart_id)->get_mmu();
  std::vector<T> data;
  try {
    for (size_t i = 0; i < bytes_to_read; i += sizeof(T))
      data.push_back(mmu->load<T>(virtual_address + i));
  } catch (trap_t &t) {
    return std::nullopt;
  }
  return data;
}

template <typename T>
int Difftest::set_mem(addr_t virtual_address, const std::vector<T> &data,
                      uint8_t hart_id) {
  mmu_t *mmu = simulator->get_core(hart_id)->get_mmu();

  logger->trace("(W) {:x} data size: {:x}", virtual_address, data.size());
  try {
    for (auto &i : data) {
      mmu->store<T>(virtual_address, i);
      virtual_address += sizeof(T);
    }
  } catch (trap_t &t) {
    return EACCES;
  }
  return 0;
}

uint32_t Difftest::set_breakpoint(addr_t addr, uint8_t hart_id) {
  uint32_t original_inst = get_mem<uint32_t>(addr, 4, hart_id).value().at(0);
  const std::vector<uint32_t> EBREAK{0x100073U};
  // assert(set_mem(addr, EBREAK, hart_id) == 0);
  return original_inst;
}

reg_t Difftest::get_gpr(uint8_t gpr_id, uint8_t hart_id) {
  assert(gpr_id < NXPR);

  processor_t *hart = simulator->get_core(hart_id);
  state_t *hart_state = hart->get_state();

  return hart_state->XPR[gpr_id];
}

void Difftest::set_gpr(uint8_t gpr_id, reg_t new_gpr_value, uint8_t hart_id) {
  assert(gpr_id < NXPR);

  processor_t *hart = simulator->get_core(hart_id);
  state_t *hart_state = hart->get_state();

  hart_state->XPR.write(gpr_id, new_gpr_value);
}

reg_t Difftest::get_pc(uint8_t hart_id) {
  processor_t *hart = simulator->get_core(hart_id);
  state_t *hart_state = hart->get_state();

  return hart_state->pc;
}

void Difftest::set_pc(reg_t new_pc_value, uint8_t hart_id) {
  processor_t *hart = simulator->get_core(hart_id);
  state_t *hart_state = hart->get_state();

  hart_state->pc = new_pc_value;
}

void Difftest::single_step(uint8_t hart_id) {
  processor_t *hart = simulator->get_core(hart_id);
  hart->step(1);
}

std::optional<bp_type_t> Difftest::cont(void) { exit(1); }

void Difftest::get_state(uint8_t hart_id) {
  processor_t *hart = simulator->get_core(hart_id);
  hart->halted();
}

extern "C" {
typedef uint32_t word_t;

struct DbgState {
  Difftest *diff;
  struct Breakpoint {
    addr_t addr;
    uint32_t original_instruction;

    bool operator==(const Breakpoint &b) { return this->addr == b.addr; }
  };
  std::vector<Breakpoint> breakpoints;
};

arch_info_t spike_isa_arch_info = {
    .target_desc = TARGET_RV32,
    .reg_num = 32,
    .reg_byte = 4,
};
bool spike_do_difftest = true;
bool spike_dbg_state_size = sizeof(DbgState);

int spike_init(void *args) {
  DbgState *dbg = (DbgState *)args;
  dbg->diff = new Difftest{};
  return dbg->diff != nullptr;
}

int spike_read_mem(void *args, size_t addr, size_t len, void *val) {
  DbgState *dbg = (DbgState *)args;

  try {
    std::vector<word_t> mem{dbg->diff->get_mem<word_t>(addr, len).value()};
    memcpy(val, mem.data(), len);
  } catch (std::bad_optional_access) {
    return EACCES;
  }

  return 0;
}

int spike_write_mem(void *args, size_t addr, size_t len, void *val) {
  DbgState *dbg = (DbgState *)args;
  size_t dlen = len / sizeof(word_t);
  word_t *dataptr = (word_t *)val;
  std::vector<word_t> data{dataptr, dataptr + dlen};

  word_t remains = 0;
  memcpy(&remains, dataptr + dlen, len % sizeof(word_t));
  if (len % sizeof(word_t))
    data.push_back(remains);

  return dbg->diff->set_mem(addr, data);
}

int spike_read_reg(void *args, int regno, size_t *value) {
  DbgState *dbg = (DbgState *)args;
  if (regno == 32) {
    *value = (word_t)dbg->diff->get_pc();
  } else {
    *value = (word_t)dbg->diff->get_gpr(regno);
  }
  return 0;
}

int spike_write_reg(void *args, int regno, size_t value) {
  DbgState *dbg = (DbgState *)args;
  if (regno == 32) {
    dbg->diff->set_pc(value);
  } else {
    dbg->diff->set_gpr(regno, value);
  }
  return 0;
}

void spike_stepi(void *args, gdb_action_t *res) {
  DbgState *dbg = (DbgState *)args;
  static bool on_breakpoint = false;
  if (!on_breakpoint) {
    for (const auto &bp : dbg->breakpoints) {
      if ((word_t)dbg->diff->get_pc() == bp.addr) {
        res->reason = gdb_action_t::ACT_BREAKPOINT;
        on_breakpoint = true;
        return;
      } else if ((word_t)dbg->diff->get_pc() == 0) {
        res->reason = gdb_action_t::ACT_SHUTDOWN;
        return;
      }
    }
  }
  dbg->diff->single_step();
  res->reason = gdb_action_t::ACT_NONE;
  on_breakpoint = false;
}

__attribute__((visibility("default"))) void spike_cont(void *args,
                                                       gdb_action_t *res) {
  ;
}

bool spike_set_bp(void *args, size_t addr, bp_type_t bp) {
  DbgState *dbg = (DbgState *)args;
  dbg->breakpoints.emplace_back(
      DbgState::Breakpoint{addr, dbg->diff->set_breakpoint(addr)});
  return true;
}
__attribute__((visibility("default"))) void
spike_del_bp(void *args, size_t addr, bp_type_t type) {
  DbgState *dbg = (DbgState *)args;
  auto &bps = dbg->breakpoints;
  auto it = std::find(bps.begin(), bps.end(), DbgState::Breakpoint{addr, type});
  std::swap(*it, *bps.rbegin());
  bps.pop_back();
}

__attribute__((visibility("default"))) void spike_on_interrupt() { ; }

} // extern "C"