#include "tb_uart.h"

#include "tb_cli.h"
#include "tb_constants.h"

#include <algorithm>
#include <cstdlib>
#include <cerrno>
#include <cstring>
#include <iostream>

#include <fcntl.h>
#include <netinet/tcp.h>
#include <sys/socket.h>
#include <unistd.h>

namespace {

void close_fd(int &fd) {
	if (fd < 0)
		return;
	(void)::close(fd);
	fd = -1;
}

bool set_nonblocking(int fd) {
	int flags = fcntl(fd, F_GETFL, 0);
	if (flags < 0)
		return false;
	return fcntl(fd, F_SETFL, flags | O_NONBLOCK) == 0;
}

int send_flags() {
	int flags = 0;
#ifdef MSG_NOSIGNAL
	flags |= MSG_NOSIGNAL;
#endif
	return flags;
}

} // namespace

tb_uart_state::tb_uart_state(const tb_cli_args &args) {
	uarts[0].init(args.uart0_port);
	uarts[1].init(args.uart1_port);
	if (args.uart0_port != 0)
		std::cout << "UART0 listening on port " << args.uart0_port << "\n";
	if (args.uart1_port != 0)
		std::cout << "UART1 listening on port " << args.uart1_port << "\n";
}

tb_uart_state::~tb_uart_state() {
	for (uint32_t i = 0; i < N_UARTS; ++i)
		uarts[i].close();
}

void tb_uart_state::uart::init(uint16_t port_) {
	port = port_;
	if (port == 0)
		return;

	server_fd = socket(AF_INET, SOCK_STREAM, 0);
	if (server_fd < 0) {
		std::cerr << "UART socket creation failed: " << strerror(errno) << "\n";
		exit(-1);
	}

	int opt = 1;
	if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
		std::cerr << "UART setsockopt(SO_REUSEADDR) failed: " << strerror(errno) << "\n";
		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, &opt, sizeof(opt));
#endif

	bind_addr.sin_family = AF_INET;
	bind_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	bind_addr.sin_port = htons(port);
	if (bind(server_fd, (struct sockaddr *)&bind_addr, sizeof(bind_addr)) < 0) {
		std::cerr << "UART bind failed: " << strerror(errno) << "\n";
		exit(-1);
	}
	if (listen(server_fd, 1) < 0) {
		std::cerr << "UART listen failed: " << strerror(errno) << "\n";
		exit(-1);
	}
	if (!set_nonblocking(server_fd)) {
		std::cerr << "UART fcntl(O_NONBLOCK) failed: " << strerror(errno) << "\n";
		exit(-1);
	}
}

void tb_uart_state::uart::close() {
	close_fd(client_fd);
	close_fd(server_fd);
	port = 0;
	overrun = false;
	rx_fifo.clear();
	tx_fifo.clear();
}

void tb_uart_state::uart::poll_accept(uint32_t uart_idx) {
	if (server_fd < 0 || client_fd >= 0)
		return;

	sockaddr_in peer {};
	socklen_t peer_len = sizeof(peer);
	int fd = accept(server_fd, (struct sockaddr *)&peer, &peer_len);
	if (fd < 0) {
		if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
			return;
		std::cerr << "UART" << uart_idx << " accept failed: " << strerror(errno) << "\n";
		return;
	}

	client_fd = fd;
	(void)set_nonblocking(client_fd);

	// Low-latency local socket traffic helps interactivity.
	int flag = 1;
	setsockopt(client_fd, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(flag));
#ifdef SO_NOSIGPIPE
	// Best-effort: avoid SIGPIPE on some platforms.
	(void)setsockopt(client_fd, SOL_SOCKET, SO_NOSIGPIPE, &flag, sizeof(flag));
#endif

	std::cout << "UART" << uart_idx << " connected\n";
}

void tb_uart_state::uart::poll_rx(uint32_t uart_idx) {
	poll_accept(uart_idx);
	if (client_fd < 0)
		return;

	uint8_t buf[1024];
	while (true) {
		ssize_t n = recv(client_fd, buf, sizeof(buf), MSG_DONTWAIT);
		if (n > 0) {
			for (ssize_t i = 0; i < n; ++i) {
				if (rx_fifo.size() >= rx_capacity) {
					overrun = true;
					continue;
				}
				rx_fifo.push_back(buf[i]);
			}
			continue;
		}
		if (n == 0) {
			close_fd(client_fd);
			std::cout << "UART" << uart_idx << " disconnected\n";
			return;
		}
		if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
			return;
		std::cerr << "UART" << uart_idx << " recv failed: " << strerror(errno) << "\n";
		close_fd(client_fd);
		std::cout << "UART" << uart_idx << " disconnected\n";
		return;
	}
}

void tb_uart_state::uart::poll_tx(uint32_t uart_idx) {
	poll_accept(uart_idx);
	if (client_fd < 0 || tx_fifo.empty())
		return;

	uint8_t buf[1024];
	while (!tx_fifo.empty()) {
		const size_t chunk = std::min(tx_fifo.size(), sizeof(buf));
		for (size_t i = 0; i < chunk; ++i)
			buf[i] = tx_fifo[i];

		ssize_t n = send(client_fd, buf, chunk, send_flags());
		if (n > 0) {
			for (ssize_t i = 0; i < n; ++i)
				tx_fifo.pop_front();
			continue;
		}
		if (n == 0)
			return;
		if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
			return;
		std::cerr << "UART" << uart_idx << " send failed: " << strerror(errno) << "\n";
		close_fd(client_fd);
		std::cout << "UART" << uart_idx << " disconnected\n";
		return;
	}
}

void tb_uart_state::step() {
	for (uint32_t i = 0; i < N_UARTS; ++i) {
		if (!uarts[i].tx_fifo.empty())
			uarts[i].poll_tx(i);
	}
}

uint32_t tb_uart_state::read_status(uint32_t uart_idx) {
	if (uart_idx >= N_UARTS)
		return 0;
	uart &u = uarts[uart_idx];
	u.poll_rx(uart_idx);
	u.poll_tx(uart_idx);
	uint32_t status = TB_UART_STATUS_TX_READY;
	if (!u.rx_fifo.empty())
		status |= TB_UART_STATUS_RX_AVAIL;
	if (u.connected())
		status |= TB_UART_STATUS_CONNECTED;
	if (u.overrun)
		status |= TB_UART_STATUS_OVERRUN;
	return status;
}

uint32_t tb_uart_state::read_data(uint32_t uart_idx) {
	if (uart_idx >= N_UARTS)
		return 0xffffffffu;
	uart &u = uarts[uart_idx];
	u.poll_rx(uart_idx);
	if (u.rx_fifo.empty())
		return 0xffffffffu;
	const uint8_t byte = u.rx_fifo.front();
	u.rx_fifo.pop_front();
	return byte;
}

void tb_uart_state::write_data(uint32_t uart_idx, uint8_t byte) {
	if (uart_idx >= N_UARTS)
		return;
	uart &u = uarts[uart_idx];
	if (u.server_fd < 0)
		return;
	if (u.tx_fifo.size() >= u.tx_capacity && !u.tx_fifo.empty())
		u.tx_fifo.pop_front();
	u.tx_fifo.push_back(byte);
	u.poll_tx(uart_idx);
}

void tb_uart_state::write_ctrl(uint32_t uart_idx, uint32_t value) {
	if (uart_idx >= N_UARTS)
		return;
	uart &u = uarts[uart_idx];
	if (value & TB_UART_CTRL_CLR_OVERRUN)
		u.overrun = false;
}