#include "tb_cli.h"
#include "tb_jtag.h"

#include <cstring>
#include <cerrno>
#include <stdio.h>
#include <iostream>
#include <netinet/tcp.h>

// This file contains socket management logic and parsing of OpenOCD JTAG
// protocols:
// - remote_bitbang (simple, lockstep with CPU clock)
// - jtag_vpi (batched, higher throughput)


static int wait_for_connection_bitbang(int server_fd, uint16_t port, struct sockaddr *sock_addr, socklen_t *sock_addr_len) {
	int sock_fd;
	printf("Waiting for connection on port %u\n", port);
	if (listen(server_fd, 3) < 0) {
		fprintf(stderr, "listen failed\n");
		exit(-1);
	}
	sock_fd = accept(server_fd, sock_addr, sock_addr_len);
	if (sock_fd < 0) {
		fprintf(stderr, "accept failed\n");
		exit(-1);
	}
	printf("Connected\n");
	return sock_fd;
}

static int wait_for_connection_vpi(int server_fd, uint16_t port, struct sockaddr *sock_addr, socklen_t *sock_addr_len) {
	int sock_fd;
	printf("Listening on port %u\n", port);
	if (listen(server_fd, 3) < 0) {
		fprintf(stderr, "listen failed\n");
		exit(-1);
	}
	sock_fd = accept(server_fd, sock_addr, sock_addr_len);
	if (sock_fd < 0) {
		fprintf(stderr, "accept failed\n");
		exit(-1);
	}
	printf("Connected\n");
	return sock_fd;
}

static uint32_t load_le32(const uint8_t *p) {
	return (uint32_t)p[0]
		| ((uint32_t)p[1] << 8)
		| ((uint32_t)p[2] << 16)
		| ((uint32_t)p[3] << 24);
}

static void store_le32(uint8_t *p, uint32_t v) {
	p[0] = v & 0xffu;
	p[1] = (v >> 8) & 0xffu;
	p[2] = (v >> 16) & 0xffu;
	p[3] = (v >> 24) & 0xffu;
}

static bool get_bit_lsb0(const uint8_t *buf, uint32_t bit_idx) {
	return (buf[bit_idx / 8] >> (bit_idx % 8)) & 0x1;
}

static void set_bit_lsb0(uint8_t *buf, uint32_t bit_idx, bool value) {
	const uint8_t mask = 1u << (bit_idx % 8);
	if (value) {
		buf[bit_idx / 8] |= mask;
	} else {
		buf[bit_idx / 8] &= ~mask;
	}
}

// Perform one JTAG clock cycle (TCK low -> high -> low) and return the TDO bit
// sampled *before* the rising edge. This matches OpenOCD's scan semantics.
static bool jtag_clock(tb_top &tb, bool tms, bool tdi) {
	// Sample TDO (stable between the previous negedge and the next posedge),
	// then generate a posedge+negedge pair with the new TMS/TDI values.
	const bool tdo = tb.get_tdo();
	tb.set_tms(tms);
	tb.set_tdi(tdi);
	tb.set_tck(true);
	tb.eval();

	tb.set_tck(false);
	tb.eval();

	return tdo;
}

static bool send_all(int fd, const uint8_t *buf, size_t len) {
	size_t sent = 0;
	while (sent < len) {
		ssize_t n = send(fd, buf + sent, len - sent, 0);
		if (n < 0) {
			if (errno == EINTR)
				continue;
			return false;
		}
		sent += (size_t)n;
	}
	return true;
}

tb_jtag_state::tb_jtag_state(const tb_cli_args &_args) {
	args = _args;
	transport = transport_t::none;
	server_fd = -1;
	sock_fd = -1;
	sock_opt = 1;
	sock_addr_len = sizeof(sock_addr);
	rx_ptr = 0;
	rx_remaining = 0;
	tx_ptr = 0;
	vpi_rx_count = 0;
	vpi_poll_ctr = 0;
	vpi_poll_backoff = 0;

	if (args.vpi_port != 0) {
		transport = transport_t::jtag_vpi;
		server_fd = socket(AF_INET, SOCK_STREAM, 0);
		if (server_fd < 0) {
			fprintf(stderr, "socket creation failed: %s\n", strerror(errno));
			exit(-1);
		}

		if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &sock_opt, sizeof(sock_opt)) < 0) {
			fprintf(stderr, "setsockopt(SO_REUSEADDR) failed: %s\n", strerror(errno));
			exit(-1);
		}
#ifdef SO_REUSEPORT
		// Best-effort: can fail on some kernels/configs, but is not required.
		(void)setsockopt(server_fd, SOL_SOCKET, SO_REUSEPORT, &sock_opt, sizeof(sock_opt));
#endif

		sock_addr.sin_family = AF_INET;
		sock_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
		sock_addr.sin_port = htons(args.vpi_port);
		if (bind(server_fd, (struct sockaddr *)&sock_addr, sizeof(sock_addr)) < 0) {
			fprintf(stderr, "bind failed: %s\n", strerror(errno));
			exit(-1);
		}

		sock_fd = wait_for_connection_vpi(server_fd, args.vpi_port, (struct sockaddr *)&sock_addr, &sock_addr_len);

		// Low-latency local socket traffic helps performance a lot.
		int flag = 1;
		setsockopt(sock_fd, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(flag));
	} else if (args.port != 0) {
		transport = transport_t::remote_bitbang;
		server_fd = socket(AF_INET, SOCK_STREAM, 0);
		if (server_fd < 0) {
			fprintf(stderr, "socket creation failed: %s\n", strerror(errno));
			exit(-1);
		}

		if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &sock_opt, sizeof(sock_opt)) < 0) {
			fprintf(stderr, "setsockopt(SO_REUSEADDR) failed: %s\n", strerror(errno));
			exit(-1);
		}
#ifdef SO_REUSEPORT
		// Best-effort: can fail on some kernels/configs, but is not required.
		(void)setsockopt(server_fd, SOL_SOCKET, SO_REUSEPORT, &sock_opt, sizeof(sock_opt));
#endif

		sock_addr.sin_family = AF_INET;
		sock_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
		sock_addr.sin_port = htons(args.port);
		if (bind(server_fd, (struct sockaddr *)&sock_addr, sizeof(sock_addr)) < 0) {
			fprintf(stderr, "bind failed: %s\n", strerror(errno));
			exit(-1);
		}

		sock_fd = wait_for_connection_bitbang(server_fd, args.port, (struct sockaddr *)&sock_addr, &sock_addr_len);
	} else if (args.replay_jtag) {
		transport = transport_t::remote_bitbang;
	}

	if (args.dump_jtag) {
		jtag_dump_fd.open(args.jtag_dump_path);
		if (!jtag_dump_fd.is_open()) {
			std::cerr << "Failed to open \"" << args.jtag_dump_path << "\"\n";
			exit(-1);
		}
	}

	if (args.replay_jtag) {
		jtag_replay_fd.open(args.jtag_replay_path);
		if (!jtag_replay_fd.is_open()) {
			std::cerr << "Failed to open \"" << args.jtag_replay_path << "\"\n";
			exit(-1);
		}
	}

}

// Return true if an exit command was received
bool tb_jtag_state::step(tb_top &tb, mem_io_state *memio, int64_t *cycle_count, bool *timed_out, uint32_t *core_cycles_advanced) {
	if (transport == transport_t::none) {
		if (core_cycles_advanced)
			*core_cycles_advanced = 0;
		return false;
	}

	if (core_cycles_advanced)
		*core_cycles_advanced = 0;

	const bool can_advance_core = transport == transport_t::jtag_vpi
		&& memio != nullptr
		&& cycle_count != nullptr
		&& timed_out != nullptr
		&& args.vpi_clk_per_tck != 0;

	auto advance_one_core_cycle = [&]() -> bool {
		if (!can_advance_core)
			return false;
		if (args.max_cycles != 0 && *cycle_count >= args.max_cycles) {
			*timed_out = true;
			return true;
		}
		memio->step(tb);
		tb.step(args, *memio);
		++(*cycle_count);
		if (core_cycles_advanced)
			++(*core_cycles_advanced);
		if (memio->exit_req)
			return true;
		if (args.max_cycles != 0 && *cycle_count >= args.max_cycles) {
			*timed_out = true;
			return true;
		}
		return false;
	};

	if (transport == transport_t::jtag_vpi) {
		// Socket is blocking, but reads are non-blocking via MSG_DONTWAIT so the
		// CPU can free-run when openocd is idle.
		//
		// Important: a recv() syscall every core cycle is very expensive. When
		// OpenOCD is idle (no data available), back off for N core cycles
		// (configurable). When data is available, process immediately (no
		// throughput throttling between packets).
		if (vpi_rx_count == 0 && args.vpi_poll_cycles != 0 && vpi_poll_ctr != 0) {
			--vpi_poll_ctr;
			return false;
		}

		// Protocol constants must match OpenOCD's jtag_vpi driver.
		static constexpr uint32_t CMD_RESET = 0;
		static constexpr uint32_t CMD_TMS_SEQ = 1;
		static constexpr uint32_t CMD_SCAN_CHAIN = 2;
		static constexpr uint32_t CMD_SCAN_CHAIN_FLIP_TMS = 3;
		static constexpr uint32_t CMD_STOP_SIMU = 4;

		static constexpr int OFF_CMD = 0;
		static constexpr int OFF_OUT = 4;
		static constexpr int OFF_IN = 4 + VPI_XFERT_MAX_SIZE;
		static constexpr int OFF_LEN = 4 + VPI_XFERT_MAX_SIZE + VPI_XFERT_MAX_SIZE;
		static constexpr int OFF_NB_BITS = OFF_LEN + 4;

		while (vpi_rx_count < VPI_PKT_SIZE) {
			ssize_t n = recv(sock_fd, vpi_rxbuf + vpi_rx_count, VPI_PKT_SIZE - vpi_rx_count, MSG_DONTWAIT);
			if (n < 0) {
				if (errno == EAGAIN || errno == EWOULDBLOCK) {
					// No data available right now. If we don't have a partial
					// packet, back off for a while to reduce syscall overhead.
					//
					// Use an exponential backoff up to vpi_poll_cycles. This
					// avoids large fixed delays between back-to-back OpenOCD
					// packets (e.g. GDB single-step), while still reducing
					// syscall rate when OpenOCD is truly idle.
					if (vpi_rx_count == 0 && args.vpi_poll_cycles != 0) {
						if (vpi_poll_backoff == 0) {
							vpi_poll_backoff = 1;
						} else if (vpi_poll_backoff < args.vpi_poll_cycles) {
							uint32_t next = vpi_poll_backoff * 2;
							if (next < vpi_poll_backoff)
								next = args.vpi_poll_cycles;
							if (next > args.vpi_poll_cycles)
								next = args.vpi_poll_cycles;
							vpi_poll_backoff = next;
						}
						vpi_poll_ctr = vpi_poll_backoff;
					}
					return false;
				}
				fprintf(stderr, "jtag_vpi recv failed: %s\n", strerror(errno));
				return true;
			}
			if (n == 0) {
				printf("jtag_vpi connection closed\n");
				::close(sock_fd);
				sock_fd = wait_for_connection_vpi(server_fd, args.vpi_port, (struct sockaddr *)&sock_addr, &sock_addr_len);
				int flag = 1;
				setsockopt(sock_fd, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(flag));
				vpi_rx_count = 0;
				vpi_poll_ctr = 0;
				vpi_poll_backoff = 0;
				return false;
			}
			vpi_rx_count += (int)n;
			// Any received data means OpenOCD is active again.
			vpi_poll_backoff = 0;
		}
		vpi_poll_ctr = 0;
		vpi_poll_backoff = 0;

		const uint32_t cmd = load_le32(vpi_rxbuf + OFF_CMD);
		const uint32_t length = load_le32(vpi_rxbuf + OFF_LEN);
		const uint32_t nb_bits = load_le32(vpi_rxbuf + OFF_NB_BITS);
		const uint8_t *buf_out = vpi_rxbuf + OFF_OUT;

		bool got_exit_cmd = false;
		bool core_stop = false;
		switch (cmd) {
		case CMD_RESET: {
			// Best-effort: reset the TAP via TRST, and also provide 5 TMS=1 clocks
			// to force Test-Logic-Reset on designs without TRST.
			tb.set_trst_n(false);
			tb.eval();
			tb.set_trst_n(true);
			tb.eval();
			for (int i = 0; i < 5; ++i) {
				jtag_clock(tb, true, false);
			}
			break;
		}
		case CMD_TMS_SEQ: {
			if (length > (uint32_t)VPI_XFERT_MAX_SIZE || nb_bits > (uint32_t)VPI_XFERT_MAX_SIZE * 8 || length * 8 < nb_bits) {
				fprintf(stderr, "jtag_vpi: invalid tms_seq length=%u nb_bits=%u\n", length, nb_bits);
				got_exit_cmd = true;
				break;
			}
			bool stop_due_to_core = false;
			for (uint32_t i = 0; i < nb_bits; ++i) {
				const bool tms = get_bit_lsb0(buf_out, i);
				jtag_clock(tb, tms, false);
				if (can_advance_core) {
					for (uint32_t j = 0; j < args.vpi_clk_per_tck; ++j) {
						if (advance_one_core_cycle()) {
							stop_due_to_core = true;
							break;
						}
					}
				}
				if (stop_due_to_core)
					break;
			}
			if (stop_due_to_core) {
				// Simulation requested stop (timeout/exit); caller will handle.
				core_stop = true;
			}
			break;
		}
		case CMD_SCAN_CHAIN:
		case CMD_SCAN_CHAIN_FLIP_TMS: {
			if (length > (uint32_t)VPI_XFERT_MAX_SIZE || nb_bits > (uint32_t)VPI_XFERT_MAX_SIZE * 8 || length * 8 < nb_bits) {
				fprintf(stderr, "jtag_vpi: invalid scan length=%u nb_bits=%u\n", length, nb_bits);
				got_exit_cmd = true;
				break;
			}

			uint8_t resp[VPI_PKT_SIZE];
			memset(resp, 0, sizeof(resp));
			store_le32(resp + OFF_CMD, cmd);
			store_le32(resp + OFF_LEN, length);
			store_le32(resp + OFF_NB_BITS, nb_bits);

			uint8_t *buf_in = resp + OFF_IN;
			for (uint32_t i = 0; i < nb_bits; ++i) {
				const bool tdi = get_bit_lsb0(buf_out, i);
				const bool tms = (cmd == CMD_SCAN_CHAIN_FLIP_TMS) && (i + 1 == nb_bits);
				const bool tdo = jtag_clock(tb, tms, tdi);
				set_bit_lsb0(buf_in, i, tdo);
			}

			if (!send_all(sock_fd, resp, sizeof(resp))) {
				fprintf(stderr, "jtag_vpi send failed: %s\n", strerror(errno));
				got_exit_cmd = true;
			}
			break;
		}
		case CMD_STOP_SIMU:
			printf("OpenOCD requested stop simulation\n");
			got_exit_cmd = true;
			break;
		default:
			fprintf(stderr, "jtag_vpi: unknown cmd=%u\n", cmd);
			got_exit_cmd = true;
			break;
		}

		vpi_rx_count = 0;
		if (core_stop)
			return false;
		return got_exit_cmd;
	}

	// If JTAG is enabled, we run the simulator in lockstep with the remote
	// bitbang commands, to get more consistent simulation traces. This slows
	// down simulation quite a bit compared with normal free-running.
	//
	// Most bitbang commands complete in one cycle (e.g. TCK/TMS/TDI writes)
	// but reads take 0 cycles, step=false.
	bool got_exit_cmd = false;
	bool step = false;
	while (!step) {
		if (rx_remaining > 0) {
			char c = rxbuf[rx_ptr++];
			--rx_remaining;

			if (c == 'r' || c == 's') {
				tb.set_trst_n(true);
				step = true;
			} else if (c == 't' || c == 'u') {
				tb.set_trst_n(false);
			} else if (c >= '0' && c <= '7') {
				int mask = c - '0';
				tb.set_tck(mask & 0x4);
				tb.set_tms(mask & 0x2);
				tb.set_tdi(mask & 0x1);
				step = true;
			} else if (c == 'R') {
				if (!args.replay_jtag) {
					txbuf[tx_ptr++] = tb.get_tdo() ? '1' : '0';
					if (tx_ptr >= TCP_BUF_SIZE || rx_remaining == 0) {
						send(sock_fd, txbuf, tx_ptr, 0);
						tx_ptr = 0;
					}
				}
			} else if (c == 'Q') {
				printf("OpenOCD sent quit command\n");
				got_exit_cmd = true;
				step = true;
			}
		} else {
			// Potentially the last command was not a read command, but
			// OpenOCD is still waiting for a last response from its
			// last command packet before it sends us any more, so now is
			// the time to flush TX.
			if (tx_ptr > 0) {
				if (!args.replay_jtag)
					send(sock_fd, txbuf, tx_ptr, 0);
				tx_ptr = 0;
			}
			rx_ptr = 0;
			if (args.replay_jtag) {
				rx_remaining = jtag_replay_fd.readsome(rxbuf, TCP_BUF_SIZE);
			} else {
				rx_remaining = read(sock_fd, &rxbuf, TCP_BUF_SIZE);
			}
			if (args.dump_jtag && rx_remaining > 0) {
				jtag_dump_fd.write(rxbuf, rx_remaining);
			}
			if (rx_remaining == 0) {
				if (args.port == 0) {
					// Presumably EOF, so quit.
					got_exit_cmd = true;
				} else {
					// The socket is closed. Wait for another connection.
					sock_fd = wait_for_connection_bitbang(server_fd, args.port, (struct sockaddr *)&sock_addr, &sock_addr_len);
				}
			}
		}
	}
	return got_exit_cmd;
}

void tb_jtag_state::close() {
	if (sock_fd >= 0)
		::close(sock_fd);
	if (server_fd >= 0)
		::close(server_fd);
	if (args.dump_jtag) {
		jtag_dump_fd.close();
	}
	if (args.replay_jtag) {
		jtag_replay_fd.close();
	}
}
