#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

std::ofstream GLogFile;

bool GShouldExit = false;

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

void Tick()
{
	using namespace std::chrono_literals;

	// 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,
	};

	static std::map<std::string, EIMEConversionMode> CachedIMEMode;

	static HWND LastWindow = nullptr;

	static std::string CachedProcess;

	if (LastWindow != nullptr)
	{
		do
		{
			if (CachedIMEMode.contains(CachedProcess) && CachedIMEMode[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 (!CachedIMEMode.contains(CachedProcess) || CachedIMEMode[CachedProcess] != CurrentMode)
				{
					CachedIMEMode[CachedProcess] = CurrentMode;

					PrintLog("Update the conversion mode cached for process '" + CachedProcess + "' to " +
						(CurrentMode == EIMEConversionMode::English ? "English" :
						 CurrentMode == EIMEConversionMode::Chinese ? "Chinese" : "Default"));
				}
			}
		}
		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 (!CachedIMEMode.contains(CachedProcess) || CachedIMEMode[CachedProcess] == EIMEConversionMode::Default) break;

		std::this_thread::sleep_for(50ms);

		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 = CachedIMEMode[CachedProcess] == EIMEConversionMode::Chinese ? 0x0401 : 0x0000;

			if (ConversionMode != TargetMode)
			{
				if (SendMessage(IMEWindow, WM_IME_CONTROL, IMC_SETCONVERSIONMODE, TargetMode) == 0)
				{
					PrintLog("Successfully set conversion mode for process '" + CachedProcess + "'");
				}
				else
				{
					PrintLog("Failed to set 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);

	const HWINEVENTHOOK 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

	GLogFile.open("Log.txt", std::ios::out | std::ios::app);

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

		CloseHandle(Mutex);

		return 1;
	}

	PrintLog("FocusIME started successfully");

	MSG Message = { };

	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 is shutting down...");

	GLogFile.close();

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

	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 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];

	GLogFile << Buffer << Text << std::endl;

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

// ReSharper restore CppDeprecatedEntity
// ReSharper restore CppClangTidyCertErr33C