#include "net_stream.h"

#include <arpa/inet.h>
#include <cerrno>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <fcntl.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <sys/socket.h>
#include <unistd.h>

namespace mtnet {

namespace {
constexpr uint16_t MAX_PACKET_LENGTH = 65000;

uint32_t now_ms() {
	using namespace std::chrono;
	return static_cast<uint32_t>(
			duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count());
}

void set_nonblocking(int fd) {
	int fl = fcntl(fd, F_GETFL, 0);
	fcntl(fd, F_SETFL, fl | O_NONBLOCK);
}
} // namespace

NetStream::~NetStream() {
	disconnect();
}

void NetStream::set_state(State s) {
	if (m_state == s) {
		return;
	}
	m_state = s;
	on_state_change(s);
}

bool NetStream::connect(const std::string &host, uint16_t port, uint32_t limit_ms) {
	if (!SecureCipher::ensure_sodium_init()) {
		m_last_error = "sodium_init failed";
		return false;
	}
	disconnect();

	addrinfo hints{};
	hints.ai_family = AF_INET;
	hints.ai_socktype = SOCK_STREAM;
	addrinfo *res = nullptr;
	char portbuf[8];
	std::snprintf(portbuf, sizeof(portbuf), "%u", port);
	if (getaddrinfo(host.c_str(), portbuf, &hints, &res) != 0 || !res) {
		m_last_error = "getaddrinfo(" + host + ") failed";
		return false;
	}

	m_sock = ::socket(res->ai_family, res->ai_socktype, res->ai_protocol);
	if (m_sock < 0) {
		m_last_error = std::string("socket: ") + std::strerror(errno);
		freeaddrinfo(res);
		return false;
	}
	set_nonblocking(m_sock);
	int one = 1;
	setsockopt(m_sock, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));

	int rc = ::connect(m_sock, res->ai_addr, res->ai_addrlen);
	freeaddrinfo(res);
	if (rc == 0) {
		m_connect_deadline_ms = 0;
		set_state(State::Online);
	} else if (errno == EINPROGRESS || errno == EWOULDBLOCK) {
		m_connect_deadline_ms = now_ms() + limit_ms;
		set_state(State::Connecting);
	} else {
		m_last_error = std::string("connect: ") + std::strerror(errno);
		::close(m_sock);
		m_sock = -1;
		return false;
	}
	m_last_error.clear();
	return true;
}

void NetStream::disconnect() {
	if (m_sock >= 0) {
		::close(m_sock);
		m_sock = -1;
	}
	m_recv.clear();
	m_send.clear();
	m_connect_deadline_ms = 0;
	m_cipher.clean_up();
	if (m_state != State::Offline) {
		set_state(State::Offline);
	}
}

bool NetStream::send_packet(const void *struct_bytes, size_t n) {
	if (m_sock < 0 || n < PACKET_HEADER_SIZE) {
		return false;
	}
	// copy so we can encrypt in place without touching the caller's struct
	const auto *src = static_cast<const uint8_t *>(struct_bytes);
	uint8_t stackbuf[512];
	uint8_t *tmp = n <= sizeof(stackbuf) ? stackbuf : new uint8_t[n];
	std::memcpy(tmp, src, n);
	if (m_trace) {
		const size_t shown = n < 96 ? n : 96;
		std::fprintf(stderr, "[tx-plain %zuB]", n);
		for (size_t i = 0; i < shown; ++i) {
			std::fprintf(stderr, " %02x", tmp[i]);
		}
		if (shown != n) {
			std::fprintf(stderr, " ...");
		}
		std::fprintf(stderr, "\n");
	}
	m_cipher.encrypt_in_place(tmp, n); // no-op until activated; advances tx counter
	if (m_trace) {
		const size_t shown = n < 96 ? n : 96;
		std::fprintf(stderr, "[tx-wire %zuB]", n);
		for (size_t i = 0; i < shown; ++i) {
			std::fprintf(stderr, " %02x", tmp[i]);
		}
		if (shown != n) {
			std::fprintf(stderr, " ...");
		}
		std::fprintf(stderr, "\n");
	}
	m_send.write(tmp, n);
	if (tmp != stackbuf) {
		delete[] tmp;
	}
	flush_send();
	return true;
}

bool NetStream::recv_into_buffer() {
	uint8_t chunk[8192];
	ssize_t r = ::recv(m_sock, chunk, sizeof(chunk), 0);
	if (r > 0) {
		if (m_cipher.is_activated()) {
			m_cipher.decrypt_in_place(chunk, static_cast<size_t>(r));
		}
		if (m_trace) {
			size_t n = (size_t)r < 48 ? (size_t)r : 48;
			std::fprintf(stderr, "[rx %zdB]", (ssize_t)r);
			for (size_t i = 0; i < n; ++i) {
				std::fprintf(stderr, " %02x", chunk[i]);
			}
			std::fprintf(stderr, "\n");
		}
		m_recv.write(chunk, static_cast<size_t>(r));
		return true;
	}
	if (r == 0) {
		m_last_error = "peer closed";
		return false;
	}
	if (errno == EWOULDBLOCK || errno == EAGAIN) {
		return true;
	}
	m_last_error = std::string("recv: ") + std::strerror(errno);
	return false;
}

bool NetStream::flush_send() {
	while (m_send.readable() > 0) {
		ssize_t w = ::send(m_sock, m_send.read_ptr(), m_send.readable(), 0);
		if (w > 0) {
			m_send.discard(static_cast<size_t>(w));
			continue;
		}
		if (w < 0 && (errno == EWOULDBLOCK || errno == EAGAIN)) {
			return true; // try again next process()
		}
		m_last_error = std::string("send: ") + std::strerror(errno);
		return false;
	}
	return true;
}

void NetStream::process() {
	if (m_sock < 0) {
		return;
	}
	if (m_state == State::Connecting && m_connect_deadline_ms != 0 &&
			static_cast<int32_t>(now_ms() - m_connect_deadline_ms) >= 0) {
		m_last_error = "connect timeout";
		on_disconnect();
		disconnect();
		return;
	}

	fd_set rfd, wfd;
	FD_ZERO(&rfd);
	FD_ZERO(&wfd);
	FD_SET(m_sock, &rfd);
	FD_SET(m_sock, &wfd);
	timeval tv{0, 0};
	if (select(m_sock + 1, &rfd, &wfd, nullptr, &tv) < 0) {
		return;
	}

	if (m_state == State::Connecting) {
		if (FD_ISSET(m_sock, &wfd)) {
			int err = 0;
			socklen_t l = sizeof(err);
			getsockopt(m_sock, SOL_SOCKET, SO_ERROR, &err, &l);
			if (err != 0) {
				m_last_error = std::string("connect: ") + std::strerror(err);
				on_disconnect();
				disconnect();
				return;
			}
			set_state(State::Online);
		} else {
			return;
		}
	}

	if (FD_ISSET(m_sock, &wfd) && m_send.readable() > 0) {
		if (!flush_send()) {
			on_disconnect();
			disconnect();
			return;
		}
	}
	if (FD_ISSET(m_sock, &rfd)) {
		if (!recv_into_buffer()) {
			on_disconnect();
			disconnect();
			return;
		}
	}
	if (!dispatch()) {
		on_disconnect();
		disconnect();
	}
}

bool NetStream::dispatch() {
	for (;;) {
		PacketHeaderPeek hp;
		if (!m_recv.peek(&hp, sizeof(hp))) {
			return true; // need more bytes
		}
		// skip zero padding from cipher block alignment
		if (hp.header == 0) {
			uint16_t z;
			m_recv.read(&z, sizeof(z));
			continue;
		}

		// give a subclass first crack at framing this header (guild-mark stream).
		switch (on_raw(hp.header)) {
			case Raw::NeedMore:
				return true;
			case Raw::Consumed:
				continue;
			case Raw::Error:
				return false;
			case Raw::NotHandled:
				break;
		}

		DynHeader dh;
		if (!m_recv.peek(&dh, sizeof(dh))) {
			return true;
		}
		if (dh.length < PACKET_HEADER_SIZE || dh.length > MAX_PACKET_LENGTH) {
			m_last_error = "bad packet length " + std::to_string(dh.length) +
					" for header 0x" + std::to_string(dh.header);
			return false;
		}
		if (!m_recv.has(dh.length)) {
			return true; // wait for the whole packet
		}

		if (m_trace) {
			std::fprintf(stderr, "[frame] header=0x%04X length=%u avail=%zu\n", dh.header,
					dh.length, m_recv.readable());
		}
		bool ok = true;
		switch (dh.header) {
			case GC_PHASE:
				ok = handle_phase();
				break;
			case GC_PING:
				ok = handle_ping();
				break;
			case GC_KEY_CHALLENGE:
				ok = handle_key_challenge();
				break;
			case GC_KEY_COMPLETE:
				ok = handle_key_complete();
				break;
			default: {
				size_t before = m_recv.readable();
				ok = on_packet(dh.header, dh.length);
				size_t after = m_recv.readable();
				if (m_trace) {
					std::fprintf(stderr, "[frame]   consumed %zu (expected %u)\n", before - after,
							dh.length);
				}
				// safety: on_packet must consume at least the 4-byte frame header,
				// otherwise dispatch() would spin on the same bytes forever.
				if (ok && before == after) {
					m_last_error = "on_packet consumed nothing for header 0x" +
							std::to_string(dh.header);
					return false;
				}
				break;
			}
		}
		if (!ok) {
			return false;
		}
	}
}

bool NetStream::handle_phase() {
	GCPhase p;
	if (!m_recv.read(&p, sizeof(p))) {
		return false;
	}
	on_phase(p.phase);
	return true;
}

bool NetStream::handle_ping() {
	GCPing p;
	if (!m_recv.read(&p, sizeof(p))) {
		return false;
	}
	CGPong pong{CG_PONG, sizeof(CGPong)};
	return send_packet(&pong, sizeof(pong));
}

bool NetStream::handle_key_challenge() {
	GCKeyChallenge kc;
	if (!m_recv.read(&kc, sizeof(kc))) {
		return false;
	}
	if (!m_cipher.initialize() || !m_cipher.compute_client_keys(kc.server_pk)) {
		m_last_error = "cipher key exchange failed";
		return false;
	}
	CGKeyResponse resp{};
	resp.header = CG_KEY_RESPONSE;
	resp.length = sizeof(resp);
	m_cipher.get_public_key(resp.client_pk);
	m_cipher.compute_challenge_response(kc.challenge, resp.challenge_response);
	return send_packet(&resp, sizeof(resp));
}

bool NetStream::handle_key_complete() {
	GCKeyComplete kc;
	if (!m_recv.read(&kc, sizeof(kc))) {
		return false;
	}
	uint8_t token[SecureCipher::SESSION_TOKEN_SIZE];
	if (!m_cipher.decrypt_token(kc.encrypted_token, sizeof(kc.encrypted_token), kc.nonce, token)) {
		m_last_error = "session token decrypt failed";
		return false;
	}
	m_cipher.set_session_token(token);
	m_cipher.set_activated(true);
	// bytes already buffered after this packet are the first ciphertext bytes;
	// rx counter is 0, decrypt them in place now.
	size_t pending = m_recv.readable();
	if (pending > 0) {
		m_cipher.decrypt_in_place(m_recv.mutable_unread(), pending);
	}
	on_cipher_active();
	return true;
}

} // namespace mtnet