/* NetHalt - NetHalt tray program
 * Copyright (C) 2008 Daniel Collins <solemnwarning@solemnwarning.net>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *	* Redistributions of source code must retain the above copyright
 *	  notice, this list of conditions and the following disclaimer.
 *
 *	* Redistributions in binary form must reproduce the above copyright
 *	  notice, this list of conditions and the following disclaimer in the
 *	  documentation and/or other materials provided with the distribution.
 *
 *	* Neither the name of the author nor the names of its contributors may
 *	  be used to endorse or promote products derived from this software
 *	  without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <windows.h>
#include <winsock2.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <stdint.h>

#include "nhtray.h"
#include "lib.h"
#include "gui.h"

#define TICON_UID 1
#define APPWM_TICON (WM_APP+1)

struct packet {
	int64_t sdtime;
	int32_t warning;
	int32_t abort;
	int32_t delay;
} __attribute__((__packed__));

static NOTIFYICONDATA ticon;
static HKEY reghandle = NULL;
static int sockfd = -1;
static HICON redpower16, redpower32;
static struct packet cfg = {0,0,0,0};

static HWND mwnd = NULL;	/* Main window */
static HWND wdlg = NULL;	/* Shutdown warning dialog */
static HWND adlg = NULL;	/* About dialog */

static LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam);
static void tray_menu(HWND hwnd);
static int ipc_connect(void);
static void ipc_read(void *buf, int size, unsigned long nblock);
static void ipc_close(void);
static void ipc_command(char cmd);
static void show_about(void);
static INT_PTR CALLBACK DialogProc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam);
static void show_warning(int show);

int main(int argc, char **argv) {
	WNDCLASSEX wclass;
	MSG msg;
	DWORD dwval;
	WSADATA wsinfo;
	
	dwval = WSAStartup(MAKEWORD(2,2), &wsinfo);
	if(dwval) {
		die("Failed to initialize winsock: %s", w32_error(dwval));
	}
	
	dwval = RegOpenKeyEx(
		HKEY_LOCAL_MACHINE,
		"SOFTWARE\\NetHalt",
		0,
		KEY_READ | KEY_WRITE,
		&reghandle
	);
	if(dwval != ERROR_SUCCESS) {
		die("Failed to open the HKLM\\SOFTWARE\\NetHalt key: %s", w32_error(dwval));
	}
	
	redpower16 = LoadImage(
		GetModuleHandle(NULL),
		MAKEINTRESOURCE(ICO_REDPOWER),
		IMAGE_ICON,
		16, 16,
		LR_DEFAULTCOLOR | LR_LOADTRANSPARENT
	);
	
	redpower32 = LoadImage(
		GetModuleHandle(NULL),
		MAKEINTRESOURCE(ICO_REDPOWER),
		IMAGE_ICON,
		32, 32,
		LR_DEFAULTCOLOR | LR_LOADTRANSPARENT
	);
	
	wclass.cbSize = sizeof(WNDCLASSEX);
	wclass.style = 0;
	wclass.lpfnWndProc = &WndProc;
	wclass.cbClsExtra = 0;
	wclass.cbWndExtra = 0;
	wclass.hInstance = GetModuleHandle(NULL);
	wclass.hIcon = NULL;
	wclass.hCursor = LoadCursor(NULL, IDC_ARROW);
	wclass.hbrBackground = NULL;
	wclass.lpszMenuName = NULL;
	wclass.lpszClassName = "nhtray";
	wclass.hIconSm = NULL;
	
	if(!RegisterClassEx(&wclass)) {
		die(
			"Failed to register window class: %s",
			w32_error(GetLastError())
		);
	}
	
	mwnd = CreateWindowEx(
		0,
		"nhtray",
		"NetHalt",
		WS_CAPTION | WS_SYSMENU,
		200, 200,
		450, 215,
		NULL,
		NULL,
		GetModuleHandle(NULL),
		NULL
	);
	if(!mwnd) {
		die(
			"Failed to create window: %s",
			w32_error(GetLastError())
		);
	}
	
	ticon.cbSize = sizeof(ticon);
	ticon.uID = TICON_UID;
	ticon.uFlags = NIF_MESSAGE | NIF_ICON | NIF_TIP;
	ticon.uCallbackMessage = APPWM_TICON;
	ticon.hIcon = redpower16;
	strcpy(ticon.szTip, "Not connected");
	ticon.uVersion = NOTIFYICON_VERSION;
	ticon.hWnd = mwnd;
	
	Shell_NotifyIcon(NIM_ADD, &ticon);
	
	while(GetMessage(&msg, NULL, 0, 0) > 0) {
		if(wdlg && IsDialogMessage(wdlg, &msg)) {
			continue;
		}
		if(adlg && IsDialogMessage(adlg, &msg)) {
			continue;
		}
		
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}
	
	ipc_close();
	RegCloseKey(reghandle);
	
	return 0;
}

static LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
	int n, remain;
	
	if(msg == WM_CREATE) {
		SetTimer(hwnd, TIM_TIMER, 100, NULL);
		SendMessage(hwnd, WM_TIMER, TIM_TIMER, GetTickCount());
	}
	if(msg == WM_CLOSE) {
		return 0;
	}
	if(msg == APPWM_TICON) {
		if(lParam == WM_LBUTTONDBLCLK) {
			show_about();
		}
		if(lParam == WM_RBUTTONUP) {
			tray_menu(hwnd);
		}
		
		return 0;
	}
	if(msg == WM_COMMAND) {
		if(wParam == MNU_ABOUT) {
			show_about();
		}
		if(wParam == MNU_EXIT) {
			n = msgboxf(
				MB_YESNO | MB_ICONWARNING | MB_DEFBUTTON2,
				"NetHalt",
				"If you exit NetHalt you will not be informed of impending shutdowns.\n"
				"Are you sure you want to exit?"
			);
			
			if(n == IDNO) {
				return 0;
			}
			
			Shell_NotifyIcon(NIM_DELETE, &ticon);
			DestroyWindow(hwnd);
		}
		
		return 0;
	}
	if(msg == WM_TIMER) {
		KillTimer(hwnd, TIM_TIMER);
		
		ipc_connect();
		ipc_read(&cfg, sizeof(cfg), 1);
		
		if(sockfd == -1) {
			strcpy(ticon.szTip, "Not connected");
			show_warning(0);
		}else if(!cfg.sdtime) {
			strcpy(ticon.szTip, "No shutdown scheduled");
			show_warning(0);
		}else{
			remain = (cfg.sdtime - time(NULL));
			
			sprintf(ticon.szTip, "Time until shutdown: %s", fmt_time(remain));
			
			if(cfg.warning && remain <= cfg.warning) {
				show_warning(1);
			}else{
				show_warning(0);
			}
		}
		
		Shell_NotifyIcon(NIM_MODIFY, &ticon);
		SetTimer(hwnd, TIM_TIMER, 100, NULL);
	}
	if(msg == WM_DESTROY) {
		PostQuitMessage(0);
	}
	
	return DefWindowProc(hwnd, msg, wParam, lParam);
}

/* Display error in message box and exit */
void die_r(char const *file, int line, char const *fmt, ...) {
	char msg[1024];
	va_list argv;
	
	va_start(argv, fmt);
	vsnprintf(msg, 1024, fmt, argv);
	va_end(argv);
	
	msgboxf(MB_OK | MB_ICONERROR, "NetHalt - Error", "%s", msg);
	
	exit(1);
}

/* Display tray icon menu */
static void tray_menu(HWND hwnd) {
	HMENU menu;
	POINT cursor_pos;
	
	GetCursorPos(&cursor_pos);
	
	SetForegroundWindow(hwnd);
	
	menu = CreatePopupMenu();
	
	InsertMenu(menu, -1, MF_BYPOSITION | MF_STRING, MNU_ABOUT, "About NetHalt");
	InsertMenu(menu, -1, MF_BYPOSITION | MF_STRING, MNU_EXIT, "Exit");
	
	SetMenuDefaultItem(menu, MNU_ABOUT, FALSE);
	
	SetFocus(hwnd);
	
	TrackPopupMenu(
		menu,
		TPM_LEFTALIGN | TPM_BOTTOMALIGN,
		cursor_pos.x, cursor_pos.y,
		0,
		hwnd,
		NULL
	);
	
	DestroyMenu(menu);
}

/* Connect to the nhservice IPC socket and read the on-connect values
 *
 * Returns 1 if the connection was successful, or already open
 * Returns 0 on failure
*/
static int ipc_connect(void) {
	DWORD dwval, psize, port;
	struct sockaddr_in caddr;
	
	if(sockfd >= 0) {
		return 1;
	}
	
	psize = sizeof(port);
	dwval = RegQueryValueEx(
		reghandle,
		"ipc_port",
		NULL,
		NULL,
		(BYTE*)&port,
		&psize
	);
	if(dwval == ERROR_FILE_NOT_FOUND) {
		return 0;
	}
	if(dwval != ERROR_SUCCESS) {
		die("Failed to read ipc_port value: %s", w32_error(dwval));
	}
	
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if(sockfd == -1) {
		die("Failed to create IPC socket: %s", w32_error(WSAGetLastError()));
	}
	
	struct linger lopt;
	lopt.l_onoff = 0;
	setsockopt(sockfd, SOL_SOCKET, SO_LINGER, (const char*)&lopt, sizeof(lopt));
	
	caddr.sin_family = AF_INET;
	caddr.sin_addr.s_addr = inet_addr("127.0.0.1");
	caddr.sin_port = htons(port);
	
	if(connect(sockfd, (struct sockaddr*)&caddr, sizeof(caddr)) == -1) {
		closesocket(sockfd);
		sockfd = -1;
		
		return 0;
	}
	
	ipc_read(&cfg, sizeof(cfg), 0);
	
	return 1;
}

/* Read data from the IPC socket
 *
 * Loops until size bytes have been read into buf from the IPC socket, if
 * nblock is nonzero, the loop will abort if no data is available on the first
 * recv() call.
 *
 * Closes the IPC socket if recv() returns zero or error.
*/
static void ipc_read(void *buf, int size, unsigned long nblock) {
	int rbytes = 0, rval;
	DWORD dwval;
	
	if(sockfd == -1) {
		return;
	}
	
	ioctlsocket(sockfd, FIONBIO, &nblock);
	
	while(rbytes < size) {
		rval = recv(sockfd, buf+rbytes, size-rbytes, 0);
		
		if(rval == -1) {
			dwval = WSAGetLastError();
			
			if(dwval == WSAEWOULDBLOCK) {
				break;
			}
			if(
				dwval == WSAECONNABORTED ||
				dwval == WSAETIMEDOUT ||
				dwval == WSAECONNRESET
			) {
				ipc_close();
				return;
			}
			
			/* TODO: Better error handling */
			
			die("Failed to read IPC socket: %s", w32_error(dwval));
		}
		if(rval == 0) {
			ipc_close();
			return;
		}
		
		rbytes += rval;
		
		nblock = 0;
		ioctlsocket(sockfd, FIONBIO, &nblock);
	}
}

/* Close the IPC socket */
static void ipc_close(void) {
	if(sockfd == -1) {
		return;
	}
	
	closesocket(sockfd);
	sockfd = -1;
}

/* Send a command to nhservice over the IPC socket */
static void ipc_command(char cmd) {
	int rval = 0;
	
	if(sockfd == -1) {
		return;
	}
	
	while(rval != 1) {
		rval = send(sockfd, &cmd, 1, 0);
		
		if(rval == -1) {
			die(
				"Failed to write to IPC socket: %s",
				w32_error(WSAGetLastError())
			);
		}
	}
}

/* Show the about dialog */
static void show_about(void) {
	if(!adlg) {
		adlg = CreateDialog(
			GetModuleHandle(NULL),
			MAKEINTRESOURCE(DLG_ABOUT),
			NULL,
			&DialogProc
		);
	}
	
	SetForegroundWindow(adlg);
}

/* Dialog event handler */
static INT_PTR CALLBACK DialogProc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam) {
	if(msg == WM_INITDIALOG) {
		SendMessage(hwnd, WM_SETICON, ICON_SMALL, (LPARAM)redpower16);
		SendMessage(hwnd, WM_SETICON, ICON_BIG, (LPARAM)redpower32);
		
		return TRUE;
	}
	if(msg == WM_CLOSE) {
		if(hwnd == adlg) {
			DestroyWindow(hwnd);
		}
		
		return TRUE;
	}
	if(msg == WM_DESTROY) {
		if(hwnd == wdlg) {
			wdlg = NULL;
		}
		if(hwnd == adlg) {
			adlg = NULL;
		}
		
		return TRUE;
	}
	if(msg == WM_COMMAND) {
		if(wparam == BTN_DELAY) {
			ipc_command('D');
		}
		if(wparam == BTN_ABORT) {
			ipc_command('A');
		}
		if(wparam == BTN_OK) {
			DestroyWindow(hwnd);
		}
		if(wparam == IDCANCEL && hwnd == adlg) {
			DestroyWindow(hwnd);
		}
		
		return TRUE;
	}
	
	return FALSE;
}

/* Show, update or hide the warning dialog */
static void show_warning(int show) {
	if(!show) {
		if(wdlg) {
			DestroyWindow(wdlg);
		}
		
		return;
	}
	
	if(!wdlg) {
		wdlg = CreateDialog(
			GetModuleHandle(NULL),
			MAKEINTRESOURCE(DLG_WARNING),
			NULL,
			&DialogProc
		);
		
		/* Grey out the close button */
		EnableMenuItem(
			GetSystemMenu(wdlg, FALSE),
			SC_CLOSE,
			MF_BYCOMMAND | MF_GRAYED
		);
	}
	
	gui_stext(wdlg, TXT_TIMER, "%s", fmt_time(cfg.sdtime - time(NULL)));
	gui_enable(wdlg, BTN_DELAY, cfg.delay);
	gui_enable(wdlg, BTN_ABORT, cfg.abort);
}