#include <stdlib.h>
#include <stdarg.h>
#include <stdio.h>
#include <assert.h>
#include <Windows.h>
#include <CommCtrl.h>
#include <WinSock2.h>
#include <dns_sd.h>

#include "resource.h"
#include "list.h"
#include "mounter.h"

#define DIALOG_CAPTIONS "Remote optical disc"
#define DRIVE_LETTER 'O'
#define EVENT_SOCKET_ACTIVITY (WM_APP+0)

typedef struct _DNSSocket {
	struct _DNSSocket *next;
	SOCKET sock;
	DNSServiceRef ref;
} DNSSocket;

typedef struct _Disc {
	struct _Disc *next;
	struct _Host *host;

	char *name;
	int list_index;
} Disc;

typedef struct _Host {
	struct _Host *next;

	DNSServiceRef resolve_ref;
	DNSServiceRef query_ref;

	char *name;
	USHORT port;

	struct _Disc discs;
} Host;

static HWND dialog_hwnd = NULL;
static DNSServiceRef search_ref = NULL;
static DNSSocket *sock_list = NULL;
static Host *host_list = NULL;

static void __cdecl odprintf(const char *format, ...)
{
	char    buf[4096], *p = buf;
	va_list args;
	int     n;

	va_start(args, format);
	n = _vsnprintf(p, sizeof buf - 3, format, args); // buf-3 is room for CR/LF/NUL
	va_end(args);

	p += (n < 0) ? sizeof buf - 3 : n;

	while ( p > buf  &&  isspace(p[-1]) )
		*--p = '\0';

	*p++ = '\r';
	*p++ = '\n';
	*p   = '\0';

	OutputDebugString(buf);
}

static void dnssocket_monitor(DNSServiceRef ref)
{
	DNSSocket *d = (DNSSocket*) malloc(sizeof(DNSSocket));
	int err;
	d->ref = ref;
	d->sock = (SOCKET) DNSServiceRefSockFD(ref);

	assert(d->ref);

	err = WSAAsyncSelect(d->sock, dialog_hwnd, EVENT_SOCKET_ACTIVITY, FD_READ | FD_CLOSE);
	if (err) {
		MessageBox(dialog_hwnd, "Failed to start WSAAsyncSelect", DIALOG_CAPTIONS, MB_ICONWARNING | MB_OK);
		free(d);
		return;
	}

	list_add(DNSSocket, sock_list, d);
}

static void dnssocket_close(DNSServiceRef ref)
{
	DNSSocket **p;
	for (p = &sock_list; *p; p = &((*p)->next)) {
		if ((*p)->ref == ref) {
			DNSSocket *d = *p;
			*p = d->next;
			DNSServiceRefDeallocate(d->ref);
			free(d);
			return;
		}
	}
}

static void dnssocket_fire(SOCKET sock)
{
	DNSSocket *d;
	list_foreach(DNSSocket, sock_list, d) {
		if (d->sock == sock) {
			DNSServiceProcessResult(d->ref);
			return;
		}
	}
}

static Disc *disc_alloc()
{
	return (Disc*) calloc(1, sizeof(Disc));
}

static void disc_free(Disc *disc)
{
	if (disc->name) {
		free((void*) disc->name);
	}
	free(disc);
}

static Host *host_alloc()
{
	return (Host*) calloc(1, sizeof(Host));
}

static void host_free(Host *host)
{
	if (host->resolve_ref) {
		dnssocket_close(host->resolve_ref);
	}
	if (host->query_ref) {
		dnssocket_close(host->query_ref);
	}
	if (host->name) {
		free((void*) host->name);
	}
	free(host);
}

static void center_dialog(HWND hwndDlg)
{
	HWND desktop = GetDesktopWindow();
	RECT rc, rcDlg, rcOwner;

	GetWindowRect(desktop, &rcOwner); 
    GetWindowRect(hwndDlg, &rcDlg); 
    CopyRect(&rc, &rcOwner); 

    OffsetRect(&rcDlg, -rcDlg.left, -rcDlg.top); 
    OffsetRect(&rc, -rc.left, -rc.top); 
    OffsetRect(&rc, -rcDlg.right, -rcDlg.bottom); 

    SetWindowPos(hwndDlg, 
                 HWND_TOP, 
                 rcOwner.left + (rc.right / 2), 
                 rcOwner.top + (rc.bottom / 2), 
                 0, 0,          // Ignores size arguments. 
                 SWP_NOSIZE);	
}

static int disc_connect(Disc *disc)
{
	static char url[1024];
	_snprintf(url, sizeof(url), "http://%s:%d/%s.dmg", disc->host->name, disc->host->port, disc->name);
	odprintf("Connecting to URL %s\n", url);
	return HttpDiskMountUrl(0, url, DRIVE_LETTER, TRUE);
}

static void CALLBACK query_reply(DNSServiceRef sdRef, DNSServiceFlags flags, uint32_t interfaceIndex, DNSServiceErrorType errorCode, const char *fullname, uint16_t rrtype, uint16_t rrclass, uint16_t rdlen, const void *rdata, uint32_t ttl, void *context)
{
	Host *host = (Host*) context;
	DNSServiceErrorType err;

	if (errorCode == kDNSServiceErr_NoError) {
		if (flags & kDNSServiceFlagsAdd) {
			int count = TXTRecordGetCount(rdlen, rdata);
			int i;
			for (i = 0; i < count; i++) {
				char *key = (char*)malloc(256);
				uint8_t value_len;
				void *value;
				err = TXTRecordGetItemAtIndex(rdlen, rdata, i, 256, key, &value_len, &value);
				if (err != kDNSServiceErr_NoError) continue;
				if (strcmp(key, "sys") == 0) {
					// Sys record, ignore for now
					free(key);
				} else {
					Disc *disc = disc_alloc();
					disc->host = host;
					disc->name = key;

					disc->list_index = SendDlgItemMessage(dialog_hwnd, IDC_LIST, LB_ADDSTRING, 0, (LPARAM) key);
					SendDlgItemMessage(dialog_hwnd, IDC_LIST, LB_SETITEMDATA, (WPARAM) disc->list_index, (LPARAM) disc);
				}
			}
		}
	} else {
		MessageBox(NULL, "Bonjour query error occurred", DIALOG_CAPTIONS, MB_ICONWARNING | MB_OK);
	}
}

static void CALLBACK resolve_reply(DNSServiceRef sdRef, DNSServiceFlags flags, uint32_t interfaceIndex, DNSServiceErrorType errorCode, const char *fullname, const char *hosttarget, uint16_t port, uint16_t txtLen, const unsigned char *txtRecord, void *context)
{
	Host *host = (Host*)context;
	DNSServiceErrorType err;

	if (errorCode == kDNSServiceErr_NoError) {
		odprintf("Resolved %s\n", fullname);

		dnssocket_close(host->resolve_ref);
		host->resolve_ref = NULL;

		host->name = _strdup(hosttarget);
		host->port = ntohs(port);

		err = DNSServiceQueryRecord(&host->query_ref, 0, 0, fullname, kDNSServiceType_TXT, kDNSServiceClass_IN, query_reply, host);
		if (err != kDNSServiceErr_NoError) {
			MessageBox(dialog_hwnd, "Failed to start DNSServiceQueryRecord", DIALOG_CAPTIONS, MB_ICONWARNING | MB_OK);
			return;
		}

		dnssocket_monitor(host->query_ref);
	} else {
		MessageBox(NULL, "Bonjour resolver error occurred", DIALOG_CAPTIONS, MB_ICONWARNING | MB_OK);
	}
}

static void start_resolve(const char *service, const char *regtype, const char *domain)
{
	Host *host = host_alloc();
	DNSServiceErrorType err;

	odprintf("Resolving %s\n", service);

	err = DNSServiceResolve(&host->resolve_ref, 0, 0, service, regtype, domain, resolve_reply, host);
	if (err != kDNSServiceErr_NoError) {
		MessageBox(dialog_hwnd, "Failed to start DNSServiceResolve", DIALOG_CAPTIONS, MB_ICONWARNING | MB_OK);
		host_free(host);
		return;
	}

	dnssocket_monitor(host->resolve_ref);

	list_add(Host, host_list, host);
}

static void CALLBACK search_reply(DNSServiceRef sdRef, DNSServiceFlags flags, uint32_t interfaceIndex, DNSServiceErrorType errorCode, const char *serviceName, const char *regtype, const char *replyDomain, void *context)
{
	HWND hwndDlg = (HWND) context;

	if (errorCode == kDNSServiceErr_NoError) {
		if (flags & kDNSServiceFlagsAdd) {
			start_resolve(serviceName, regtype, replyDomain);
		}
	} else {
		MessageBox(hwndDlg, "Bonjour browser error occurred", DIALOG_CAPTIONS, MB_ICONWARNING | MB_OK);
	}
}

static void start_search()
{
	DNSServiceErrorType err;

	err = DNSServiceBrowse(&search_ref, 0, 0, "_odisk._tcp", NULL, search_reply, NULL);
	if (err != kDNSServiceErr_NoError) {
		MessageBox(dialog_hwnd, "Failed to start DNSServiceBrowse", DIALOG_CAPTIONS, MB_ICONWARNING | MB_OK);
		return;
	}
	
	dnssocket_monitor(search_ref);
}

static void stop_search()
{
	dnssocket_close(search_ref);
}

static void handle_ok()
{
	Disc *disc;
	int sel = SendDlgItemMessage(dialog_hwnd, IDC_LIST, LB_GETCURSEL, 0, 0);
	if (sel == LB_ERR) {
		MessageBox(dialog_hwnd, "Please select a disc", DIALOG_CAPTIONS, MB_ICONINFORMATION | MB_OK);
		return;
	}

	disc = (Disc*) SendDlgItemMessage(dialog_hwnd, IDC_LIST, LB_GETITEMDATA, (WPARAM) sel, 0);
	assert(disc);
	if (disc_connect(disc) == 0) {
		EndDialog(dialog_hwnd, IDOK);
	}
}

static void handle_disconnect()
{
	HttpDiskUmount(DRIVE_LETTER);
}

BOOL CALLBACK ConnectDlgProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
	switch (msg)
	{
	case WM_INITDIALOG:
		dialog_hwnd = hwnd;
		center_dialog(hwnd);
		start_search();
		return TRUE;
	case WM_COMMAND:
		switch (LOWORD(wParam)) {
		case IDOK:
			handle_ok();
			return TRUE;
		case IDCANCEL:
			EndDialog(hwnd, IDCANCEL);
			return TRUE;
		case IDC_DISCONNECT:
			handle_disconnect();
			return TRUE;
		default:
			return FALSE;
		}
	case EVENT_SOCKET_ACTIVITY:
		dnssocket_fire((SOCKET) wParam);
		return TRUE;
	default:
		return FALSE;
	}
}

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow)
{
	INITCOMMONCONTROLSEX icex;
	icex.dwSize = sizeof(icex);
	icex.dwICC = ICC_STANDARD_CLASSES;
	InitCommonControlsEx(&icex);

	DialogBox(GetModuleHandle(NULL), MAKEINTRESOURCE(IDD_CONNECT), NULL, ConnectDlgProc);

	return EXIT_SUCCESS;
}