// Common Verilator Simulation/test constructs.
// Includes a simulation-fixture useful for writing tests.
#pragma once
#include "verilated.h"
#include "verilated_vcd_c.h"
#include <cstdint>
#include <memory>
#include <queue>
#include <span>
#include <vector>

// represents a generic co-simulated device.
// While this abstract class is very simple, it enables dynamic behavior to be
// added to the simulation fixture.
class CosimulatedDevice {
public:
  virtual ~CosimulatedDevice(){};

  virtual void tick() = 0;
};

// Simple stimulus class used to trigger basic operations
class PulseStimulus : public CosimulatedDevice {
  unsigned long width;
  unsigned char &signal;

  unsigned long counter = 0;

public:
  PulseStimulus(unsigned char &signal, unsigned long width) : signal(signal) {
    this->width = width;
  }
  virtual void tick() override {
    if (counter < width) {
      signal = 1;
    } else {
      signal = 0;
    }
    counter++;
  }
};

class FakeBRAM : public CosimulatedDevice {
  std::array<unsigned long, 512> ram{};

  std::queue<unsigned long> addr_lookup_q;

  unsigned short &addr_in;
  unsigned long &data_out;

  unsigned char &clk;
  unsigned char prev_clk;

public:
  FakeBRAM(int latency, unsigned char &clk, unsigned short &addr_in,
           unsigned long &data_out)
      : addr_lookup_q(), addr_in(addr_in), data_out(data_out), clk(clk) {
    for (int i = 0; i < latency; i++) {
      addr_lookup_q.push(0);
    }
  };

  FakeBRAM(int latency, unsigned char &clk, unsigned short &addr_in,
           unsigned long &data_out, std::array<unsigned long, 512> ram)
      : addr_lookup_q(), addr_in(addr_in), data_out(data_out), clk(clk),
        ram(ram) {
    for (int i = 0; i < latency; i++) {
      addr_lookup_q.push(0);
    }
    prev_clk = clk;
  }

  void tick() {
    // we push and pop in the same tick: this way we keep the queue the same
    // size, acting as a pipeline delay.
    if (prev_clk == 0 && clk == 1) { // rising edge
      addr_lookup_q.push(addr_in);

      auto addr_to_load = addr_lookup_q.front();
      addr_lookup_q.pop();
      data_out = ram.at(addr_to_load);
    }
    prev_clk = clk;
  }

  // TODO: allow accessing/setting the data
  const std::span<uint64_t, 512> get() { return this->ram; }

  void write(unsigned short addr, unsigned long data) { ram[addr] = data; }
};

// test fixture to reduce amount of runtime code.
// Supports:
// adding external modules
// running the test
template <typename DUT> class VerilatorTestFixture {
public:
  enum class FinishReason { Ok, Timeout };

private:
  // we call .tick() on all of these. They are bound externally to the DUT.
  std::vector<std::shared_ptr<CosimulatedDevice>> external_devices;
  std::unique_ptr<VerilatedContext> ctx;
  std::unique_ptr<DUT> dut;

  unsigned long timeout =
      1000000; // clock cycles to execute before ending the simulation.
  unsigned long simtime = 0;
  unsigned long posedge = 0;

  // stores the termination-condition for the simulation.
  // if false, it means we timed out. If true, our done_condition returned true.
  enum FinishReason reason;

  typedef std::function<bool(DUT &, unsigned long)> done_callback;
  // callback function to determine execution completion.
  // This can be used to add arbitrary finish-conditions to the execution.
  // If the function returns true, we set the done_set boolean to true;
  done_callback done_check;

  std::unique_ptr<VerilatedVcdC> trace;

public:
  // Create a new test fixture with a given timeout. Everything else (including
  // done-condition) can be set later.
  VerilatorTestFixture() {
    ctx = std::make_unique<VerilatedContext>();
    dut = std::make_unique<DUT>(ctx.get(), "dut");
    dut->eval(); // let values settle before adding modules.
  }

  void set_timeout(unsigned long new_timeout) { this->timeout = new_timeout; }

  void set_done_callback(done_callback d) { this->done_check = d; }

  void add_module(std::shared_ptr<CosimulatedDevice> device) {
    external_devices.push_back(device);
  }

  void enable_trace(std::string name) {
    if (!trace) {
      ctx->traceEverOn(true);
      trace = std::make_unique<VerilatedVcdC>();
      dut->trace(trace.get(), 99);
      trace->open(name.c_str());
    }
  }

  void exec() {
    bool done = false;
    dut->eval(); // pre-eval.
    while (!done) {
      dut->clk ^= 1;
      dut->eval();

      if (trace) {
        trace->dump(10 * simtime);
      }

      if (dut->clk == 1) {
        posedge++;
      }
      if (done_check) {
        if (done_check(*dut, posedge)) {
          reason = FinishReason::Ok;
          done = true;
        }
      }

      if (posedge >= timeout) {
        reason = FinishReason::Timeout;
        done = true;
      }
      // run our external devices
      for (auto dev : external_devices) {
        dev->tick();
      }
      if (trace) {
        trace->dump(10 * simtime + 5);
      }
      dut->eval(); // allow combinational logic to settle if it's being set on
                   // the negative clock.
      if (trace) {
        trace->dump(10 * simtime + 6);
      }
      simtime++;
    }

    if (trace) {
      // close the trace.
      trace->close();
    }
  }

  // return a reference to the DUT itself. Useful for bespoke tests.
  const DUT &get() { return *dut; }

  const FinishReason get_reason() { return reason; }
};