/*
 * Copyright 2005, Ingo Weinhold <bonefish@cs.tu-berlin.de>.
 * All rights reserved. Distributed under the terms of the MIT License.
 */

#include <boot/net/RemoteDisk.h>

#include <new>

#include <endian.h>
#include <stdio.h>

#include <OS.h>
#include <SupportDefs.h>

#include <boot/net/UDP.h>


static const bigtime_t kRequestTimeout = 100000LL;


using std::nothrow;


#if __BYTE_ORDER == __LITTLE_ENDIAN

static inline
uint64_t swap_uint64(uint64_t data)
{
	return ((data & 0xff) << 56)
		| ((data & 0xff00) << 40)
		| ((data & 0xff0000) << 24)
		| ((data & 0xff000000) << 8)
		| ((data >> 8) & 0xff000000)
		| ((data >> 24) & 0xff0000)
		| ((data >> 40) & 0xff00)
		| ((data >> 56) & 0xff);
}

#define host_to_net64(data)	swap_uint64(data)
#define net_to_host64(data)	swap_uint64(data)

#endif

#if __BYTE_ORDER == __BIG_ENDIAN
#define host_to_net64(data)	(data)
#define net_to_host64(data)	(data)
#endif

#undef htonll
#undef ntohll
#define htonll(data)	host_to_net64(data)
#define ntohll(data)	net_to_host64(data)


// constructor
RemoteDisk::RemoteDisk()
	: fServerAddress(INADDR_ANY),
	  fServerPort(0),
	  fImageSize(0),
	  fRequestID(0),
	  fSocket(NULL),
	  fPacket(NULL)
{
}

// destructor
RemoteDisk::~RemoteDisk()
{
	delete fSocket;
	delete fPacket;
}

// Init
status_t
RemoteDisk::Init(ip_addr_t serverAddress, uint16 serverPort, off_t imageSize)
{
	fServerAddress = serverAddress;
	fServerPort = serverPort;
	fImageSize = imageSize;

	// create and bind socket
	fSocket = new(nothrow) UDPSocket;
	if (!fSocket)
		return B_NO_MEMORY;
	
	status_t error = fSocket->Bind(INADDR_ANY, 6666);
	if (error != B_OK)
		return error;

	return B_OK;
}

// ReadAt
ssize_t
RemoteDisk::ReadAt(void */*cookie*/, off_t pos, void *_buffer,
	size_t bufferSize)
{
	if (!fSocket)
		return B_NO_INIT;

	uint8 *buffer = (uint8*)_buffer;
	if (!buffer || pos < 0)
		return B_BAD_VALUE;

	if (bufferSize == 0)
		return 0;

	// Check whether the current packet already contains the beginning of the
	// data to read.
	ssize_t bytesRead = _ReadFromPacket(pos, buffer, bufferSize);

	// If there still remains something to be read, we need to get it from the
	// server.
	status_t error = B_OK;
	while (bufferSize > 0) {
		// prepare request
		remote_disk_header request;
		request.offset = htonll(pos);
		uint32 toRead = min_c(bufferSize, REMOTE_DISK_BLOCK_SIZE);
		request.size = htons(toRead);
		request.command = REMOTE_DISK_READ_REQUEST;

		// send request
		UDPPacket *packet;
		error = _SendRequest(&request, sizeof(request), REMOTE_DISK_READ_REPLY,
			&packet);
		if (error != B_OK)
			break;

		// check for errors
		int16 packetSize = ntohs(((remote_disk_header*)packet->Data())->size);
		if (packetSize < 0) {
			if (packetSize == REMOTE_DISK_IO_ERROR)
				error = B_IO_ERROR;
			else if (packetSize == REMOTE_DISK_BAD_REQUEST)
				error = B_BAD_VALUE;
			break;
		}

		// make the reply packet the current packet
		delete fPacket;
		fPacket = packet;

		// read from the packet
		size_t packetBytesRead = _ReadFromPacket(pos, buffer, bufferSize);
		if (packetBytesRead == 0)
			break;
		bytesRead += packetBytesRead;
	}

	// only return an error, when we were not able to read anything at all
	return (bytesRead == 0 ? error : bytesRead);
}

// WriteAt
ssize_t
RemoteDisk::WriteAt(void */*cookie*/, off_t pos, const void *buffer,
	size_t bufferSize)
{
	// Not needed in the boot loader.
	return B_PERMISSION_DENIED;
}

// GetName
status_t
RemoteDisk::GetName(char *nameBuffer, size_t bufferSize) const
{
	if (!nameBuffer)
		return B_BAD_VALUE;

	snprintf(nameBuffer, bufferSize,
		"RemoteDisk:%" B_PRIu32 ".%" B_PRIu32 ".%" B_PRIu32 ".%" B_PRIu32 ":%hd",
		(fServerAddress >> 24) & 0xff, (fServerAddress >> 16) & 0xff,
		(fServerAddress >> 8) & 0xff, fServerAddress & 0xff, fServerPort);

	return B_OK;
}

// Size
off_t
RemoteDisk::Size() const
{
	return fImageSize;
}

ip_addr_t
RemoteDisk::ServerIPAddress() const
{
	return fServerAddress;
}

uint16
RemoteDisk::ServerPort() const
{
	return fServerPort;
}

// FindAnyRemoteDisk
RemoteDisk *
RemoteDisk::FindAnyRemoteDisk()
{
	// create a socket and bind it
	UDPSocket socket;
	status_t error = socket.Bind(INADDR_ANY, 6665);
	if (error != B_OK) {
		printf("RemoteDisk::GetAnyRemoteDisk(): Failed to bind socket.\n");
		return NULL;
	}

	// prepare request
	remote_disk_header request;
	request.command = REMOTE_DISK_HELLO_REQUEST;

	// send request
	UDPPacket *packet;
	error = _SendRequest(&socket, INADDR_BROADCAST, REMOTE_DISK_SERVER_PORT,
		&request, sizeof(request), REMOTE_DISK_HELLO_REPLY, &packet);
	if (error != B_OK) {
		printf("RemoteDisk::GetAnyRemoteDisk(): Got no server reply.\n");
		return NULL;
	}
	remote_disk_header *reply = (remote_disk_header*)packet->Data();

	// create a RemoteDisk object
	RemoteDisk *remoteDisk = new(nothrow) RemoteDisk;
	if (remoteDisk) {
		error = remoteDisk->Init(packet->SourceAddress(), ntohs(reply->port),
			ntohll(reply->offset));
		if (error != B_OK) {
			delete remoteDisk;
			remoteDisk = NULL;
		}
	}

	delete packet;

	return remoteDisk;
}

// _ReadFromPacket
ssize_t
RemoteDisk::_ReadFromPacket(off_t &pos, uint8 *&buffer, size_t &bufferSize)
{
	if (!fPacket)
		return 0;

	remote_disk_header *header = (remote_disk_header*)fPacket->Data();
	uint64 packetOffset = ntohll(header->offset);
	uint32 packetSize = ntohs(header->size);
	if (packetOffset > (uint64)pos || packetOffset + packetSize <= (uint64)pos)
		return 0;
		
	// we do indeed have some bytes already
	size_t toCopy = size_t(packetOffset + packetSize - (uint64)pos);
	if (toCopy > bufferSize)
		toCopy = bufferSize;
	memcpy(buffer, header->data + (pos - packetOffset), toCopy);
	
	pos += toCopy;
	buffer += toCopy;
	bufferSize -= toCopy;
	return toCopy;
}

// _SendRequest
status_t
RemoteDisk::_SendRequest(UDPSocket *socket, ip_addr_t serverAddress,
	uint16 serverPort, remote_disk_header *request, size_t size,
	uint8 expectedReply, UDPPacket **_packet)
{
	request->port = htons(socket->Port());

	// try sending the request 3 times at most
	for (int i = 0; i < 3; i++) {
		// send request
		status_t error = socket->Send(serverAddress, serverPort, request, size);
		if (error != B_OK)
			return error;

		// receive reply
		bigtime_t timeout = system_time() + kRequestTimeout;
		do {
			UDPPacket *packet;
			error = socket->Receive(&packet, timeout - system_time());
			if (error == B_OK) {
				// got something; check, if it is looks good
				if (packet->DataSize() >= sizeof(remote_disk_header)) {
					remote_disk_header *reply
						= (remote_disk_header*)packet->Data();
					if (reply->request_id == request->request_id
						&& reply->command == expectedReply) {
						*_packet = packet;
						return B_OK;
					}
				}

				// reply not OK
				delete packet;
			} else if (error != B_TIMED_OUT && error != B_WOULD_BLOCK)
				return error;

		} while (timeout > system_time());
	}

	// no reply
	return B_ERROR;
}

// _SendRequest
status_t
RemoteDisk::_SendRequest(remote_disk_header *request, size_t size,
	uint8 expectedReply, UDPPacket **packet)
{
	request->request_id = fRequestID++;
	return _SendRequest(fSocket, fServerAddress, fServerPort, request, size,
		expectedReply, packet);
}