#include <stdlib.h>

#include <glib.h>
#include <glib-object.h>

#include <avahi-gobject/ga-client.h>
#include <avahi-gobject/ga-entry-group.h>

#include <libsoup/soup.h>

#include "udisks.h"
#include "udisks-device.h"

#define RODISC_MDNS_SERVICE "_odisk._tcp"

#define RODISC_IMAGE_MIME "application/octet-stream"

#define RODISC_TYPE_GENERIC "public.optical-storage-media"
#define RODISC_TYPE_CD "public.cd-media"
#define RODISC_TYPE_DVD "public.dvd-media"
#define RODISC_TYPE_BD "public.optical-storage-media"

/** Buffer size for a single read() call.
 * Any request larger than this will be served in chunks. */
#define MAX_BUFFER_SIZE (32*1024*1024UL)

static GMainLoop *main_loop;

static SoupServer *server;

static OrgFreedesktopUDisks *monitor;

static GaClient *mdns_client;
static GaEntryGroup *mdns_group;
static GaEntryGroupService *mdns_service;

static GHashTable *discs;
static guint disc_change_count = 0;

static gchar **files = NULL;

static GOptionEntry entries[] =
{
    { "file", 'f', 0, G_OPTION_ARG_FILENAME_ARRAY, &files, "Image files to export as remote discs", "FILE"},
    { NULL }
};

typedef struct {
	gchar *id;
	gchar *uri;
	gchar *file_path;
	GFile *file;
	gchar *label;
	const gchar *type;
	guint64 size;
} RODisc;

typedef struct {
	RODisc *disc;
	gsize start;
	gsize length;
	gpointer buffer;
	SoupMessage *msg;
	GInputStream *stream;
	GCancellable *cancel;
} RODiscReadOp;

static void mdns_service_freeze()
{
	g_return_if_fail(mdns_service);
	ga_entry_group_service_freeze(mdns_service);
}

static void mdns_service_thaw()
{
	GError *error = NULL;
	g_return_if_fail(mdns_service);
	if (!ga_entry_group_service_thaw(mdns_service, &error)) {
		g_warning("Could not update service TXT entries: %s", error->message);
		g_error_free(error);
	}
}

static void mdns_service_update()
{
	GError *error = NULL;

	g_return_if_fail(mdns_service);

	// "sys=waMA=00:00:00:00:00:00,adVF=0x4,adDT=0x2,adCC=0"
	// waMA = MAC address
	// adVF = Volume flags (0x200 "ask me first")
	// adDT = Supported media?
	// adCC = Disc change count?

	guint flags = 0;
	guint media = 2;
	gchar *record =
	        g_strdup_printf("waMA=00:00:00:00:00:00,adVF=0x%x,adDT=0x%x,adCC=%u",
	                        flags, media, disc_change_count);

	if (!ga_entry_group_service_set(mdns_service, "sys", record, &error))
	{
		g_warning("Could not update main TXT record: %s", error->message);
		g_error_free(error);
	}

	g_free(record);
}

static void mdns_service_update_disc(RODisc *disc)
{
	GError *error = NULL;

	g_return_if_fail(mdns_service);

	// "CdRom0=adVN=DiscLabel,adVT=public.cd-media"
	gchar *record = g_strdup_printf("adVN=%s,adVT=%s", disc->label, disc->type);

	const gchar *uri_basename = &disc->uri[1]; // Skip first '/'
	if (!ga_entry_group_service_set(mdns_service, uri_basename, record, &error))
	{
		g_warning("Could not update TXT record for disc at '%s': %s",
		          disc->uri, error->message);
		g_error_free(error);
	}

	g_free(record);
}

static void mdns_service_remove_disc(RODisc *disc)
{
	GError *error = NULL;
	g_return_if_fail(mdns_service);
	const gchar *uri_basename = &disc->uri[1];
	if (!ga_entry_group_service_remove_key(mdns_service, uri_basename, &error)) {
		g_warning("Could not update TXT record for disc at '%s': %s",
		          disc->uri, error->message);
		g_error_free(error);
	}
}

static void mdns_service_update_discs_func(gpointer key, gpointer value, gpointer user_data)
{
	RODisc *disc = (RODisc*) value;
	mdns_service_update_disc(disc);
}

static void mdns_service_update_discs()
{
	g_hash_table_foreach(discs, mdns_service_update_discs_func, NULL);
}

static void mdns_register_service()
{
	GError * error = NULL;
	if (!mdns_group) {
		mdns_group = ga_entry_group_new();

		if (!ga_entry_group_attach(mdns_group, mdns_client, &error)) {
			g_warning("Could not attach MDNS group to client: %s", error->message);
			g_error_free(error);
			return;
		}
	}

	const gchar *name = avahi_client_get_host_name(mdns_client->avahi_client);
	guint port = soup_server_get_port(server);
	mdns_service = ga_entry_group_add_service(mdns_group,
	                                          name, RODISC_MDNS_SERVICE,
	                                          port, &error,
	                                          NULL);
	if (!mdns_service) {
		g_warning("Could not create service: %s", error->message);
		g_error_free(error);
		return;
	}

	// Create TXT records, disc records, etc.
	mdns_service_update();
	mdns_service_update_discs();

	if (!ga_entry_group_commit(mdns_group, &error)) {
		g_warning("Could not announce MDNS service: %s", error->message);
		g_error_free(error);
		return;
	}
}

static void mdns_client_state_changed_cb(GaClient *client, GaClientState state, gpointer user_data)
{
	switch (state) {
	case GA_CLIENT_STATE_FAILURE:
		g_warning("MDNS client state failure");
		break;
	case GA_CLIENT_STATE_S_RUNNING:
		g_debug("MDNS client found server running");
		mdns_register_service();
		break;
	case GA_CLIENT_STATE_S_COLLISION:
	case GA_CLIENT_STATE_S_REGISTERING:
		g_message("MDNS collision");
		if (mdns_group) {
			ga_entry_group_reset(mdns_group, NULL);
			mdns_service = 0;
		}
		break;
	default:
		// Do nothing
		break;
	}
}

static void server_disc_cb(SoupServer *server, SoupMessage *msg, const char *path,
						   GHashTable *query, SoupClientContext *client, gpointer user_data);

static inline RODiscReadOp* server_disc_read_op_new()
{
	RODiscReadOp *op = g_slice_new(RODiscReadOp);
	return op;
}

static void server_disc_read_op_free(RODiscReadOp *op)
{
	if (op->stream) g_object_unref(op->stream);
	if (op->buffer) g_free(op->buffer);
	if (op->cancel) g_object_unref(op->cancel);
	g_slice_free(RODiscReadOp, op);
}

static void server_disc_finished_cb(SoupMessage *msg, gpointer user_data)
{
	RODiscReadOp *op = (RODiscReadOp*) user_data;
	g_warn_if_fail(op->msg == msg);

	g_debug(" request finished");

	if (op->cancel) {
		// Cancel pending operation if any
		g_cancellable_cancel(op->cancel);
		g_object_unref(op->cancel);
		op->cancel = NULL;
	}

	server_disc_read_op_free(op);
}

static void server_disc_read_cb(GInputStream *stream, GAsyncResult *res, gpointer user_data)
{
	RODiscReadOp *op = (RODiscReadOp*) user_data;
	GError *error = NULL;
	gssize size = g_input_stream_read_finish(stream, res, &error);
	if (size == -1) {
		if (g_error_matches(error, G_IO_ERROR, G_IO_ERROR_CANCELLED)) {
			g_debug(" request canceled");
			return;
		}

		g_warning("Failed to read from %s: %s",
		          op->disc->file_path, error->message);
		soup_message_set_status(op->msg, SOUP_STATUS_INTERNAL_SERVER_ERROR);
		soup_server_unpause_message(server, op->msg);
		g_error_free(error);
		server_disc_read_op_free(op);
		return;
	}

	g_debug(" read %ld bytes of %lu bytes", size, op->length);
	g_warn_if_fail(size > 0);

	if (size >= op->length || size == 0) {
		// We have read everything we needed, so this finishes the request!
		g_debug(" finishing request");
		g_object_unref(op->cancel);
		op->cancel = NULL;
		op->length = 0;

		soup_message_body_append(op->msg->response_body, SOUP_MEMORY_TAKE,
								 op->buffer, size);
		soup_message_body_complete(op->msg->response_body);
		soup_server_unpause_message(server, op->msg);

		op->buffer = NULL; // Passed ownership of buffer to libsoup
		// 'op' will be freed when 'finished' signal @msg is triggered
	} else {
		soup_message_body_append(op->msg->response_body, SOUP_MEMORY_COPY,
								 op->buffer, size);
		soup_server_unpause_message(server, op->msg);

		// Launch the read for the remaining bytes
		op->start += size;
		op->length -= size;

		gsize read_size = MIN(op->length, MAX_BUFFER_SIZE);
		g_debug(" going to read %lu bytes", read_size);
		g_return_if_fail(read_size > 0);

		g_input_stream_read_async(op->stream, op->buffer, read_size,
								  G_PRIORITY_DEFAULT, op->cancel,
								  (GAsyncReadyCallback) server_disc_read_cb, op);
	}
}

static void server_disc_perform_read_range(SoupMessage *msg, RODisc *disc, goffset start, goffset end)
{
	GError *error = NULL;

	g_debug("Opening %s", disc->file_path);

	// File is opened on every read because of asynchronous reads
	GFileInputStream * stream = g_file_read(disc->file, NULL, &error);

	if (!stream) {
		g_warning("Could not open file '%s' for reading: %s",
		          disc->file_path, error->message);
		soup_message_set_status(msg, SOUP_STATUS_INTERNAL_SERVER_ERROR);
		g_error_free(error);
		return;
	}

	gsize size = (end - start) + 1;
	g_debug(" reading range %lu-%lu size %lu", start, end, size);
	g_warn_if_fail(size > 0);

	if (start != 0 &&
			!g_seekable_seek(G_SEEKABLE(stream), start, G_SEEK_SET, NULL, &error)) {
		g_warning("Could not seek in file '%s': %s",
				  disc->file_path, error->message);
		soup_message_set_status(msg, SOUP_STATUS_INTERNAL_SERVER_ERROR);
		g_error_free(error);
		return;
	}

	gsize read_size = MIN(size, MAX_BUFFER_SIZE);
	g_debug(" going to read %lu bytes", read_size);

	RODiscReadOp *op = server_disc_read_op_new();
	op->disc = disc;
	op->start = start;
	op->length = size;
	op->buffer = g_malloc(read_size);
	op->msg = msg;
	op->stream = G_INPUT_STREAM(stream);
	op->cancel = g_cancellable_new();

	g_signal_connect(msg, "finished",
					 G_CALLBACK(server_disc_finished_cb), op);

	g_input_stream_read_async(op->stream, op->buffer, read_size,
							  G_PRIORITY_DEFAULT, op->cancel,
	                          (GAsyncReadyCallback) server_disc_read_cb, op);
}

static void server_disc_perform_read(SoupMessage *msg, RODisc *disc)
{
    server_disc_perform_read_range(msg, disc, 0, disc->size - 1);
}

static void server_disc_cb(SoupServer *server, SoupMessage *msg, const char *path,
                           GHashTable *query, SoupClientContext *client, gpointer user_data)
{

	RODisc *disc = (RODisc*) user_data;
	if (!g_str_has_suffix(path,".dmg")) {
		g_debug("Not found (%s)", path);
		soup_message_set_status(msg, SOUP_STATUS_NOT_FOUND);
		return;
	}
	soup_message_headers_append(msg->response_headers, "Server", "RODisc/1.0");
	if (msg->method == SOUP_METHOD_HEAD) {
		g_debug("Head on %s", path);
		soup_message_set_status(msg, SOUP_STATUS_OK);
		soup_message_headers_set_content_length(msg->response_headers, disc->size);
		soup_message_headers_set_content_type(msg->response_headers, RODISC_IMAGE_MIME, NULL);
		soup_message_headers_replace(msg->response_headers, "Accept-Ranges", "bytes");
	} else if (msg->method == SOUP_METHOD_GET) {
		g_debug("Get on %s", path);
		soup_message_headers_set_content_type(msg->response_headers, RODISC_IMAGE_MIME, NULL);
		soup_message_headers_replace(msg->response_headers, "Accept-Ranges", "bytes");
		soup_message_body_set_accumulate(msg->response_body, FALSE);
		SoupRange *ranges;
		int length;
		if (soup_message_headers_get_ranges(msg->request_headers, disc->size, &ranges, &length)) {
			if (length != 1) {
				g_warning("Multi-range not yet supported");
				soup_message_set_status(msg, SOUP_STATUS_INVALID_RANGE);
				return;
			}
			goffset start = ranges[0].start, end = ranges[0].end;
			soup_message_set_status(msg, SOUP_STATUS_PARTIAL_CONTENT);
			soup_message_headers_set_content_range(msg->response_headers,
												   start, end,
												   disc->size);
			soup_server_pause_message(server, msg);
			server_disc_perform_read_range(msg, disc, start, end);
		} else {
			soup_message_headers_set_content_length(msg->response_headers, disc->size);
			soup_message_set_status(msg, SOUP_STATUS_OK);
			soup_server_pause_message(server, msg);
			server_disc_perform_read(msg, disc);
		}
	} else {
		soup_message_set_status(msg, SOUP_STATUS_NOT_IMPLEMENTED);
		g_warning("Unimplemented method on %s", path);
	}
}

static void server_cb(SoupServer *server, SoupMessage *msg, const char *path,
                      GHashTable *query, SoupClientContext *client, gpointer user_data)
{
	g_message("Unknown path requested: %s", path);
	soup_message_set_status(msg, SOUP_STATUS_NOT_FOUND);
}

static inline RODisc *rodisc_lookup(const gchar *id)
{
	return g_hash_table_lookup(discs, id);
}

static inline RODisc *rodisc_new()
{
	return g_slice_new0(RODisc);
}

static void rodisc_destroy(RODisc *disc)
{
	g_free(disc->id);
	g_free(disc->uri);
	g_object_unref(disc->file);
	g_free(disc->file_path);
	g_free(disc->label);
	g_slice_free(RODisc, disc);
}

static void rodisc_export(RODisc *disc)
{
	g_debug("Exporting %s to %s (volume '%s' type '%s')",
	        disc->file_path, disc->uri, disc->label, disc->type);
	g_hash_table_insert(discs, disc->id, disc);

	soup_server_add_handler(server, disc->uri, server_disc_cb, disc, NULL);

	mdns_service_freeze();
	disc_change_count++;
	mdns_service_update_disc(disc);
	mdns_service_update();
	mdns_service_thaw();
}

static void rodisc_remove(RODisc *disc)
{
	g_debug("Unexporting %s", disc->uri);

	mdns_service_freeze();
	disc_change_count++;
	mdns_service_remove_disc(disc);
	mdns_service_update();
	mdns_service_thaw();

	soup_server_remove_handler(server, disc->uri);
	g_hash_table_remove(discs, disc->id);
}

static void rodisc_refresh(RODisc *disc)
{
	mdns_service_freeze();
	disc_change_count++;
	mdns_service_update_disc(disc);
	mdns_service_update();
	mdns_service_thaw();
}

static void monitor_set_disc_attrs(RODisc *disc, OrgFreedesktopUDisksDevice *device)
{
	const gchar * disc_type = org_freedesktop_udisks_device_get_drive_media(device);

	disc->label = org_freedesktop_udisks_device_dup_id_label(device);
	if (!disc->label || disc->label[0] == '\0') {
		g_free(disc->label);
		disc->label = g_strdup("Disc");
	}

	if (g_str_has_prefix(disc_type, "optical_cd")) {
		disc->type = RODISC_TYPE_CD;
	} else if (g_str_has_prefix(disc_type, "optical_dvd")) {
		disc->type = RODISC_TYPE_DVD;
	} else {
		disc->type = RODISC_TYPE_GENERIC;
	}

	disc->size = org_freedesktop_udisks_device_get_device_size(device);
}

static void monitor_try_add_device(const gchar *device_path)
{
	GError *error = NULL;
	OrgFreedesktopUDisksDevice *device =
	        org_freedesktop_udisks_device_proxy_new_for_bus_sync(
	            G_BUS_TYPE_SYSTEM, G_DBUS_PROXY_FLAGS_NONE,
	            "org.freedesktop.UDisks", device_path,
	            NULL, &error);
	if (!device) {
		g_warning("Could not add UDisk device: %s", error->message);
		g_error_free(error);
		return;
	}
	if (!org_freedesktop_udisks_device_get_device_is_optical_disc(device)) {
		g_object_unref(device);
		return;
	}

	RODisc *disc = rodisc_new();

	const gchar * device_file_path = org_freedesktop_udisks_device_get_device_file(device);
	const gchar * const * id_paths = org_freedesktop_udisks_device_get_device_file_by_id(device);
	disc->id = g_strdup(device_path);
	disc->uri = g_strconcat("/", g_path_get_basename(id_paths[0]), NULL);
	disc->file = g_file_new_for_path(device_file_path);
	disc->file_path = g_file_get_path(disc->file);

	monitor_set_disc_attrs(disc, device);

	rodisc_export(disc);
}

static void monitor_try_update_remove_device(const gchar *device_path, RODisc *disc)
{
	OrgFreedesktopUDisksDevice *device =
	        org_freedesktop_udisks_device_proxy_new_for_bus_sync(
	            G_BUS_TYPE_SYSTEM, G_DBUS_PROXY_FLAGS_NONE,
	            "org.freedesktop.UDisks", device_path,
	            NULL, NULL);
	if (!device) {
		g_debug("Device removal");
		rodisc_remove(disc);
		return;
	}
	if (!org_freedesktop_udisks_device_get_device_is_optical_disc(device)) {
		g_debug("Media removal");
		g_object_unref(device);
		rodisc_remove(disc);
		return;
	}
	if (g_strcasecmp(org_freedesktop_udisks_device_get_id_label(device),
	                 disc->label) != 0) {
		g_debug("Media change");
		g_free(disc->label);
		monitor_set_disc_attrs(disc, device);
		rodisc_refresh(disc);
	}
}

static void monitor_device_added_cb(GObject *source_object, gchar *device,
                                    gpointer user_data)
{
	g_debug("Device added: %s", device);
	RODisc *disc = rodisc_lookup(device);
	if (disc) {
		g_warning("Disk already added: %s", device);
	} else {
		monitor_try_add_device(device);
	}
}

static void monitor_device_changed_cb(GObject *source_object, gchar *device,
                                    gpointer user_data)
{
	g_debug("Device changed: %s", device);
	RODisc *disc = rodisc_lookup(device);
	// Hopefully we detected the eject event, otherwise we'll mess up.
	if (disc) {
		monitor_try_update_remove_device(device, disc);
	} else {
		monitor_try_add_device(device);
	}
}

static void monitor_device_removed_cb(GObject *source_object, gchar *device,
                                    gpointer user_data)
{
	g_debug("Device removed: %s", device);
	RODisc *disc = rodisc_lookup(device);
	if (disc) {
		monitor_try_update_remove_device(device, disc);
	}
}


static void monitor_enumerate_devices_cb(GObject *source_object,
                                         GAsyncResult *res,
                                         gpointer user_data)
{
	GError *error = NULL;
	gchar **devices;
	if (org_freedesktop_udisks_call_enumerate_devices_finish(monitor, &devices, res, &error)) {
		gchar **s;
		for (s = devices; *s; s++) {
			monitor_try_add_device(*s);
		}
		g_strfreev(devices);
	} else {
		g_warning("Could not enumerate devices using UDisks: %s", error->message);
		g_error_free(error);
	}
}

static void file_add_disc(const gchar *path)
{
	static int num = 0;
	GFile *file = g_file_new_for_path(path);
	GFileInfo *info = g_file_query_info(file,
										G_FILE_ATTRIBUTE_STANDARD_DISPLAY_NAME ","
										G_FILE_ATTRIBUTE_STANDARD_SIZE,
	                                    G_FILE_QUERY_INFO_NONE, NULL, NULL);
	g_return_if_fail(info);
	const int my_num = ++num;
	RODisc *disc = rodisc_new();
	disc->id = g_strdup_printf("file%d", my_num);
	disc->uri = g_strdup_printf("/file%d", my_num);
	disc->file_path = g_file_get_path(file);
	disc->file = file;
	disc->label = g_strdup(g_file_info_get_display_name(info));
	disc->type = RODISC_TYPE_GENERIC;
	disc->size = g_file_info_get_size(info);
	rodisc_export(disc);
	g_object_unref(info);
}

static void files_add()
{
	gchar **f;
	if (!files) return;
	for (f = files; *f; f++) {
		file_add_disc(*f);
	}
}

int main(int argc, char * argv[])
{
	GError *error = NULL;
	GOptionContext *context = g_option_context_new("- remote optical disc service");

	g_type_init();
	main_loop = g_main_loop_new(NULL, FALSE);

	discs = g_hash_table_new_full(g_str_hash, g_str_equal,
	                              NULL, (GDestroyNotify) rodisc_destroy);

	g_option_context_add_main_entries(context, entries, NULL);
	if (!g_option_context_parse(context, &argc, &argv, &error)) {
		g_printerr("Option parsing failed: %s\n", error->message);
		return EXIT_FAILURE;
	}

	SoupAddress *addr = soup_address_new_any(SOUP_ADDRESS_FAMILY_IPV4, SOUP_ADDRESS_ANY_PORT);
	server = soup_server_new(SOUP_SERVER_INTERFACE, addr, NULL);
	g_object_unref(addr);
	if (!server) {
		g_warning("Could not create HTTP server");
		return EXIT_FAILURE;
	}

	soup_server_add_handler(server, NULL, server_cb, NULL, NULL);
	soup_server_run_async(server);

	mdns_client = ga_client_new(GA_CLIENT_FLAG_NO_FLAGS);
	g_signal_connect(mdns_client, "state-changed",
	                 G_CALLBACK(mdns_client_state_changed_cb), NULL);
	if (!ga_client_start(mdns_client, &error)) {
		g_warning("Could not start MDNS client");
		g_error_free(error);
		return EXIT_FAILURE;
	}

	monitor = org_freedesktop_udisks_proxy_new_for_bus_sync(
	            G_BUS_TYPE_SYSTEM, G_DBUS_PROXY_FLAGS_NONE,
	            "org.freedesktop.UDisks", "/org/freedesktop/UDisks",
	            NULL, &error);
	if (!monitor) {
		g_warning("Could not create proxy to the UDisks service: %s", error->message);
		g_error_free(error);
		return EXIT_FAILURE;
	}
	g_signal_connect(monitor, "device-added",
	                 G_CALLBACK(monitor_device_added_cb), NULL);
	g_signal_connect(monitor, "device-changed",
	                 G_CALLBACK(monitor_device_changed_cb), NULL);
	g_signal_connect(monitor, "device-removed",
	                 G_CALLBACK(monitor_device_removed_cb), NULL);
	org_freedesktop_udisks_call_enumerate_devices(monitor, NULL,
	                                              monitor_enumerate_devices_cb, NULL);

	files_add();

	g_message("Listening on %d", soup_server_get_port(server));

	g_main_loop_run(main_loop);

	g_hash_table_destroy(discs);
	g_object_unref(monitor);
	g_object_unref(mdns_client);
	g_object_unref(server);
	g_main_loop_unref(main_loop);

	return EXIT_SUCCESS;
}