/* See LICENSE file for copyright and license details. */
#define WINVER 0x0601
#define _WIN32_WINNT 0x0601
#define WIN32_LEAN_AND_MEAN

#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <wchar.h>
#include <windows.h>

#include "config.h"

#define LEN(x) (sizeof(x) / sizeof((x)[0]))

enum {
	WM_SEL = WM_APP + 1
};

struct Sel {
	POINT down;
	POINT up;
	int dragged;
	int clicks;
	int shift;
	int kbd;
};

static char *argv0;
static int verbose;
static HWND hwnd;
static HHOOK mousehook;
static HHOOK kbdhook;
static HANDLE quit;
static struct Sel sel;
static struct Sel pending;

static void die(const char *fmt, ...);
static void usage(void);
static void attachio(int force);
static int samewin(HWND a, HWND b);
static int isclient(HWND hw, POINT pt);
static int ispass(HWND hw);
static int skipctrlc(HWND hw);
static int goodsel(const struct Sel *s);
static void sendctrlc(void);
static void docopy(void);
static void request(int kbd);
static LRESULT CALLBACK mouseproc(int code, WPARAM wp, LPARAM lp);
static LRESULT CALLBACK kbdproc(int code, WPARAM wp, LPARAM lp);
static LRESULT CALLBACK wndproc(HWND hw, UINT msg, WPARAM wp,
		LPARAM lp);
static int runkey(int add);
static void signalquit(void);
static int loop(HINSTANCE inst);

static void
die(const char *fmt, ...)
{
	va_list ap;
	size_t n;

	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	n = strlen(fmt);
	if (n && fmt[n - 1] == ':')
		fprintf(stderr, " %lu\n", GetLastError());
	else
		fputc('\n', stderr);
	exit(1);
}

static void
usage(void)
{
	fprintf(stderr, "usage: %s [-i] [-k] [-u] [-v]\n",
	    argv0 ? argv0 : "selc");
	exit(2);
}

static void
attachio(int force)
{
	if (!AttachConsole(ATTACH_PARENT_PROCESS)) {
		if (!force)
			return;
		AllocConsole();
	}
	freopen("CONOUT$", "w", stdout);
	freopen("CONOUT$", "w", stderr);
}

static int
samewin(HWND a, HWND b)
{
	if (!a || !b)
		return 0;
	if (a == b)
		return 1;
	return GetAncestor(a, GA_ROOT) == GetAncestor(b, GA_ROOT);
}

static int
isclient(HWND hw, POINT pt)
{
	DWORD_PTR ht;
	LPARAM lp;

	lp = MAKELPARAM((short)pt.x, (short)pt.y);
	if (!SendMessageTimeoutA(hw, WM_NCHITTEST, 0, lp,
	    SMTO_ABORTIFHUNG | SMTO_NORMAL, HITTO, &ht))
		return 0;
	return ht == HTCLIENT;
}

static int
ispass(HWND hw)
{
	char cls[32];
	LONG style;

	if (!hw)
		return 0;
	if (!GetClassNameA(hw, cls, (int)sizeof(cls)))
		return 0;
	if (strcmp(cls, "Edit") != 0)
		return 0;
	style = GetWindowLongA(hw, GWL_STYLE);
	return (style & ES_PASSWORD) != 0;
}

static int
skipctrlc(HWND hw)
{
	char cls[64];
	HWND root;
	int i;

	root = GetAncestor(hw, GA_ROOT);
	if (!root)
		root = hw;
	if (!GetClassNameA(root, cls, (int)sizeof(cls)))
		return 0;
	for (i = 0; i < (int)LEN(nocctrlc); i++)
		if (strcmp(cls, nocctrlc[i]) == 0)
			return 1;
	return 0;
}

static int
goodsel(const struct Sel *s)
{
	HWND a, b;

	a = WindowFromPoint(s->down);
	b = WindowFromPoint(s->up);
	if (!a || !isclient(a, s->down))
		return 0;
	if (!samewin(a, b))
		return 0;
	return 1;
}

static void
sendctrlc(void)
{
	INPUT in[4];
	int n, ctrl;

	ctrl = (GetAsyncKeyState(VK_CONTROL) & 0x8000) != 0;
	memset(in, 0, sizeof(in));
	n = 0;
	if (!ctrl) {
		in[n].type = INPUT_KEYBOARD;
		in[n].ki.wVk = VK_CONTROL;
		n++;
	}
	in[n].type = INPUT_KEYBOARD;
	in[n].ki.wVk = 'C';
	n++;
	in[n].type = INPUT_KEYBOARD;
	in[n].ki.wVk = 'C';
	in[n].ki.dwFlags = KEYEVENTF_KEYUP;
	n++;
	if (!ctrl) {
		in[n].type = INPUT_KEYBOARD;
		in[n].ki.wVk = VK_CONTROL;
		in[n].ki.dwFlags = KEYEVENTF_KEYUP;
		n++;
	}
	SendInput((UINT)n, in, sizeof(INPUT));
}

static void
docopy(void)
{
	HWND fg, foc;
	DWORD tid, mine, seq;
	DWORD_PTR dummy;
	int attached;

	if (GetAsyncKeyState(VK_MENU) & 0x8000)
		return;
	fg = GetForegroundWindow();
	if (!fg)
		return;
	mine = GetCurrentThreadId();
	tid = GetWindowThreadProcessId(fg, NULL);
	attached = 0;
	if (tid && tid != mine)
		attached = AttachThreadInput(mine, tid, TRUE);
	foc = GetFocus();
	if (!foc)
		foc = fg;
	if (ispass(foc))
		goto done;
	seq = GetClipboardSequenceNumber();
	SendMessageTimeoutA(foc, WM_COPY, 0, 0,
	    SMTO_ABORTIFHUNG | SMTO_NORMAL, COPYTO, &dummy);
	if (GetClipboardSequenceNumber() != seq) {
		if (verbose)
			fprintf(stderr, "selc: WM_COPY\n");
		goto done;
	}
	if (skipctrlc(foc) || skipctrlc(fg)) {
		if (verbose)
			fprintf(stderr, "selc: skip Ctrl+C\n");
		goto done;
	}
	sendctrlc();
	if (verbose)
		fprintf(stderr, "selc: Ctrl+C\n");
done:
	if (attached)
		AttachThreadInput(mine, tid, FALSE);
}

static void
request(int kbd)
{
	sel.kbd = kbd;
	pending = sel;
	if (hwnd)
		PostMessageA(hwnd, WM_SEL, 0, 0);
}

static LRESULT CALLBACK
mouseproc(int code, WPARAM wp, LPARAM lp)
{
	MSLLHOOKSTRUCT *ms;
	static POINT down, lastpt;
	static int dragged, clicks, armed;
	static DWORD lastup;
	long dx, dy;
	int dt, cx, cy;

	if (code != HC_ACTION)
		return CallNextHookEx(mousehook, code, wp, lp);
	ms = (MSLLHOOKSTRUCT *)lp;
	if (wp == WM_LBUTTONDOWN) {
		down = ms->pt;
		dragged = 0;
		armed = 1;
		dt = (int)(ms->time - lastup);
		cx = GetSystemMetrics(SM_CXDOUBLECLK);
		cy = GetSystemMetrics(SM_CYDOUBLECLK);
		if (DBLCOPY && dt >= 0 &&
		    dt <= (int)GetDoubleClickTime() &&
		    abs(ms->pt.x - lastpt.x) <= cx &&
		    abs(ms->pt.y - lastpt.y) <= cy)
			clicks++;
		else
			clicks = 1;
		lastpt = ms->pt;
	} else if (wp == WM_MOUSEMOVE) {
		if (armed && !dragged) {
			dx = (long)ms->pt.x - down.x;
			dy = (long)ms->pt.y - down.y;
			if (dx * dx + dy * dy >=
			    (long)MINPIX * MINPIX)
				dragged = 1;
		}
	} else if (wp == WM_LBUTTONUP) {
		lastup = ms->time;
		lastpt = ms->pt;
		if (armed) {
			sel.down = down;
			sel.up = ms->pt;
			sel.dragged = dragged;
			sel.clicks = clicks;
			sel.shift = (GetAsyncKeyState(VK_SHIFT) &
			    0x8000) != 0;
			if (sel.dragged ||
			    (DBLCOPY && sel.clicks >= 2) ||
			    (SHIFTCLICK && sel.shift &&
			    sel.clicks == 1))
				request(0);
		}
		armed = 0;
		dragged = 0;
	} else if (wp == WM_RBUTTONDOWN || wp == WM_MBUTTONDOWN) {
		armed = 0;
		dragged = 0;
	}
	return CallNextHookEx(mousehook, code, wp, lp);
}

static LRESULT CALLBACK
kbdproc(int code, WPARAM wp, LPARAM lp)
{
	KBDLLHOOKSTRUCT *kb;
	int shift, ctrl, alt, nav, akey;

	if (code != HC_ACTION)
		return CallNextHookEx(kbdhook, code, wp, lp);
	if (wp != WM_KEYUP && wp != WM_SYSKEYUP)
		return CallNextHookEx(kbdhook, code, wp, lp);
	kb = (KBDLLHOOKSTRUCT *)lp;
	shift = (GetAsyncKeyState(VK_SHIFT) & 0x8000) != 0;
	ctrl = (GetAsyncKeyState(VK_CONTROL) & 0x8000) != 0;
	alt = (GetAsyncKeyState(VK_MENU) & 0x8000) != 0;
	nav = 0;
	switch (kb->vkCode) {
	case VK_LEFT:
	case VK_RIGHT:
	case VK_UP:
	case VK_DOWN:
	case VK_HOME:
	case VK_END:
	case VK_PRIOR:
	case VK_NEXT:
		nav = 1;
		break;
	}
	akey = CTRLA && kb->vkCode == 'A' && ctrl && !shift && !alt;
	if ((KEYSEL && nav && shift && !alt) || akey)
		request(1);
	return CallNextHookEx(kbdhook, code, wp, lp);
}

static LRESULT CALLBACK
wndproc(HWND hw, UINT msg, WPARAM wp, LPARAM lp)
{
	switch (msg) {
	case WM_SEL:
		if (COPYWAIT > 0)
			SetTimer(hw, 1, COPYWAIT, NULL);
		else if (pending.kbd || goodsel(&pending))
			docopy();
		return 0;
	case WM_TIMER:
		if (wp != 1)
			break;
		KillTimer(hw, 1);
		if (pending.kbd || goodsel(&pending))
			docopy();
		return 0;
	case WM_DESTROY:
		PostQuitMessage(0);
		return 0;
	}
	return DefWindowProcA(hw, msg, wp, lp);
}

static int
runkey(int add)
{
	HKEY k;
	LONG e;
	wchar_t path[MAX_PATH], val[MAX_PATH + 3];
	DWORD n;
	const wchar_t *sub =
	    L"Software\\Microsoft\\Windows\\CurrentVersion\\Run";

	e = RegCreateKeyExW(HKEY_CURRENT_USER, sub, 0, NULL, 0,
	    KEY_SET_VALUE | KEY_QUERY_VALUE, NULL, &k, NULL);
	if (e != ERROR_SUCCESS) {
		SetLastError((DWORD)e);
		die("RegCreateKeyEx:");
	}
	if (!add) {
		RegDeleteValueW(k, L"selc");
		RegCloseKey(k);
		if (verbose)
			fprintf(stderr, "selc: uninstalled Run key\n");
		return 0;
	}
	n = GetModuleFileNameW(NULL, path, (DWORD)LEN(path));
	if (n == 0 || n >= LEN(path)) {
		RegCloseKey(k);
		die("GetModuleFileName:");
	}
	if (wcschr(path, L' ')) {
		val[0] = L'"';
		wcsncpy(val + 1, path, LEN(val) - 3);
		val[LEN(val) - 2] = L'\0';
		n = (DWORD)wcslen(val);
		if (n < LEN(val) - 1) {
			val[n] = L'"';
			val[n + 1] = L'\0';
		}
	} else {
		wcsncpy(val, path, LEN(val));
		val[LEN(val) - 1] = L'\0';
	}
	n = (DWORD)((wcslen(val) + 1) * sizeof(wchar_t));
	e = RegSetValueExW(k, L"selc", 0, REG_SZ, (BYTE *)val, n);
	RegCloseKey(k);
	if (e != ERROR_SUCCESS) {
		SetLastError((DWORD)e);
		die("RegSetValueEx:");
	}
	if (verbose)
		fprintf(stderr, "selc: installed Run key\n");
	return 0;
}

static void
signalquit(void)
{
	HANDLE e;

	e = OpenEventA(EVENT_MODIFY_STATE, FALSE, "Local\\selc-quit");
	if (!e)
		return;
	SetEvent(e);
	CloseHandle(e);
}

static int
loop(HINSTANCE inst)
{
	WNDCLASSA wc;
	MSG msg;
	DWORD r;

	memset(&wc, 0, sizeof(wc));
	wc.lpfnWndProc = wndproc;
	wc.hInstance = inst;
	wc.lpszClassName = "selc";
	if (!RegisterClassA(&wc))
		die("RegisterClass:");
	hwnd = CreateWindowExA(0, "selc", "selc", WS_OVERLAPPED,
	    0, 0, 0, 0, HWND_MESSAGE, NULL, inst, NULL);
	if (!hwnd)
		die("CreateWindowEx:");
	mousehook = SetWindowsHookExA(WH_MOUSE_LL, mouseproc, inst, 0);
	if (!mousehook)
		die("SetWindowsHookEx mouse:");
	if (KEYSEL || CTRLA) {
		kbdhook = SetWindowsHookExA(WH_KEYBOARD_LL, kbdproc,
		    inst, 0);
		if (!kbdhook)
			die("SetWindowsHookEx kbd:");
	}
	if (verbose)
		fprintf(stderr, "selc: running\n");
	for (;;) {
		r = MsgWaitForMultipleObjects(1, &quit, FALSE,
		    INFINITE, QS_ALLINPUT);
		if (r == WAIT_OBJECT_0)
			break;
		if (r != WAIT_OBJECT_0 + 1)
			die("MsgWaitForMultipleObjects:");
		while (PeekMessageA(&msg, NULL, 0, 0, PM_REMOVE)) {
			if (msg.message == WM_QUIT)
				goto done;
			TranslateMessage(&msg);
			DispatchMessageA(&msg);
		}
	}
done:
	if (kbdhook)
		UnhookWindowsHookEx(kbdhook);
	if (mousehook)
		UnhookWindowsHookEx(mousehook);
	return 0;
}

int
main(int argc, char *argv[])
{
	HINSTANCE inst;
	HANDLE mtx;
	char *a, *p;
	int i, iflag, kflag, uflag;

	argv0 = argv[0] ? argv[0] : "selc";
	iflag = kflag = uflag = 0;
	attachio(0);
	for (i = 1; i < argc; i++) {
		a = argv[i];
		if (a[0] != '-' || a[1] == '\0')
			usage();
		if (a[1] == '-' && a[2] == '\0') {
			i++;
			break;
		}
		for (p = a + 1; *p; p++)
			switch (*p) {
			case 'i':
				iflag = 1;
				break;
			case 'k':
				kflag = 1;
				break;
			case 'u':
				uflag = 1;
				break;
			case 'v':
				verbose = 1;
				break;
			case 'h':
				usage();
				break; /* NOTREACHED */
			default:
				usage();
			}
	}
	if (i != argc)
		usage();
	if (verbose)
		attachio(1);
	if (kflag) {
		signalquit();
		return 0;
	}
	if (uflag) {
		runkey(0);
		signalquit();
		return 0;
	}
	if (iflag)
		runkey(1);

	inst = GetModuleHandleA(NULL);
	quit = CreateEventA(NULL, TRUE, FALSE, "Local\\selc-quit");
	if (!quit)
		die("CreateEvent:");
	mtx = CreateMutexA(NULL, TRUE, "Local\\selc-inst");
	if (!mtx)
		die("CreateMutex:");
	if (GetLastError() == ERROR_ALREADY_EXISTS)
		return 0;
	return loop(inst);
}