#include <cstdio>

#include "socket.h"
#include "tunnel.h"



Tunnel::Tunnel():
	time(monotonicTime()),
	udpSendDuration(1000),
	udpPartialSendDuration(500000),
	udpResendDuration(udpSendDuration * PACKETS_COUNT),
	udpConfirmDuration(udpResendDuration/8),
	udpSilenceDuration(300000000),
	udpKeepAliveDuration(udpSilenceDuration/32),
	pollDuration(5000000),
	udpSockId(-1),
	tcpSockId(-1),
	nextTermSendTime(time + udpConfirmDuration)
	{ }


Tunnel::~Tunnel()
	{ close(); }


bool Tunnel::initUdpServer(const Address &address, const Address &remoteTcpAddress, const char *key) {
	closeUdpServer();
	printf("Init UDP server at address: "); address.print();
	if (remoteTcpAddress.port)
		{ printf(", with remote TCP address: "); remoteTcpAddress.print(); }
	printf("\n");
	
	if (!crypt.setKey(key))
		{ printf("bad key, cancelled"); return false; }
	
	udpSockId = Socket::udpBind(address);
	if (udpSockId < 0)
		{ crypt.resetKey(); printf("cannot bind, cancelled"); return false; }
	
	printf("success");
	this->remoteTcpAddress = remoteTcpAddress;
	return true;
}


bool Tunnel::initTcpServer(const Address &address, const Address &remoteUdpAddress) {
	closeTcpServer();
	printf("Init TCP server at address: "); address.print(); printf("\n");
	
	if (udpSockId < 0)
		{ printf("UDP server was not initialized, cancelled\n"); return false; }
	
	tcpSockId = Socket::tcpListen(address);
	if (tcpSockId < 0)
		{ printf("cannot listen, cancelled"); return false; }
	
	printf("success");
	this->remoteUdpAddress = remoteUdpAddress;
	return true;
}


void Tunnel::closeUdpServer() {
	if (udpSockId < 0) return;
	closeTcpServer();
	
	printf("Close UDP server");
	for(ConnMap::iterator i = connections.begin(); i != connections.end(); ++i)
		Socket::close(i->second.tcpSockId);
	connections.clear();
	termQueue.clear();
	Socket::close(udpSockId);
	udpSockId = -1;
	remoteTcpAddress = Address();
	crypt.resetKey();
}


void Tunnel::closeTcpServer() {
	if (tcpSockId < 0) return;
	printf("Close TCP server");
	Socket::close(tcpSockId);
	tcpSockId = -1;
	remoteUdpAddress = Address();
}


void Tunnel::close()
	{ closeUdpServer(); }


bool Tunnel::recvUdpPacket() {
	ConnId id;
	unsigned char raw[CRYPT_PACKET_SIZE] = {};
	int rawSize = Socket::udpRecvAsync(udpSockId, id.address, raw, sizeof(raw));
	if (rawSize <= 0)
		return false;
	
	unsigned char data[CRYPT_PACKET_SIZE] = {};
	int dataSize = CRYPT_PACKET_SIZE;
	if (!crypt.decrypt(raw, rawSize, data, dataSize))
		return true;
	
	void *ptr = data;
	Packet &packet = *(Packet*)ptr;
	
	int cmd = packet.getCommand();
	if (cmd < 0 || cmd > CMDMAX)
		return true;

	int size = packet.getSize();
	if (size + HEADER_SIZE > dataSize)
		return true;
	
	id.id = packet.connId;
	ConnMap::iterator i = connections.find(id);
	if (i == connections.end() && cmd == CMD_DATA && remoteTcpAddress.port && !closedConnections.count(id)) {
		printf("Open incoming connection from UDP: ");
		id.address.print();
		printf(", id: %u\n", id.id);
		int sockId = Socket::tcpConnect(remoteTcpAddress);
		if (sockId < 0) {
			printf("cannot init TCP connection, cancelled\n");
		} else {
			i = connections.emplace(id, Connection::Args(*this, id, sockId)).first;
		}
	}
	
	if (i == connections.end()) {
		termQueue.insert(id);
		return true;
	}
	
	Connection &conn = i->second;
	if (!conn.udpActive) {
		if (cmd != CMD_TERM) termQueue.insert(id);
		return true;
	}
	
	if (cmd == CMD_DATA) {
		conn.lastUdpInputTime = time;
		conn.udpRecvQueue.recvUdp(packet.index, packet.payload, size);
	} else
	if (cmd == CMD_CONFIRM) {
		conn.lastUdpInputTime = time;
		conn.udpSendQueue.confirm(packet.index, packet.payload, size*8);
	} else
	if (cmd == CMD_TERM) {
		conn.udpActive = false;
	}
	
	return true;
}


void Tunnel::sendUdpPacket(const Address &address, const Packet &packet) {
	unsigned char raw[CRYPT_PACKET_SIZE] = {};
	int rawSize = CRYPT_PACKET_SIZE;
	if (crypt.encrypt(&packet, HEADER_SIZE + packet.getSize(), raw, rawSize))
		Socket::udpSend(udpSockId, address, raw, rawSize);
}


bool Tunnel::iteration() {
	if (udpSockId < 0)
		return false;
	
	// prepare poll
	
	poll.list.clear();
	
	poll.list.emplace_back();
	Poll::Entry &e = poll.list.back();
	e.sockId = udpSockId;
	e.wantRead = true;
	
	if (tcpSockId >= 0) {
		Poll::Entry &e = poll.list.back();
		e.sockId = tcpSockId;
		e.wantRead = true;
	}
	
	Time pollTime = time + pollDuration;
	for(ConnMap::iterator i = connections.begin(); i != connections.end(); ++i) {
		Connection &conn = i->second;
		if (conn.tcpSockId >= 0) {
			poll.list.emplace_back();
			Poll::Entry &e = poll.list.back();
			e.connection = &conn;
			e.sockId = conn.tcpSockId;
			e.wantRead = conn.udpActive && conn.wantReadFromTcp();
			e.wantWrite = conn.wantWriteToTcp();
		}
		if (conn.udpActive) {
			Time t = conn.whenWriteToUdp();
			if (timeLess(pollTime, t)) pollTime = t;
		}
	}
	
	// wait for events
	poll.wait(pollTime);
	time = monotonicTime();
	crypt.setTime(time);
	
	if (poll.list[0].closed || (tcpSockId >= 0 && poll.list[1].closed)) {
		poll.list.clear();
		close();
		return false;
	}
	
	// read and parse udp
	if (poll.list[0].canRead)
		while(recvUdpPacket());
	
	if (tcpSockId >= 0 && poll.list[1].canRead) {
		// accept new tcp connection
		int sockId = Socket::tcpAccept(tcpSockId);
		if (sockId >= 0) {
			ConnId id;
			id.address = remoteUdpAddress;
			id.id = ConnId::generateId();
			Connection &conn = connections.emplace(id, Connection::Args(*this, id, tcpSockId)).first->second;
			conn.tcpRecvQueue.readFromTcp();
		}
	}
	
	// tcp read and write
	int first = tcpSockId < 0 ? 1 : 2;
	for(Poll::List::const_iterator i = poll.list.begin() + first; i != poll.list.end(); ++i) {
		Connection &conn = *i->connection;
		conn.udpRecvQueue.sendTcp();
		if (i->canWrite) while(true) {
			int size = conn.tcpSendQueue.writeToTcp();
			if (size <= 0) break;
			conn.udpRecvQueue.sendTcp();
		}
		if (i->canRead && conn.udpActive)
			conn.tcpRecvQueue.readFromTcp();
		if (i->closed || (!conn.udpActive && !conn.tcpSendQueue.busySize())) {
			Socket::close(conn.tcpSockId);
			conn.tcpSockId = -1;
		}
	}
	
	// write to udp
	for(ConnMap::iterator i = connections.begin(); i != connections.end(); ++i) {
		Connection &conn = i->second;
		if (conn.udpActive) {
			bool sent = false;
			if (conn.udpSendQueue.sendUdp()) sent = true;
			if (conn.udpRecvQueue.sendUdpConfirm()) sent = true;
			if (sent) conn.lastUdpOutputTime = time;
			if (timeLequal(conn.lastUdpOutputTime + udpKeepAliveDuration, time)) {
				conn.udpRecvQueue.sendUdpConfirm(true);
				conn.lastUdpOutputTime = time;
			}
		}
	}
	
	// close connections	
	for(ConnMap::iterator i = connections.begin(); i != connections.end();) {
		ConnMap::iterator j = i++;
		Connection &conn = j->second;
		
		if (conn.udpActive && timeLess(conn.lastUdpInputTime + udpSilenceDuration, time))
			conn.udpActive = false;
		
		if (!conn.udpActive && !conn.wantWriteToTcp()) {
			Socket::close(conn.tcpSockId);
			conn.tcpSockId = -1;
			closedConnections[conn.id] = time + 2*udpSilenceDuration;
			connections.erase(i);
			continue;
		}
		
		if (conn.tcpSockId < 0 && !conn.udpSendQueue.canSentToUdp()) {
			conn.udpActive = false;
			closedConnections[conn.id] = time + 2*udpSilenceDuration;
			connections.erase(i);
			continue;
		}
		
		if (!conn.udpActive && conn.tcpSockId < 0) {
			closedConnections[conn.id] = time + 2*udpSilenceDuration;
			connections.erase(i);
			continue;
		}
	}
	
	// send udp term packets
	if (timeLequal(nextTermSendTime, time)) {
		for(ConnSet::const_iterator i = termQueue.begin(); i != termQueue.end(); ++i) {
			Packet packet;
			packet.init(i->id, 0, CMD_TERM, 0);
			sendUdpPacket(i->address, packet);
		}
		termQueue.clear();
		nextTermSendTime = time + nextTermSendTime;
	}
	
	poll.list.clear();
	return true;
}


void Tunnel::run() {
	time = monotonicTime();
	nextTermSendTime = time + udpConfirmDuration;
	while(iteration());
}