#include "models/lt_bus.h"
#include "utils/log.h"

#include <tlm_utils/simple_initiator_socket.h>
#include <tlm_utils/simple_target_socket.h>

using sc_core::SC_NS;
using sc_core::sc_time;
using sc_core::sc_time_stamp;
using sc_core::SC_ZERO_TIME;

// -- TARGET -------------------------------------------------------------------

struct target : public sc_core::sc_module {
  explicit target(sc_core::sc_module_name nm) : sc_module(std::move(nm)) {
    p_sock.register_b_transport(this, &target::b_transport);
    p_sock.register_get_direct_mem_ptr(this, &target::get_direct_mem_ptr);
  }

  tlm_utils::simple_target_socket<target> p_sock{"sock"};

 private:
  void b_transport(tlm::tlm_generic_payload& tx, sc_core::sc_time&) {
    CLOGM(YELLOW, "transport 0x%llx w: %d r: %d", tx.get_address(),
          tx.is_write(), tx.is_read());

    // Miss-use byte enable len for sleep duration.
    if (u32 len = tx.get_byte_enable_length()) {
      CLOGM(YELLOW, "wait(%u)", len);
      wait(sc_time(len, SC_NS));
    }

    tx.set_response_status(tlm::TLM_OK_RESPONSE);
  }

  bool get_direct_mem_ptr(tlm::tlm_generic_payload&, tlm::tlm_dmi&) {
    return true;
  }
};

// -- INITIATOR ----------------------------------------------------------------

struct initiator : public sc_core::sc_module {
  explicit initiator(sc_core::sc_module_name nm) : sc_module(std::move(nm)) {
    SC_THREAD(run0);
    SC_THREAD(run1);
    SC_THREAD(run2);
  }

  tlm_utils::simple_initiator_socket<initiator> p_sock0{"sock0"};
  tlm_utils::simple_initiator_socket<initiator> p_sock1{"sock1"};
  tlm_utils::simple_initiator_socket<initiator> p_sock2{"sock2"};

 private:
  void setup_tx(tlm::tlm_generic_payload& tx, u64 addr, u32 len) const {
    tx.set_command(tlm::TLM_WRITE_COMMAND);
    tx.set_address(addr);
    tx.set_data_ptr(nullptr);
    tx.set_data_length(4);
    tx.set_byte_enable_ptr(nullptr);
    tx.set_byte_enable_length(len);
    tx.set_response_status(tlm::TLM_INCOMPLETE_RESPONSE);
    tx.set_dmi_allowed(false);
  }

  void send_tx(tlm_utils::simple_initiator_socket<initiator>& sock,
               u64 addr,
               u32 len,
               bool lock = false) {
    tlm::tlm_generic_payload tx;
    setup_tx(tx, addr, len);

    auto* bl = new bus_lock;
    bl->is_lock = lock;
    tx.set_extension(bl);

    CLOGM(CYAN, "[%s] ACCESS @0x%lx lock=%d", sock.name(), addr, lock);

    sc_core::sc_time t;
    sock->b_transport(tx, t);

    CLOGM(CYAN, "[%s] ACCESS ok: %d (%s)", sock.name(), tx.is_response_ok(),
          tx.get_response_string().c_str());
  }

  bool get_dmi(tlm_utils::simple_initiator_socket<initiator>& sock, u64 addr) {
    tlm::tlm_generic_payload tx;
    setup_tx(tx, addr, 1 /* len */);

    tlm::tlm_dmi dmi;
    return sock->get_direct_mem_ptr(tx, dmi);
  }

  void run0() {
    // Lock the bus.
    send_tx(p_sock0, 0x1000, 0, true);  // @0ns
    // Spend some time, keep the bus locked.
    wait(500, SC_NS);
    // Unlock the bus.
    send_tx(p_sock0, 0x1000, 0, false);  // @500ns
    assert(sc_time_stamp() == sc_time(500, SC_NS));

    // -- 500ns --

    wait(10, SC_NS);
    LOGM("------------------------------------------------------------");
    wait(90, SC_NS);

    // Waiting for bus lock, as run1 & run2 have currently an outstanding tx.
    assert(sc_time_stamp() == sc_time(600, SC_NS));
    send_tx(p_sock0, 0x1000, 0, true);  // @600ns
    assert(sc_time_stamp() == sc_time(750, SC_NS));
  }

  void run1() {
    wait(10, SC_NS);
    // No dmi while bus is locked.
    assert(get_dmi(p_sock1, 0x2000) == false);

    send_tx(p_sock1, 0x2000, 0);  // @10ns
    assert(sc_time_stamp() == sc_time(500, SC_NS));

    // -- 500ns --

    wait(50, SC_NS);
    send_tx(p_sock1, 0x2000, 100);  // @550ns
    assert(sc_time_stamp() == sc_time(650, SC_NS));
  }

  void run2() {
    wait(20, SC_NS);
    // No dmi while bus is locked.
    assert(get_dmi(p_sock2, 0x3000) == false);

    send_tx(p_sock2, 0x3000, 0);
    assert(sc_time_stamp() == sc_time(500, SC_NS));

    // -- 500ns --

    wait(50, SC_NS);
    send_tx(p_sock1, 0x3000, 200);  // @550ns
    assert(sc_time_stamp() == sc_time(750, SC_NS));
  }
};

// -- SC_MAIN ------------------------------------------------------------------

extern "C" int sc_main(int, char*[]) {
  lt_bus bus{"bus"};

  target target1{"target1"};

  initiator init1{"init1"};

  bus.attach_target(target1.p_sock, 0x0000, 0xffff);

  bus.attach_initiator(init1.p_sock0);
  bus.attach_initiator(init1.p_sock1);
  bus.attach_initiator(init1.p_sock2);

  bus.show_mmap();

  sc_core::sc_start();

  return 0;
}