/* SPDX-License-Identifier: MIT
 *
 * Copyright (C) 2019 WireGuard LLC. All Rights Reserved.
 */

#include <stdlib.h>
#include <stdbool.h>
#include <stdint.h>
#include <string.h>
#include <assert.h>
#include <windows.h>
#include <richedit.h>
#include <richole.h>
#include <tom.h>

#include "syntaxedit.h"
#include "highlighter.h"

const GUID CDECL IID_ITextDocument = { 0x8CC497C0, 0xA1DF, 0x11CE, { 0x80, 0x98, 0x00, 0xAA, 0x00, 0x47, 0xBE, 0x5D } };

struct syntaxedit_data {
	IRichEditOle *irich;
	ITextDocument *idoc;
	bool highlight_guard;
};

static WNDPROC parent_proc;

struct span_style {
	COLORREF color;
	DWORD effects;
};

static const struct span_style stylemap[] = {
	[HighlightSection] = { .color = RGB(0x32, 0x6D, 0x74), .effects = CFE_BOLD },
	[HighlightField] = { .color = RGB(0x9B, 0x23, 0x93), .effects = CFE_BOLD },
	[HighlightPrivateKey] = { .color = RGB(0x64, 0x38, 0x20) },
	[HighlightPublicKey] = { .color = RGB(0x64, 0x38, 0x20) },
	[HighlightPresharedKey] = { .color = RGB(0x64, 0x38, 0x20) },
	[HighlightIP] = { .color = RGB(0x0E, 0x0E, 0xFF) },
	[HighlightCidr] = { .color = RGB(0x81, 0x5F, 0x03) },
	[HighlightHost] = { .color = RGB(0x0E, 0x0E, 0xFF) },
	[HighlightPort] = { .color = RGB(0x81, 0x5F, 0x03) },
	[HighlightMTU] = { .color = RGB(0x1C, 0x00, 0xCF) },
	[HighlightKeepalive] = { .color = RGB(0x1C, 0x00, 0xCF) },
	[HighlightComment] = { .color = RGB(0x53, 0x65, 0x79), .effects = CFE_ITALIC },
	[HighlightDelimiter] = { .color = RGB(0x00, 0x00, 0x00) },
#ifndef MOBILE_WGQUICK_SUBSET
	[HighlightTable] = { .color = RGB(0x1C, 0x00, 0xCF) },
	[HighlightFwMark] = { .color = RGB(0x1C, 0x00, 0xCF) },
	[HighlightSaveConfig] = { .color = RGB(0x81, 0x5F, 0x03) },
	[HighlightCmd] = { .color = RGB(0x63, 0x75, 0x89) },
#endif
	[HighlightError] = { .color = RGB(0xC4, 0x1A, 0x16), .effects = CFE_UNDERLINE }
};

static void highlight_text(HWND hWnd)
{
	GETTEXTLENGTHEX gettextlengthex = {
		.flags = GTL_NUMBYTES,
		.codepage = CP_ACP /* Probably CP_UTF8 would be better, but (wine at least) returns utf32 sizes. */
	};
	GETTEXTEX gettextex = {
		.flags = GT_NOHIDDENTEXT,
		.codepage = gettextlengthex.codepage
	};
	CHARFORMAT2 format = {
		.cbSize = sizeof(CHARFORMAT2),
		.dwMask = CFM_COLOR | CFM_CHARSET | CFM_SIZE | CFM_BOLD | CFM_ITALIC | CFM_UNDERLINE,
		.dwEffects = CFE_AUTOCOLOR,
		.yHeight = 20 * 10,
		.bCharSet = ANSI_CHARSET
	};
	struct syntaxedit_data *this = (struct syntaxedit_data *)GetWindowLongPtr(hWnd, GWLP_USERDATA);
	LRESULT msg_size;
	char *msg = NULL;
	struct highlight_span *spans = NULL;
	CHARRANGE orig_selection;
	POINT original_scroll;
	bool found_private_key = false;

	if (this->highlight_guard)
		return;
	this->highlight_guard = true;

	msg_size = SendMessage(hWnd, EM_GETTEXTLENGTHEX, (WPARAM)&gettextlengthex, 0);
	if (msg_size == E_INVALIDARG)
		return;
	gettextex.cb = msg_size + 1;

	msg = malloc(msg_size + 1);
	if (!msg)
		goto out;
	if (SendMessage(hWnd, EM_GETTEXTEX, (WPARAM)&gettextex, (LPARAM)msg) <= 0)
		goto out;

	/* By default we get CR not CRLF, so just convert to LF. */
	for (size_t i = 0; i < msg_size; ++i) {
		if (msg[i] == '\r')
			msg[i] = '\n';
	}

	spans = highlight_config(msg);
	if (!spans)
		goto out;

	this->idoc->lpVtbl->Undo(this->idoc, tomSuspend, NULL);
	SendMessage(hWnd, WM_SETREDRAW, FALSE, 0);
	SendMessage(hWnd, EM_EXGETSEL, 0, (LPARAM)&orig_selection);
	SendMessage(hWnd, EM_GETSCROLLPOS, 0, (LPARAM)&original_scroll);
	SendMessage(hWnd, EM_HIDESELECTION, TRUE, 0);
	SendMessage(hWnd, EM_SETCHARFORMAT, SCF_ALL, (LPARAM)&format);
	for (struct highlight_span *span = spans; span->type != HighlightEnd; ++span) {
		CHARRANGE selection = { span->start, span->len + span->start };
		SendMessage(hWnd, EM_EXSETSEL, 0, (LPARAM)&selection);
		format.crTextColor = stylemap[span->type].color;
		format.dwEffects = stylemap[span->type].effects;
		SendMessage(hWnd, EM_SETCHARFORMAT, SCF_SELECTION, (LPARAM)&format);
		if (span->type == HighlightPrivateKey && !found_private_key) {
			/* Rather than allocating a new string, we mangle this one, since (for now) we don't use msg again. */
			msg[span->start + span->len] = '\0';
			SendMessage(hWnd, SE_PRIVATE_KEY, 0, (LPARAM)&msg[span->start]);
			found_private_key = true;
		}
	}
	SendMessage(hWnd, EM_SETSCROLLPOS, 0, (LPARAM)&original_scroll);
	SendMessage(hWnd, EM_EXSETSEL, 0, (LPARAM)&orig_selection);
	SendMessage(hWnd, EM_HIDESELECTION, FALSE, 0);
	SendMessage(hWnd, WM_SETREDRAW, TRUE, 0);
	RedrawWindow(hWnd, NULL, NULL, RDW_ERASE | RDW_FRAME | RDW_INVALIDATE | RDW_ALLCHILDREN);
	this->idoc->lpVtbl->Undo(this->idoc, tomResume, NULL);
	if (!found_private_key)
		SendMessage(hWnd, SE_PRIVATE_KEY, 0, 0);

out:
	free(spans);
	free(msg);
	this->highlight_guard = false;
}

static void context_menu(HWND hWnd, INT x, INT y)
{
	GETTEXTLENGTHEX gettextlengthex = {
		.flags = GTL_DEFAULT,
		.codepage = CP_ACP
	};
	/* This disturbing hack grabs the system edit menu normally used for the EDIT control. */
	HMENU popup, menu = LoadMenuW(GetModuleHandleW(L"comctl32.dll"), MAKEINTRESOURCEW(1));
	CHARRANGE selection = { 0 };
	bool has_selection, can_selectall, can_undo, can_paste;
	UINT cmd;

	if (!menu)
		return;

	SendMessage(hWnd, EM_EXGETSEL, 0, (LPARAM)&selection);
	has_selection = selection.cpMax - selection.cpMin;
	can_selectall = selection.cpMin || (selection.cpMax < SendMessage(hWnd, EM_GETTEXTLENGTHEX, (WPARAM)&gettextlengthex, 0));
	can_undo = SendMessage(hWnd, EM_CANUNDO, 0, 0);
	can_paste = SendMessage(hWnd, EM_CANPASTE, CF_TEXT, 0);

	popup = GetSubMenu(menu, 0);
	EnableMenuItem(popup, WM_UNDO, MF_BYCOMMAND | (can_undo ? MF_ENABLED : MF_GRAYED));
	EnableMenuItem(popup, WM_CUT, MF_BYCOMMAND | (has_selection ? MF_ENABLED : MF_GRAYED));
	EnableMenuItem(popup, WM_COPY, MF_BYCOMMAND | (has_selection ? MF_ENABLED : MF_GRAYED));
	EnableMenuItem(popup, WM_PASTE, MF_BYCOMMAND | (can_paste ? MF_ENABLED : MF_GRAYED));
	EnableMenuItem(popup, WM_CLEAR, MF_BYCOMMAND | (has_selection ? MF_ENABLED : MF_GRAYED));
	EnableMenuItem(popup, EM_SETSEL, MF_BYCOMMAND | (can_selectall ? MF_ENABLED : MF_GRAYED));

	/* Delete items that we don't handle. */
	for (int ctl = GetMenuItemCount(popup) - 1; ctl >= 0; --ctl) {
		MENUITEMINFOW menu_item = {
			.cbSize = sizeof(MENUITEMINFOW),
			.fMask = MIIM_FTYPE | MIIM_ID
		};
		if (!GetMenuItemInfoW(popup, ctl, MF_BYPOSITION, &menu_item))
			continue;
		if (menu_item.fType & MFT_SEPARATOR)
			continue;
		switch (menu_item.wID) {
		case WM_UNDO:
		case WM_CUT:
		case WM_COPY:
		case WM_PASTE:
		case WM_CLEAR:
		case EM_SETSEL:
			continue;
		}
		DeleteMenu(popup, ctl, MF_BYPOSITION);
	}
	/* Delete trailing and adjacent separators. */
	for (int ctl = GetMenuItemCount(popup) - 1, end = true; ctl >= 0; --ctl) {
		MENUITEMINFOW menu_item = {
			.cbSize = sizeof(MENUITEMINFOW),
			.fMask = MIIM_FTYPE
		};
		if (!GetMenuItemInfoW(popup, ctl, MF_BYPOSITION, &menu_item))
			continue;
		if (!(menu_item.fType & MFT_SEPARATOR)) {
			end = false;
			continue;
		}
		if (!end && ctl) {
			if (!GetMenuItemInfoW(popup, ctl - 1, MF_BYPOSITION, &menu_item))
				continue;
			if (!(menu_item.fType & MFT_SEPARATOR))
				continue;
		}
		DeleteMenu(popup, ctl, MF_BYPOSITION);
	}

	if (x == -1 && y == -1) {
		RECT rect;
		GetWindowRect(hWnd, &rect);
		x = rect.left + (rect.right - rect.left) / 2;
		y = rect.top + (rect.bottom - rect.top) / 2;
	}

	if (GetFocus() != hWnd)
		SetFocus(hWnd);

	cmd = TrackPopupMenu(popup, TPM_LEFTALIGN | TPM_RIGHTBUTTON | TPM_RETURNCMD | TPM_NONOTIFY, x, y, 0, hWnd, NULL);
	if (cmd)
		SendMessage(hWnd, cmd, 0, cmd == EM_SETSEL ? -1 : 0);

	DestroyMenu(menu);
}

static LRESULT CALLBACK child_proc(HWND hWnd, UINT Msg, WPARAM wParam, LPARAM lParam)
{
	switch (Msg) {
	case WM_CREATE: {
		struct syntaxedit_data *this = calloc(1, sizeof(*this));
		SetWindowLong(hWnd, GWL_EXSTYLE,  GetWindowLong(hWnd, GWL_EXSTYLE) & ~WS_EX_CLIENTEDGE);
		assert(this);
		SetWindowLongPtr(hWnd, GWLP_USERDATA, (LONG_PTR)this);
		SendMessage(hWnd, EM_GETOLEINTERFACE, 0, (LPARAM)&this->irich);
		assert(this->irich);
		this->irich->lpVtbl->QueryInterface(this->irich, &IID_ITextDocument, (void **)&this->idoc);
		assert(this->idoc);
		SendMessage(hWnd, EM_SETEVENTMASK, 0, ENM_CHANGE);
		SendMessage(hWnd, EM_SETTEXTMODE, TM_SINGLECODEPAGE, 0);
		break;
	}
	case WM_DESTROY: {
		struct syntaxedit_data *this = (struct syntaxedit_data *)GetWindowLongPtr(hWnd, GWLP_USERDATA);
		this->idoc->lpVtbl->Release(this->idoc);
		this->irich->lpVtbl->Release(this->irich);
		free(this);
	}
	case WM_SETTEXT: {
		LRESULT ret = parent_proc(hWnd, Msg, wParam, lParam);
		highlight_text(hWnd);
		SendMessage(hWnd, EM_EMPTYUNDOBUFFER, 0, 0);
		return ret;
	}
	case WM_REFLECT + WM_COMMAND:
	case WM_COMMAND:
	case WM_REFLECT + WM_NOTIFY:
	case WM_NOTIFY:
		switch (HIWORD(wParam)) {
		case EN_CHANGE:
			highlight_text(hWnd);
			break;
		}
		break;
	case WM_PASTE:
		SendMessage(hWnd, EM_PASTESPECIAL, CF_TEXT, 0);
		return 0;
	case WM_KEYDOWN:
		if (!(GetKeyState(VK_CONTROL) & 0x8000))
			break;
		switch (LOWORD(wParam)) {
		case 'V':
			SendMessage(hWnd, EM_PASTESPECIAL, CF_TEXT, 0);
			return 0;
		}
		break;
	case WM_CONTEXTMENU:
		context_menu(hWnd, LOWORD(lParam), HIWORD(lParam));
		return 0;
	}
	return parent_proc(hWnd, Msg, wParam, lParam);
}

static long has_loaded = 0;

bool register_syntax_edit(void)
{
	WNDCLASSEXW class = { .cbSize = sizeof(WNDCLASSEXW) };
	WNDPROC pp;
	HANDLE lib;

	if (InterlockedCompareExchange(&has_loaded, 1, 0) != 0)
		return !!parent_proc;

	lib = LoadLibraryExW(L"msftedit.dll", NULL, LOAD_LIBRARY_SEARCH_SYSTEM32);
	if (!lib)
		return false;

	if (!GetClassInfoExW(NULL, L"RICHEDIT50W", &class))
		goto err;
	pp = class.lpfnWndProc;
	if (!pp)
		goto err;
	class.cbSize = sizeof(WNDCLASSEXW);
	class.hInstance	= GetModuleHandleW(NULL);
	class.lpszClassName = L"WgQuickSyntaxEdit";
	class.lpfnWndProc = child_proc;
	if (!RegisterClassExW(&class))
		goto err;
	parent_proc = pp;
	return true;

err:
	FreeLibrary(lib);
	return false;
}