#include "Vtb.h"
#include "Vtb___024root.h"
#include "verilated.h"

#include <iostream>
#include <algorithm>
#include <cctype>
#include <cstdint>
#include <string>
#include <stdio.h>
#include <cstring>

#include "tb.h"
#include "tb_cli.h"
#include "tb_gdb.h"
#include "tb_jtag.h"

class tb_verilator_top: public tb_top {
	VerilatedContext *contextp;
	Vtb *top;

	// Loop-carried address-phase requests
	bus_request req_i;
	bus_request req_d;
	bool req_i_vld = false;
	bool req_d_vld = false;

public:
	tb_verilator_top(const tb_cli_args &parsed_args, int argc, char **argv);
	~tb_verilator_top() {delete top; delete contextp;}

	void step(const tb_cli_args &args, mem_io_state &memio) override;
	void eval() override {top->eval();}

	void set_trst_n(bool trst_n)     override {top->trst_n = trst_n;}
	void set_tck(bool tck)           override {top->tck = tck;}
	void set_tdi(bool tdi)           override {top->tdi = tdi;}
	void set_tms(bool tms)           override {top->tms = tms;}
	bool get_tdo()                   override {return top->tdo;}
	void set_irq(uint32_t mask)      override {top->irq = mask;}
	void set_soft_irq(uint8_t mask)  override {top->soft_irq = mask;}
	void set_timer_irq(uint8_t mask) override {top->timer_irq = mask;}

	Vtb___024root *rootp() { return top->rootp; }

	void reset_dut(const tb_cli_args &args, mem_io_state &memio, uint32_t assert_cycles = 8, uint32_t release_cycles = 0);

};

tb_verilator_top::tb_verilator_top(const tb_cli_args &parsed_args, int argc, char **argv): tb_top {parsed_args} {
	contextp = new VerilatedContext;
	// Verilated context also gets a chance to parse the arguments. Any
	// argument prefixed with "+verilator+" is ignored by our tb arg parsing.
	contextp->commandArgs(argc, argv);
	top = new Vtb{contextp};

	req_i_vld = false;
	req_d_vld = false;
	req_i.reservation_id = 0;
	req_d.reservation_id = 1;

	// Set bus interfaces to generate good IDLE responses at first
	top->i_hready = true;
	top->d_hready = true;

	// Reset + initial clock pulse
	top->eval();
	top->clk = true;
	top->tck = true;
	top->eval();
	top->clk = false;
	top->tck = false;
	top->trst_n = true;
	top->rst_n = true;
	top->eval();
}

void tb_verilator_top::reset_dut(const tb_cli_args &args, mem_io_state &memio, uint32_t assert_cycles, uint32_t release_cycles) {
	// Clear any loop-carried bus phase state.
	req_i_vld = false;
	req_d_vld = false;

	// Ensure a clean default handshake state at reset release.
	top->i_hready = true;
	top->d_hready = true;

	// Assert global reset (and JTAG TAP reset) for a few core cycles.
	top->trst_n = false;
	top->rst_n = false;
	for (uint32_t i = 0; i < assert_cycles; ++i)
		step(args, memio);

	// Release reset. By default, we don't advance any cycles after release so
	// the CPU is effectively "reset+halted" until the debugger resumes.
	top->trst_n = true;
	top->rst_n = true;
	for (uint32_t i = 0; i < release_cycles; ++i)
		step(args, memio);
}

void tb_verilator_top::step(const tb_cli_args &args, mem_io_state &memio) {
	top->clk = false;
	top->eval();
	top->clk = true;
	top->eval();

	// The two bus ports are handled identically. This enables swapping out of
	// various `tb.v` hardware integration files containing:
	//
	// - A single, dual-ported processor (instruction fetch, load/store ports)
	// - A single, single-ported processor (instruction fetch + load/store muxed internally)
	// - A pair of single-ported processors, for dual-core debug tests

	if (top->d_hready) {
		// Clear bus error by default
		top->d_hresp = false;

		// Handle current data phase
		req_d.wdata = top->d_hwdata;
		bus_response resp;
		if (req_d_vld)
			resp = mem_callback_d(*this, memio, req_d);
		else
			resp.exokay = !memio.monitor_enabled;
		if (resp.err) {
			// Phase 1 of error response
			top->d_hready = false;
			top->d_hresp = true;
		}
		if (req_d_vld && !req_d.write) {
			top->d_hrdata = resp.rdata;
		} else {
			top->d_hrdata = rand();
		}
		top->d_hexokay = resp.exokay;
	} else {
		// hready=0. Currently this only happens when we're in the first
		// phase of an error response, so go to phase 2.
		top->d_hready = true;
	}

	req_d_vld = false;
	if (top->d_hready) {
		// Progress current address phase to data phase
		req_d_vld = top->d_htrans >> 1;
		req_d.write = top->d_hwrite;
		req_d.size = (bus_size_t)top->d_hsize;
		req_d.addr = top->d_haddr;
		req_d.excl = top->d_hexcl;
	}

	if (top->i_hready) {
		top->i_hresp = false;

		req_i.wdata = top->i_hwdata;
		bus_response resp;
		if (req_i_vld)
			resp = mem_callback_i(*this, memio, req_i);
		else
			resp.exokay = !memio.monitor_enabled;
		if (resp.err) {
			// Phase 1 of error response
			top->i_hready = false;
			top->i_hresp = true;
		}
		if (req_i_vld && !req_i.write) {
			top->i_hrdata = resp.rdata;
		} else {
			top->i_hrdata = rand();
		}
		top->i_hexokay = resp.exokay;
	} else {
		// hready=0. Currently this only happens when we're in the first
		// phase of an error response, so go to phase 2.
		top->i_hready = true;
	}

	req_i_vld = false;
	if (top->i_hready) {
		// Progress current address phase to data phase
		req_i_vld = top->i_htrans >> 1;
		req_i.write = top->i_hwrite;
		req_i.size = (bus_size_t)top->i_hsize;
		req_i.addr = top->i_haddr;
		req_i.excl = top->i_hexcl;
	}
}

int main(int argc, char **argv) {
	tb_cli_args args;
	tb_parse_args(argc, argv, args);

	VerilatedContext *contextp = new VerilatedContext;
	contextp->commandArgs(argc, argv);

	tb_jtag_state jtag(args);
	mem_io_state memio(args);
	tb_verilator_top tb(args, argc, argv);

	bool timed_out = false;
	int64_t cycle = 0;

	if (args.gdb_port != 0) {
		struct verilator_gdb_target final : tb_gdb_target {
			tb_verilator_top &tb;
			mem_io_state &memio;
			const tb_cli_args &args;
			Vtb___024root *root;
			int64_t &cycle;
			bool &timed_out;

			verilator_gdb_target(tb_verilator_top &tb_, mem_io_state &memio_, const tb_cli_args &args_, int64_t &cycle_, bool &timed_out_)
				: tb(tb_), memio(memio_), args(args_), root(tb_.rootp()), cycle(cycle_), timed_out(timed_out_) {}

				uint32_t read_reg(uint32_t regno) override {
					if (regno == 0)
						return 0;
					if (regno < 32) {
						// GDB single-step stops with decode PC already advanced to the
						// next instruction, while the just-computed register result can
						// still be sitting in the M stage for one cycle before it is
						// committed into the architectural register file. Overlay the
						// pending writeback so the debugger sees a coherent post-step
						// snapshot.
						if (root->tb__DOT__cpu__DOT__core__DOT__m_reg_wen_if_nonzero &&
							regno == root->tb__DOT__cpu__DOT__core__DOT__xm_rd) {
							return (uint32_t)root->tb__DOT__cpu__DOT__core__DOT__m_result;
						}
						return (uint32_t)root->tb__DOT__cpu__DOT__core__DOT__regs__DOT__real_dualport_reset__DOT__mem[regno];
					}
					if (regno == 32) {
						return (uint32_t)root->tb__DOT__cpu__DOT__core__DOT__decode_u__DOT__pc;
					}
				return 0;
			}

			void write_reg(uint32_t regno, uint32_t value) override {
				if (regno == 0)
					return;
				if (regno < 32) {
					root->tb__DOT__cpu__DOT__core__DOT__regs__DOT__real_dualport_reset__DOT__mem[regno] = value;
					tb.eval();
					return;
				}
				if (regno == 32) {
					root->tb__DOT__cpu__DOT__core__DOT__decode_u__DOT__pc = value;
					tb.eval();
					return;
				}
			}

			bool read_mem(uint32_t addr, uint8_t *dst, size_t len) override {
				if (len == 0)
					return true;
				if (addr < MEM_BASE || addr + len > MEM_BASE + MEM_SIZE)
					return false;
				memcpy(dst, memio.mem + (addr - MEM_BASE), len);
				return true;
			}

			bool write_mem(uint32_t addr, const uint8_t *src, size_t len) override {
				if (len == 0)
					return true;
				if (addr < MEM_BASE || addr + len > MEM_BASE + MEM_SIZE)
					return false;
				memcpy(memio.mem + (addr - MEM_BASE), src, len);
				return true;
			}

			tb_gdb_run_result step_instruction() override {
				while (true) {
					if (args.max_cycles != 0 && cycle >= args.max_cycles) {
						timed_out = true;
						return tb_gdb_run_result::timed_out;
					}

					const bool will_retire = root->tb__DOT__cpu__DOT__core__DOT__x_instr_ret;
					memio.step(tb);
					tb.step(args, memio);
					++cycle;

					if (memio.exit_req)
						return tb_gdb_run_result::exited;
					if (args.max_cycles != 0 && cycle >= args.max_cycles) {
						timed_out = true;
						return tb_gdb_run_result::timed_out;
					}
					if (will_retire)
						return tb_gdb_run_result::ok;
				}
			}

			uint32_t exit_code() const override { return memio.exit_code; }

			bool monitor_cmd(const std::string &cmd, std::string &out_console) override {
				auto is_space = [](unsigned char c) { return std::isspace(c) != 0; };

				// Normalise whitespace and case to make matching user-friendly.
				size_t start = 0;
				while (start < cmd.size() && is_space((unsigned char)cmd[start]))
					++start;
				size_t end = cmd.size();
				while (end > start && is_space((unsigned char)cmd[end - 1]))
					--end;

				std::string norm;
				norm.reserve(end - start);
				bool in_space = false;
				for (size_t i = start; i < end; ++i) {
					const unsigned char c = (unsigned char)cmd[i];
					if (is_space(c)) {
						in_space = true;
						continue;
					}
					if (in_space && !norm.empty())
						norm.push_back(' ');
					in_space = false;
					norm.push_back((char)std::tolower(c));
				}

				if (norm == "reset halt" || norm == "reset") {
					// Clear external testbench state (but keep RAM contents).
					memio.mtime = 0;
					memio.mtimecmp[0] = 0;
					memio.mtimecmp[1] = 0;
					memio.exit_req = false;
					memio.exit_code = 0;
					memio.monitor_enabled = false;
					memio.poison_addr = -4u;
					memio.soft_irq_state = 0;
					memio.irq_state = 0;
					for (int i = 0; i < N_RESERVATIONS; ++i) {
						memio.reservation_valid[i] = false;
						memio.reservation_addr[i] = 0;
					}
					for (uint32_t i = 0; i < tb_uart_state::N_UARTS; ++i) {
						memio.uart.uarts[i].overrun = false;
						memio.uart.uarts[i].rx_fifo.clear();
						memio.uart.uarts[i].tx_fifo.clear();
					}
					tb.set_soft_irq(0);
					tb.set_timer_irq(0);
					tb.set_irq(0);

					cycle = 0;
					timed_out = false;

					tb.reset_dut(args, memio);
					out_console = "reset halt\n";
					return true;
				}

				return false;
			}
		};

		verilator_gdb_target target(tb, memio, args, cycle, timed_out);
		while (!memio.exit_req && !timed_out) {
			tb_gdb_server gdb(args.gdb_port, target);
			const tb_gdb_run_result r = gdb.serve();
			if (r == tb_gdb_run_result::error)
				return 1;
		}
	} else {
	while (args.max_cycles == 0 || cycle < args.max_cycles) {
		uint32_t core_cycles_advanced = 0;
		bool jtag_exit_cmd = jtag.step(tb, &memio, &cycle, &timed_out, &core_cycles_advanced);

		if (memio.exit_req) {
			fprintf(tb.logfile, "CPU requested halt. Exit code %d\n", memio.exit_code);
			fprintf(tb.logfile, "Ran for " I64_FMT " cycles\n", cycle);
			break;
		}
		if (timed_out) {
			fprintf(tb.logfile, "Max cycles reached\n");
			break;
		}
		if (jtag_exit_cmd)
			break;

		// If the JTAG handler didn't advance the core clock, free-run for one cycle.
		if (core_cycles_advanced == 0) {
			memio.step(tb);
			tb.step(args, memio);
			++cycle;
			if (memio.exit_req) {
				fprintf(tb.logfile, "CPU requested halt. Exit code %d\n", memio.exit_code);
				fprintf(tb.logfile, "Ran for " I64_FMT " cycles\n", cycle);
				break;
			}
			if (args.max_cycles != 0 && cycle >= args.max_cycles) {
				fprintf(tb.logfile, "Max cycles reached\n");
				timed_out = true;
				break;
			}
		}
	}
	}

	jtag.close();

	for (auto r : args.dump_ranges) {
		fprintf(tb.logfile, "Dumping memory from %08x to %08x:\n", r.first, r.second);
		for (int i = 0; i < r.second - r.first; ++i)
			fprintf(tb.logfile, "%02x%c", memio.mem[r.first + i - MEM_BASE], i % 16 == 15 ? '\n' : ' ');
		fprintf(tb.logfile, "\n");
	}

	if (args.sig_path != "") {
		FILE *sigfile = fopen(args.sig_path.c_str(), "wb");
		for (auto r : args.dump_ranges) {
			for (uint32_t i = 0; i < r.second - r.first; i += 4) {
				fprintf(
					sigfile,
					"%02x%02x%02x%02x\n",
					memio.mem[r.first + i + 3 - MEM_BASE],
					memio.mem[r.first + i + 2 - MEM_BASE],
					memio.mem[r.first + i + 1 - MEM_BASE],
					memio.mem[r.first + i + 0 - MEM_BASE]
				);
			}
		}
		fclose(sigfile);
	}

	if (args.propagate_return_code && timed_out) {
		return -1;
	} else if (args.propagate_return_code && memio.exit_req) {
		return memio.exit_code;
	} else {
		return 0;
	}
}

// Needed on MacOS, haven't looked into why
double sc_time_stamp() {
	return 0.0;
}