/*
 * Copyright 2016 Google Inc.
 *
 * See file CREDITS for list of people who contributed to this
 * project.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2 of
 * the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but without any warranty; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 * MA 02111-1307 USA
 */

#include <assert.h>
#include <stdint.h>
#include <stdlib.h>

#include "base/container_of.h"
#include "base/die.h"
#include "base/hexdump.h"
#include "base/init_funcs.h"
#include "base/list.h"
#include "base/xalloc.h"
#include "drivers/net/net.h"
#include "net/ethernet.h"

#include "uefi/edk/Protocol/DevicePath.h"
#include "uefi/edk/Protocol/DevicePathToText.h"
#include "uefi/edk/Protocol/SimpleNetwork.h"
#include "uefi/uefi.h"

static EFI_GUID uefi_device_path_protocol_guid =
	EFI_DEVICE_PATH_PROTOCOL_GUID;

static EFI_GUID uefi_device_path_to_text_protocol_guid =
	EFI_DEVICE_PATH_TO_TEXT_PROTOCOL_GUID;

static EFI_GUID uefi_simple_network_protocol_guid =
	EFI_SIMPLE_NETWORK_PROTOCOL_GUID;

typedef struct {
	NetDevice dev;
	ListNode list_node;

	int initialized;

	EFI_SIMPLE_NETWORK_PROTOCOL *snp;
	MacAddress mac;
} UefiNetDevice;

ListNode uefinet_devices;


static int uefinet_initialize(UefiNetDevice *dev)
{
	if (dev->initialized)
		return 0;

	EFI_SIMPLE_NETWORK_PROTOCOL *snp = dev->snp;
	EFI_SIMPLE_NETWORK_MODE *mode = snp->Mode;

	EFI_STATUS status;

	if (mode->State == EfiSimpleNetworkStopped) {
		status = snp->Start(snp);
		if (status != EFI_SUCCESS) {
			printf("Failed to start network device.\n");
			return 1;
		}
	}

	if (mode->State == EfiSimpleNetworkStarted) {
		status = snp->Initialize(snp, 0, 0);
		if (status != EFI_SUCCESS) {
			printf("Failed to initialize network device.\n");
			return 1;
		}
	}

	if (mode->State != EfiSimpleNetworkInitialized) {
		printf("Device was not initialized as expected.\n");
		return 1;
	}

	mode->ReceiveFilterSetting =
		EFI_SIMPLE_NETWORK_RECEIVE_UNICAST |
		EFI_SIMPLE_NETWORK_RECEIVE_MULTICAST |
		EFI_SIMPLE_NETWORK_RECEIVE_BROADCAST |
		EFI_SIMPLE_NETWORK_RECEIVE_PROMISCUOUS |
		EFI_SIMPLE_NETWORK_RECEIVE_PROMISCUOUS_MULTICAST;

	if (mode->HwAddressSize != sizeof(dev->mac)) {
		printf("Unexpected MAC size %d.\n", mode->HwAddressSize);
		return 1;
	}
	memcpy(&dev->mac, &mode->CurrentAddress, sizeof(dev->mac));

	dev->initialized = 1;
	return 0;
}

static int uefinet_ready(NetDevice *net_dev, int *ready)
{
	UefiNetDevice *dev = container_of(net_dev, UefiNetDevice, dev);
	EFI_SIMPLE_NETWORK_PROTOCOL *snp = dev->snp;

	if (uefinet_initialize(dev))
		return 1;

	UINT32 intStatus;
	VOID *txbufPtr;
	EFI_STATUS status = snp->GetStatus(snp, &intStatus, &txbufPtr);
	if (status != EFI_SUCCESS) {
		printf("Failed to get network device status.\n");
		return 1;
	}

	if (snp->Mode->MediaPresentSupported) {
		*ready = snp->Mode->MediaPresent;
	} else {
		printf("Device doesn't support the MediaPresent field.\n");
		*ready = 1;
	}

	return 0;
}

static int uefinet_recv(NetDevice *net_dev, void *buf, uint16_t *len,
			int maxlen)
{
	UefiNetDevice *dev = container_of(net_dev, UefiNetDevice, dev);
	EFI_SIMPLE_NETWORK_PROTOCOL *snp = dev->snp;

	if (uefinet_initialize(dev))
		return 1;

	// Try to receive until there's some data to return.
	UINTN buf_size = maxlen;
	EFI_STATUS status = snp->Receive(snp, NULL, &buf_size, buf,
					 NULL, NULL, NULL);
	if (status == EFI_NOT_READY) {
		*len = 0;
		return 0;
	} else {
		*len = buf_size;
	}

	if (status == EFI_BUFFER_TOO_SMALL) {
		printf("Packet receive buffer too small.\n");
		return 1;
	} else if (status != EFI_SUCCESS) {
		printf("Failed to receive packets from the device.\n");
		return 1;
	}

	return 0;
}

static int uefinet_send(NetDevice *net_dev, void *buf, uint16_t len)
{
	UefiNetDevice *dev = container_of(net_dev, UefiNetDevice, dev);
	EFI_SIMPLE_NETWORK_PROTOCOL *snp = dev->snp;

	if (uefinet_initialize(dev))
		return 1;

	EFI_STATUS status;
	// Try to send until the transmit request is accepted.
	do {
		status = snp->Transmit(snp, 0, len, buf, NULL, NULL, NULL);
	} while (status == EFI_NOT_READY);

	UINT32 intStatus;
	VOID *txbufPtr;
	do {
		status = snp->GetStatus(snp, &intStatus, &txbufPtr);
		if (status != EFI_SUCCESS) {
			printf("Failed to get network device status.\n");
			return 1;
		}
	} while (txbufPtr != buf);
	

	if (status != EFI_SUCCESS) {
		printf("Failed to transmit packets with the device.\n");
		return 1;
	}

	return 0;
}

static const MacAddress *uefinet_get_mac(NetDevice *net_dev)
{
	UefiNetDevice *dev = container_of(net_dev, UefiNetDevice, dev);

	if (uefinet_initialize(dev))
		return NULL;

	return &dev->mac;
}


static void uefi_net_create(EFI_SIMPLE_NETWORK_PROTOCOL *snp)
{
	UefiNetDevice *net = xzalloc(sizeof(*net));

	net->dev.ready = &uefinet_ready;
	net->dev.recv = &uefinet_recv;
	net->dev.send = &uefinet_send;
	net->dev.get_mac = &uefinet_get_mac;

	net->snp = snp;

	list_insert_after(&net->list_node, &uefinet_devices);
	net_add_device(&net->dev);
}

static void uefi_net_destroy(UefiNetDevice *net)
{
	net_remove_device(&net->dev);
	list_remove(&net->list_node);
	free(net);
}

static void uefi_net_poll(NetPoller *poller)
{
	// Destroy all the existing devices and recreate them. This is the
	// easiest way to ensure stale devices go away and new ones are found.
	UefiNetDevice *prev_dev = NULL;
	UefiNetDevice *net_dev;
	list_for_each(net_dev, uefinet_devices, list_node) {
		if (prev_dev)
			uefi_net_destroy(prev_dev);
		prev_dev = net_dev;
	}
	if (prev_dev)
		uefi_net_destroy(prev_dev);

	EFI_SYSTEM_TABLE *sys = uefi_system_table_ptr();
	assert(sys);
	EFI_BOOT_SERVICES *bs = sys->BootServices;

	UINTN buf_size = 0;
	EFI_HANDLE dummy_handle;
	EFI_STATUS status = bs->LocateHandle(
		ByProtocol, &uefi_device_path_to_text_protocol_guid, NULL,
		&buf_size, &dummy_handle);
	die_if(status != EFI_BUFFER_TOO_SMALL,
	       "Error retrieving device path to text protocol handles.\n");

	EFI_HANDLE *dp2text_handles = xmalloc(buf_size);

	status = bs->LocateHandle(
		ByProtocol, &uefi_device_path_to_text_protocol_guid, NULL,
		&buf_size, dp2text_handles);
	die_if(status != EFI_SUCCESS,
	       "Error retrieving device path to text protocol handles.\n");

	EFI_DEVICE_PATH_TO_TEXT_PROTOCOL *dp2text;
	bs->HandleProtocol(dp2text_handles[0],
		&uefi_device_path_to_text_protocol_guid, (void **)&dp2text);
	die_if(status != EFI_SUCCESS,
	       "Failed to get device path to text protocol.\n");
	free(dp2text_handles);


	buf_size = 0;
	status = bs->LocateHandle(
		ByProtocol, &uefi_simple_network_protocol_guid, NULL,
		&buf_size, &dummy_handle);

	// If no network devices were found, there's nothing more to do.
	if (status == EFI_NOT_FOUND)
		return;
	die_if(status != EFI_BUFFER_TOO_SMALL,
	       "Error retrieving network protocol handles.\n");

	EFI_HANDLE *handles = xmalloc(buf_size);
	int handle_count = buf_size / sizeof(dummy_handle);

	status = bs->LocateHandle(
		ByProtocol, &uefi_simple_network_protocol_guid, NULL,
		&buf_size, handles);
	die_if(status != EFI_SUCCESS,
	       "Failed to retrieve network protocol handles.\n");

	uint16_t *last_path = NULL;
	for (int i = 0; i < handle_count; i++) {
		EFI_SIMPLE_NETWORK_PROTOCOL *snp;
		status = bs->HandleProtocol(handles[i],
			&uefi_simple_network_protocol_guid, (void **)&snp);
		die_if(status != EFI_SUCCESS,
		       "Failed to get network protocol.\n");

		EFI_DEVICE_PATH_PROTOCOL *dev_path;
		status = bs->HandleProtocol(handles[i],
			&uefi_device_path_protocol_guid, (void **)&dev_path);
		if (status != EFI_SUCCESS) {
			printf("Failed to get device path.\n");
			continue;
		}

		/*
		 * If the device path of this device is a subpath of the one
		 * we've just seen, ignore it. Apparently certain devices
		 * (e1000 was given as an example) have additional device
		 * paths for their ipv4 and ipv6 configuration information.
		 */
		uint16_t *this_path = dp2text->ConvertDevicePathToText(
			dev_path, 0, 0);
		if (!this_path) {
			printf("Failed to convert device path to text.\n");
			continue;
		}

		if (last_path) {
			int matches = 1;
			for (size_t j = 0;
			     matches && last_path[j] && this_path[j]; j++) {
				if (last_path[j] != this_path[j])
					matches = 0;
			}

			if (matches)
				continue;
		}

		if (last_path)
			bs->FreePool(last_path);
		last_path = this_path;
		uefi_net_create(snp);
	}

	if (last_path)
		bs->FreePool(last_path);
}

static NetPoller uefinet_poller = {
	.poll = &uefi_net_poll,
};

static int uefinet_install_poller(void)
{
	list_insert_after(&uefinet_poller.list_node, &net_pollers);
	return 0;
}

INIT_FUNC(uefinet_install_poller);