#include "Defines.h"

#include <windows.h>
#include <shellapi.h>
#include <psapi.h>
#include <imm.h>

#include <iostream>
#include <fstream>
#include <chrono>
#include <string>
#include <thread>
#include <ctime>
#include <map>

#pragma comment(lib, "psapi.lib")
#pragma comment(lib, "imm32.lib")

#pragma warning(disable: 4996)

// ReSharper disable CppDeprecatedEntity
// ReSharper disable CppClangTidyCertErr33C

using namespace std::chrono_literals;

constexpr bool bEnableCallback = true;
constexpr bool bEnablePolling  = true;
constexpr auto TickInterval    = 250ms;
constexpr auto IMEMessageDelay =  50ms;

const std::string LogFilename     = "FocusIME.log";
const std::string IMEModeFilename = "FocusIME.json";

static_assert(bEnableCallback || bEnablePolling);

std::ofstream GLogStream;

bool GShouldExit = false;

void PrintLog(const std::string& Text);

void Tick()
{
	// ReSharper disable CppInconsistentNaming
	// ReSharper disable CppClangTidyPerformanceEnumSize

	enum { IMC_GETCONVERSIONMODE = 0x0001 };
	enum { IMC_SETCONVERSIONMODE = 0x0002 };
	enum { IMC_GETOPENSTATUS     = 0x0005 };

	// ReSharper restore CppInconsistentNaming
	// ReSharper restore CppClangTidyPerformanceEnumSize

	enum class EIMEConversionMode : std::uint8_t
	{
		Default,
		English,
		Chinese,
	};

	using FIMEMode = std::map<std::string, EIMEConversionMode>;

	auto LoadIMEMode = []() -> FIMEMode
	{
		std::map<std::string, EIMEConversionMode> Result;

		std::ifstream File(IMEModeFilename);

		if (!File.is_open())
		{
			PrintLog("Conversion mode configuration file not found");

			return Result;
		}

		do
		{
			std::string Buffer;

			Buffer.assign(std::istreambuf_iterator<char>(File), std::istreambuf_iterator<char>());

			std::string_view View = Buffer;

			// { "ItemA": "ValueA", "ItemB": "ValueB" }

			while (!View.empty() && std::isspace(View.front())) View.remove_prefix(1);
			while (!View.empty() && std::isspace(View.back ())) View.remove_suffix(1);

			if (!View.starts_with('{') || !View.ends_with('}')) break;

			View.remove_prefix(1);
			View.remove_suffix(1);

			// "ItemA": "ValueA", "ItemB": "ValueB"

			const bool bSuccessful = [&Result, &View]
			{
				std::string_view::size_type Index = std::string_view::npos;

				while (true)
				{
					while (!View.empty() && std::isspace(View.front())) View.remove_prefix(1);

					if (!View.starts_with('"')) return false;

					View.remove_prefix(1);

					// ItemA": "ValueA", "ItemB": "ValueB"

					Index = View.find('"');

					if (Index == std::string_view::npos) return false;

					std::string_view Process = View.substr(0, Index);

					View.remove_prefix(Process.size() + 1);

					// : "ValueA", "ItemB": "ValueB"

					while (!View.empty() && std::isspace(View.front())) View.remove_prefix(1);

					if (!View.starts_with(':')) return false;

					View.remove_prefix(1);

					// "ValueA", "ItemB": "ValueB"

					while (!View.empty() && std::isspace(View.front())) View.remove_prefix(1);

					if (!View.starts_with('"')) return false;

					View.remove_prefix(1);

					// ValueA", "ItemB": "ValueB"

					Index = View.find('"');

					if (Index == std::string_view::npos) return false;

					std::string_view Mode = View.substr(0, Index);

					View.remove_prefix(Mode.size() + 1);

					// , "ItemB": "ValueB"

					while (!View.empty() && std::isspace(View.front())) View.remove_prefix(1);

					if (Mode != "English" && Mode != "Chinese" && Mode != "Default") return false;

					Result.emplace(Process,
						Mode == "English" ? EIMEConversionMode::English :
						Mode == "Chinese" ? EIMEConversionMode::Chinese : EIMEConversionMode::Default);

					if (View.empty()) break;

					if (!View.starts_with(',')) return false;

					View.remove_prefix(1);
				}

				return true;

			} ();

			if (!bSuccessful) break;

			PrintLog("Successfully loaded " + std::to_string(Result.size()) + " conversion mode items");

			return Result;
		}
		while (false);

		PrintLog("Invalid format detected in conversion mode configuration file");

		return Result;
	};

	auto SaveIMEMode = [](const FIMEMode& IMEMode)
	{
		std::ofstream File(IMEModeFilename);

		if (!File.is_open())
		{
			PrintLog("Error: Failed to save conversion mode configuration");

			return;
		}

		File << "{";

		bool bFirstItem = true;

		for (const auto& [Name, Mode] : IMEMode)
		{
			if (!bFirstItem) File << ",";

			File << "\n\t\"" << Name << "\": \"" <<
				(Mode == EIMEConversionMode::English ? "English" :
				 Mode == EIMEConversionMode::Chinese ? "Chinese" : "Default") << "\"";

			bFirstItem = false;
		}

		File << "\n}\n";

		File.close();
	};

	static FIMEMode IMEMode = LoadIMEMode();

	static HWND LastWindow = nullptr;

	static std::string CachedProcess;

	if (LastWindow != nullptr)
	{
		do
		{
			if (IMEMode.contains(CachedProcess) && IMEMode[CachedProcess] == EIMEConversionMode::Default) break;

			HKL KeyboardLayout = GetKeyboardLayout(GetWindowThreadProcessId(LastWindow, nullptr));

			LANGID LanguageID = LOWORD(KeyboardLayout);

			if (PRIMARYLANGID(LanguageID) != LANG_CHINESE) break;

			if (const HWND IMEWindow = ImmGetDefaultIMEWnd(LastWindow))
			{
				const LRESULT ConversionMode = SendMessage(IMEWindow, WM_IME_CONTROL, IMC_GETCONVERSIONMODE, 0);
				const LRESULT OpenStatus     = SendMessage(IMEWindow, WM_IME_CONTROL, IMC_GETOPENSTATUS,     0);

				if (!OpenStatus) break;

				const EIMEConversionMode CurrentMode =
					ConversionMode == 0x0000 ? EIMEConversionMode::English :
					ConversionMode == 0x0401 ? EIMEConversionMode::Chinese : EIMEConversionMode::Default;

				if (CurrentMode == EIMEConversionMode::Default) break;

				if (!IMEMode.contains(CachedProcess) || IMEMode[CachedProcess] != CurrentMode)
				{
					IMEMode[CachedProcess] = CurrentMode;

					PrintLog("Updated conversion mode for process '" + CachedProcess + "' to " +
						(CurrentMode == EIMEConversionMode::English ? "English" :
					 	 CurrentMode == EIMEConversionMode::Chinese ? "Chinese" : "Default"));

					SaveIMEMode(IMEMode);
				}
			}
		}
		while (false);
	}

	const HWND ForegroundWindow = GetForegroundWindow();

	if (ForegroundWindow == LastWindow) return;

	if (ForegroundWindow == nullptr)
	{
		LastWindow = nullptr;

		return;
	}

	do
	{
		DWORD PID = 0;

		GetWindowThreadProcessId(ForegroundWindow, &PID);

		const HANDLE ProcessHandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, PID);

		if (ProcessHandle == nullptr) break;

		char Buffer[MAX_PATH] = { };

		if (GetModuleBaseName(ProcessHandle, nullptr, Buffer, sizeof(Buffer)) == 0)
		{
			CloseHandle(ProcessHandle);

			break;
		}

		CloseHandle(ProcessHandle);

		LastWindow = ForegroundWindow;

		CachedProcess = Buffer;

		if (!IMEMode.contains(CachedProcess) || IMEMode[CachedProcess] == EIMEConversionMode::Default) break;

		// Sometimes the message will miss if we don't sleep for a little while.
		std::this_thread::sleep_for(IMEMessageDelay);

		if (const HWND IMEWindow = ImmGetDefaultIMEWnd(LastWindow))
		{
			const LRESULT ConversionMode = SendMessage(IMEWindow, WM_IME_CONTROL, IMC_GETCONVERSIONMODE, 0);
			const LRESULT OpenStatus     = SendMessage(IMEWindow, WM_IME_CONTROL, IMC_GETOPENSTATUS,     0);

			if (!OpenStatus) break;

			LPARAM TargetMode = IMEMode[CachedProcess] == EIMEConversionMode::Chinese ? 0x0401 : 0x0000;

			if (ConversionMode != TargetMode)
			{
				if (SendMessage(IMEWindow, WM_IME_CONTROL, IMC_SETCONVERSIONMODE, TargetMode) == 0)
				{
					PrintLog("Successfully applied conversion mode for process '" + CachedProcess + "'");
				}
				else
				{
					PrintLog("Error: Failed to apply conversion mode for process '" + CachedProcess + "'");
				}
			}
		}
	}
	while (false);
}

LRESULT CALLBACK MainWindowProc(HWND Window, UINT Message, WPARAM WParam, LPARAM LParam);

void CALLBACK FocusEventHookProc(HWINEVENTHOOK, DWORD, HWND, LONG, LONG, DWORD, DWORD) { Tick(); }

int WINAPI WinMain(HINSTANCE Instance, HINSTANCE, LPSTR, int)
{
	const HANDLE Mutex = CreateMutex(nullptr, TRUE, "FocusIME_SingleInstance_Mutex");

	if (Mutex == nullptr)
	{
		MessageBox(nullptr, "Failed to create mutex.", "FocusIME", MB_OK | MB_ICONERROR);

		return 1;
	}

	if (GetLastError() == ERROR_ALREADY_EXISTS)
	{
		MessageBox(nullptr, "FocusIME is already running.", "FocusIME", MB_OK | MB_ICONINFORMATION);

		CloseHandle(Mutex);

		return 0;
	}

	WNDCLASS MainWindowClass = { };

	MainWindowClass.lpfnWndProc   = MainWindowProc;
	MainWindowClass.hInstance     = Instance;
	MainWindowClass.lpszClassName = "FocusIME_MainWindow";
	MainWindowClass.hCursor       = LoadCursor(nullptr, IDC_ARROW);
	MainWindowClass.hbrBackground = reinterpret_cast<HBRUSH>(COLOR_WINDOW + 1); // NOLINT(performance-no-int-to-ptr)

	if (!RegisterClass(&MainWindowClass))
	{
		MessageBox(nullptr, "Failed to register main window class.", "FocusIME", MB_OK | MB_ICONERROR);

		CloseHandle(Mutex);

		return 1;
	}

	const HWND MainWindow = CreateWindow(
		"FocusIME_MainWindow",
		"FocusIME",
		WS_OVERLAPPEDWINDOW,
		CW_USEDEFAULT, CW_USEDEFAULT,
		CW_USEDEFAULT, CW_USEDEFAULT,
		nullptr, nullptr, Instance, nullptr);

	if (MainWindow == nullptr)
	{
		MessageBox(nullptr, "Failed to create main window.", "FocusIME", MB_OK | MB_ICONERROR);

		CloseHandle(Mutex);

		return 1;
	}

	ShowWindow(MainWindow, SW_HIDE);

	HWINEVENTHOOK FocusEventHook;

	if constexpr (bEnableCallback)
	{
		FocusEventHook = SetWinEventHook(
			EVENT_OBJECT_FOCUS, EVENT_OBJECT_FOCUS,
			nullptr,
			FocusEventHookProc,
			0, 0,
			WINEVENT_OUTOFCONTEXT | WINEVENT_SKIPOWNPROCESS);

		if (FocusEventHook == nullptr)
		{
			MessageBox(nullptr, "Failed to set focus event hook.", "FocusIME", MB_OK | MB_ICONWARNING);

			DestroyWindow(MainWindow);

			CloseHandle(Mutex);

			return 1;
		}
	}

#	if BUILD_DEBUG || BUILD_DEVELOPMENT
	{
		if (!AllocConsole())
		{
			MessageBox(nullptr, "Failed to allocate console.", "FocusIME", MB_OK | MB_ICONERROR);

			DestroyWindow(MainWindow);

			CloseHandle(Mutex);

			return 1;
		}

		freopen("CONOUT$", "w", stdout);
		freopen("CONOUT$", "w", stderr);
		freopen("CONIN$",  "r", stdin);

		SetConsoleTitle("FocusIME Debug Console");

		std::ios::sync_with_stdio(true);

		std::wcout.clear();
		std::cout.clear();
		std::wcerr.clear();
		std::cerr.clear();
		std::wcin.clear();
		std::cin.clear();
	}
#	endif

	GLogStream.open(LogFilename, std::ios::out | std::ios::app);

	if (!GLogStream.is_open())
	{
		MessageBox(nullptr, "Failed to open log file.", "FocusIME", MB_OK | MB_ICONERROR);

		CloseHandle(Mutex);

		return 1;
	}

	PrintLog("FocusIME application started successfully");

	MSG Message = { };

	if constexpr (bEnablePolling)
	{
		while (!GShouldExit)
		{
			if (PeekMessage(&Message, nullptr, 0, 0, PM_REMOVE)) {

				if (Message.message == WM_QUIT) break;

				TranslateMessage(&Message);
				DispatchMessage (&Message);

				continue;
			}

			Tick();

			std::this_thread::sleep_for(TickInterval);
		}
	}
	else
	{
		while (!GShouldExit)
		{
			const BOOL bResult = GetMessage(&Message, nullptr, 0, 0);

			if (bResult > 0)
			{
				TranslateMessage(&Message);
				DispatchMessage (&Message);
			}

			else if (bResult == 0) break;

			else
			{
				DWORD ErrorCode = GetLastError();

				LPSTR MessageBuffer = nullptr;

				FormatMessage(
					FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
					nullptr, ErrorCode, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
					reinterpret_cast<LPSTR>(&MessageBuffer), 0, nullptr);

				std::string ErrorMessage = "Message loop failed with error code " + std::to_string(ErrorCode);

				if (MessageBuffer)
				{
					ErrorMessage += ": " + std::string(MessageBuffer);

					LocalFree(MessageBuffer);
				}

				PrintLog(ErrorMessage);

				break;
			}
		}
	}

	PrintLog("FocusIME application is shutting down");

	GLogStream.close();

#	if BUILD_DEBUG || BUILD_DEVELOPMENT
	{
		FreeConsole();
	}
#	endif

	if constexpr (bEnableCallback)
	{
		UnhookWinEvent(FocusEventHook);
	}

	DestroyWindow(MainWindow);

	CloseHandle(Mutex);

	return 0;
}

LRESULT CALLBACK MainWindowProc(const HWND Window, const UINT Message, const WPARAM WParam, const LPARAM LParam)
{
	static NOTIFYICONDATA NotifyIconData;

	// ReSharper disable CppInconsistentNaming
	// ReSharper disable CppClangTidyPerformanceEnumSize

	enum { WM_TRAY_ICON = WM_USER + 1 };
	enum { ID_TRAY_EXIT = 1001        };

	// ReSharper restore CppInconsistentNaming
	// ReSharper restore CppClangTidyPerformanceEnumSize

	switch (Message)
	{
	case WM_CREATE:

		memset(&NotifyIconData, 0, sizeof(NOTIFYICONDATA));

		NotifyIconData.cbSize           = sizeof(NOTIFYICONDATA);
		NotifyIconData.hWnd             = Window;
		NotifyIconData.uID              = 1;
		NotifyIconData.uFlags           = NIF_ICON | NIF_MESSAGE | NIF_TIP;
		NotifyIconData.uCallbackMessage = WM_TRAY_ICON;

		NotifyIconData.hIcon = LoadIcon(nullptr, IDI_APPLICATION);

		strcpy(NotifyIconData.szTip, "FocusIME");

		if (!Shell_NotifyIcon(NIM_ADD, &NotifyIconData)) return -1;

		break;

	case WM_TRAY_ICON:

		switch (LParam)
		{
		case WM_RBUTTONUP:
		case WM_CONTEXTMENU:

			POINT Point;

			GetCursorPos(&Point);

			if (const HMENU Menu = CreatePopupMenu())
			{
				InsertMenu(Menu, -1, MF_BYPOSITION | MF_STRING, ID_TRAY_EXIT, "Exit(&X)");

				SetMenuDefaultItem(Menu, ID_TRAY_EXIT, FALSE);

				SetForegroundWindow(Window);

				TrackPopupMenu(Menu, TPM_LEFTALIGN | TPM_RIGHTBUTTON, Point.x, Point.y, 0, Window, nullptr);

				PostMessage(Window, WM_NULL, 0, 0);

				DestroyMenu(Menu);
			}

			break;

		default: break;
		}

		break;

	case WM_COMMAND:

		if (LOWORD(WParam) == ID_TRAY_EXIT)
		{
			PrintLog("Exit command received from system tray menu");

			GShouldExit = true;

			PostQuitMessage(0);
		}

		break;

	case WM_DESTROY:

		Shell_NotifyIcon(NIM_DELETE, &NotifyIconData);

		PostQuitMessage(0);

		break;

	default: return DefWindowProc(Window, Message, WParam, LParam);
	}

	return 0;
}

void PrintLog(const std::string& Text)
{
	const auto Now  = std::chrono::system_clock::now();
	const auto Time = std::chrono::system_clock::to_time_t(Now);

	const long long Milliseconds = (std::chrono::duration_cast<std::chrono::milliseconds>(Now.time_since_epoch()) % 1000).count();

	const std::tm* LocalTime = std::localtime(&Time); // NOLINT(concurrency-mt-unsafe)

	char Buffer[] = "[2025-06-21 20:05:04.305]: ";

	std::strftime(&Buffer[1], 20, "%Y-%m-%d %H:%M:%S", LocalTime);

	Buffer[20] = '.';

	constexpr char Digits[] = "0123456789";

	Buffer[21] = Digits[Milliseconds / 100 % 10];
	Buffer[22] = Digits[Milliseconds /  10 % 10];
	Buffer[23] = Digits[Milliseconds /   1 % 10];

	GLogStream << Buffer << Text << std::endl;

#if BUILD_DEBUG || BUILD_DEVELOPMENT
	{
		std::clog << Buffer << Text << std::endl;
	}
#endif
}

// ReSharper restore CppDeprecatedEntity
// ReSharper restore CppClangTidyCertErr33C