#include "VisualProbeCommon.h"

#include <imm.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <iostream>
#include <mutex>
#include <sstream>

#pragma comment(lib, "Imm32.lib")

namespace VisualProbe
{
	namespace
	{
		struct Anchor
		{
			const char* id;
			POINT controlOrigin;
			POINT localCaretTopLeft;
			POINT localCaretBottom;
			RECT localLineRect;
		};

		struct ImePositionEvent
		{
			std::uint64_t sequence = 0;
			std::string stage;
			std::string positionTrigger = "none";
			std::string anchorId;
			DWORD threadId = 0;
			std::uintptr_t hkl = 0;
			std::wstring imeDescription;
			std::wstring imeFileName;
			bool apiAttempted = false;
			bool compositionActiveAtAttempt = false;
			bool contextAvailable = false;
			bool compositionSucceeded = false;
			bool candidateSucceeded = false;
			POINT hwndRelativePoint{};
			RECT hwndRelativeLineRect{};
			POINT screenPoint{};
		};

		struct ImePositionState
		{
			HWND hwnd = nullptr;
			WNDPROC previousWndProc = nullptr;
			std::atomic<int> activeAnchor{ 1 };
			std::atomic<bool> compositionActive{ false };
			std::atomic<bool> closeRequested{ false };
			std::atomic<std::uint64_t> nextSequence{ 0 };
			std::mutex eventsMutex;
			std::vector<ImePositionEvent> events;
		};

		constexpr std::array<Anchor, 3> Anchors{
			Anchor{ "anchor_top_left_safe", POINT{ 80, 125 }, POINT{ 30, 35 }, POINT{ 30, 61 }, RECT{ 10, 30, 250, 66 } },
			Anchor{ "anchor_center", POINT{ 350, 305 }, POINT{ 130, 55 }, POINT{ 130, 81 }, RECT{ 10, 50, 260, 86 } },
			Anchor{ "anchor_bottom_right_safe", POINT{ 650, 520 }, POINT{ 180, 45 }, POINT{ 180, 71 }, RECT{ 10, 40, 230, 76 } }
		};

		ImePositionState* activeState = nullptr;

		POINT surfaceCaret(const Anchor& anchor) noexcept
		{
			return POINT{
				anchor.controlOrigin.x + anchor.localCaretBottom.x,
				anchor.controlOrigin.y + anchor.localCaretBottom.y
			};
		}

		std::wstring imeDescription(HKL hkl)
		{
			const UINT length = ImmGetDescriptionW(hkl, nullptr, 0);
			if (!length) return {};
			std::wstring value(length + 1, L'\0');
			const UINT written = ImmGetDescriptionW(hkl, value.data(), static_cast<UINT>(value.size()));
			value.resize(written);
			return value;
		}

		std::wstring imeFileName(HKL hkl)
		{
			const UINT length = ImmGetIMEFileNameW(hkl, nullptr, 0);
			if (!length) return {};
			std::wstring value(length + 1, L'\0');
			const UINT written = ImmGetIMEFileNameW(hkl, value.data(), static_cast<UINT>(value.size()));
			value.resize(written);
			return value;
		}

		void appendEvent(ImePositionState& state, ImePositionEvent event)
		{
			event.sequence = state.nextSequence.fetch_add(1, std::memory_order_relaxed) + 1;
			std::lock_guard<std::mutex> lock(state.eventsMutex);
			state.events.push_back(std::move(event));
		}

		void recordLifecycle(ImePositionState& state, const char* stage)
		{
			const int anchorIndex = state.activeAnchor.load(std::memory_order_acquire);
			ImePositionEvent event;
			event.stage = stage;
			event.anchorId = Anchors[static_cast<std::size_t>(anchorIndex)].id;
			event.threadId = GetCurrentThreadId();
			const HKL hkl = GetKeyboardLayout(0);
			event.hkl = reinterpret_cast<std::uintptr_t>(hkl);
			event.imeDescription = imeDescription(hkl);
			event.imeFileName = imeFileName(hkl);
			appendEvent(state, std::move(event));
		}

		void configureImePosition(ImePositionState& state, const char* stage, const char* positionTrigger)
		{
			const int anchorIndex = state.activeAnchor.load(std::memory_order_acquire);
			const Anchor& anchor = Anchors[static_cast<std::size_t>(anchorIndex)];
			ImePositionEvent event;
			event.stage = stage;
			event.positionTrigger = positionTrigger;
			event.anchorId = anchor.id;
			event.threadId = GetCurrentThreadId();
			const HKL hkl = GetKeyboardLayout(0);
			event.hkl = reinterpret_cast<std::uintptr_t>(hkl);
			event.imeDescription = imeDescription(hkl);
			event.imeFileName = imeFileName(hkl);
			event.apiAttempted = true;
			event.compositionActiveAtAttempt = state.compositionActive.load(std::memory_order_acquire);
			event.hwndRelativePoint = surfaceCaret(anchor);
			event.hwndRelativeLineRect = RECT{
				anchor.controlOrigin.x + anchor.localLineRect.left,
				anchor.controlOrigin.y + anchor.localLineRect.top,
				anchor.controlOrigin.x + anchor.localLineRect.right,
				anchor.controlOrigin.y + anchor.localLineRect.bottom
			};
			event.screenPoint = event.hwndRelativePoint;
			ClientToScreen(state.hwnd, &event.screenPoint);

			HIMC context = ImmGetContext(state.hwnd);
			event.contextAvailable = context != 0;
			if (context)
			{
				COMPOSITIONFORM composition{};
				composition.dwStyle = CFS_POINT;
				composition.ptCurrentPos = event.hwndRelativePoint;
				event.compositionSucceeded = ImmSetCompositionWindow(context, &composition) != FALSE;

				CANDIDATEFORM candidate{};
				candidate.dwIndex = 0;
				candidate.dwStyle = CFS_EXCLUDE;
				candidate.ptCurrentPos = event.hwndRelativePoint;
				candidate.rcArea = event.hwndRelativeLineRect;
				event.candidateSucceeded = ImmSetCandidateWindow(context, &candidate) != FALSE;
				ImmReleaseContext(state.hwnd, context);
			}
			appendEvent(state, std::move(event));
		}

		LRESULT CALLBACK ProbeWndProc(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam)
		{
			ImePositionState* state = activeState;
			if (!state)
				return DefWindowProcA(hwnd, message, wParam, lParam);

			if (message == WM_KEYDOWN && wParam >= VK_F1 && wParam <= VK_F3)
			{
				state->activeAnchor.store(static_cast<int>(wParam - VK_F1), std::memory_order_release);
				configureImePosition(*state, "anchor_preconfigure", "anchor_change");
				return 0;
			}
			else if (message == WM_IME_STARTCOMPOSITION)
			{
				state->compositionActive.store(true, std::memory_order_release);
				configureImePosition(*state, "ime_start", "ime_lifecycle");
			}
			else if (message == WM_IME_COMPOSITION)
			{
				configureImePosition(*state, "ime_composition", "ime_lifecycle");
			}
			else if (message == WM_IME_ENDCOMPOSITION)
			{
				recordLifecycle(*state, "ime_end");
				state->compositionActive.store(false, std::memory_order_release);
			}
			else if (message == WM_INPUTLANGCHANGE)
			{
				recordLifecycle(*state, "input_language_change");
			}
			else if (message == WM_CLOSE)
			{
				recordLifecycle(*state, "close_requested");
				state->closeRequested.store(true, std::memory_order_release);
				return 0;
			}

			return state->previousWndProc
				? CallWindowProcA(state->previousWndProc, hwnd, message, wParam, lParam)
				: DefWindowProcA(hwnd, message, wParam, lParam);
		}

		void drawAnchorScreen(int activeAnchorIndex)
		{
			SetWorkingImage(nullptr);
			BeginBatchDraw();
			setbkcolor(RGB(238, 241, 246));
			cleardevice();
			setbkmode(TRANSPARENT);
			settextstyle(22, 0, "Microsoft YaHei UI");
			settextcolor(RGB(30, 45, 70));
			outtextxy(30, 25, "IME position probe: F1/F2/F3 select anchor; close window when done");
			outtextxy(30, 58, "Use only fixed non-sensitive input: ceshi; inspect candidate position; Esc cancels");
			for (std::size_t index = 0; index < Anchors.size(); ++index)
			{
				const Anchor& anchor = Anchors[index];
				const bool selected = static_cast<int>(index) == activeAnchorIndex;
				const COLORREF color = selected ? RGB(220, 60, 45) : RGB(45, 125, 190);
				const POINT caret = surfaceCaret(anchor);
				setlinecolor(color);
				setlinestyle(PS_SOLID, selected ? 3 : 1);
				rectangle(anchor.controlOrigin.x, anchor.controlOrigin.y,
					anchor.controlOrigin.x + anchor.localLineRect.right,
					anchor.controlOrigin.y + anchor.localLineRect.bottom + 12);
				line(caret.x, caret.y - 26, caret.x, caret.y + 4);
				line(caret.x - 5, caret.y, caret.x + 5, caret.y);
				settextcolor(color);
				const std::string label = std::string("F") + std::to_string(index + 1) + " " + anchor.id;
				outtextxy(anchor.controlOrigin.x + 8, anchor.controlOrigin.y + 5, label.c_str());
			}
			EndBatchDraw();
		}

		std::string normalizedRelative(const std::filesystem::path& path)
		{
			std::string value = path.generic_string();
			std::replace(value.begin(), value.end(), '\\', '/');
			return value;
		}

		void writeAnchor(std::ostream& output, const Anchor& anchor)
		{
			const POINT surfaceTopLeft{
				anchor.controlOrigin.x + anchor.localCaretTopLeft.x,
				anchor.controlOrigin.y + anchor.localCaretTopLeft.y
			};
			const POINT surfaceBottom = surfaceCaret(anchor);
			RECT surfaceLine{
				anchor.controlOrigin.x + anchor.localLineRect.left,
				anchor.controlOrigin.y + anchor.localLineRect.top,
				anchor.controlOrigin.x + anchor.localLineRect.right,
				anchor.controlOrigin.y + anchor.localLineRect.bottom
			};
			output << "{\"syntheticAnchorId\":" << jsonString(anchor.id)
				<< ",\"controlOrigin\":";
			writePointJson(output, anchor.controlOrigin);
			output << ",\"controlLocalCaretTopLeft\":";
			writePointJson(output, anchor.localCaretTopLeft);
			output << ",\"controlLocalCaretBottom\":";
			writePointJson(output, anchor.localCaretBottom);
			output << ",\"controlLocalLineRect\":";
			writeRectJson(output, anchor.localLineRect);
			output << ",\"easyxSurfaceCaretTopLeft\":";
			writePointJson(output, surfaceTopLeft);
			output << ",\"easyxSurfaceCaretBottom\":";
			writePointJson(output, surfaceBottom);
			output << ",\"easyxSurfaceLineRect\":";
			writeRectJson(output, surfaceLine);
			output << ",\"targetHwndRelativePoint\":";
			writePointJson(output, surfaceBottom);
			output << '}';
		}

		void writeEvent(std::ostream& output, const ImePositionEvent& event)
		{
			output << "{\"sequence\":" << event.sequence
				<< ",\"stage\":" << jsonString(event.stage)
				<< ",\"positionTrigger\":" << jsonString(event.positionTrigger)
				<< ",\"syntheticAnchorId\":" << jsonString(event.anchorId)
				<< ",\"threadId\":" << event.threadId
				<< ",\"hkl\":" << jsonString(hexValue(event.hkl))
				<< ",\"imeDescriptionUtf16\":" << utf16UnitsJson(event.imeDescription)
				<< ",\"imeFileNameUtf16\":" << utf16UnitsJson(event.imeFileName)
				<< ",\"apiAttempted\":";
			writeBoolean(output, event.apiAttempted);
			output << ",\"compositionActiveAtAttempt\":";
			writeBoolean(output, event.compositionActiveAtAttempt);
			output << ",\"contextAvailable\":";
			writeBoolean(output, event.contextAvailable);
			output << ",\"compositionForm\":{\"style\":" << CFS_POINT
				<< ",\"apiResult\":";
			writeBoolean(output, event.compositionSucceeded);
			output << "},\"candidateForm\":{\"style\":" << CFS_EXCLUDE
				<< ",\"index\":0,\"apiResult\":";
			writeBoolean(output, event.candidateSucceeded);
			output << ",\"excludeRect\":";
			writeRectJson(output, event.hwndRelativeLineRect);
			output << "},\"targetHwndRelativePoint\":";
			writePointJson(output, event.hwndRelativePoint);
			output << ",\"diagnosticScreenPoint\":";
			writePointJson(output, event.screenPoint);
			output << '}';
		}
	}

	int runImePositionProbe(const Options& options, const DpiConfiguration& dpiConfiguration)
	{
		HWND hwnd = initgraph(CanvasWidth, CanvasHeight, EX_SHOWCONSOLE);
		if (!hwnd)
		{
			std::cerr << "initgraph failed\n";
			return 3;
		}
		SetWindowTextA(hwnd, "StellarX 0C IME Position Probe");
		ShowWindow(hwnd, SW_SHOWNORMAL);
		SetForegroundWindow(hwnd);
		Sleep(80);
		const EnvironmentSnapshot environment = captureEnvironment(hwnd);
		DWORD windowProcessId = 0;
		const DWORD windowOwnerThread = GetWindowThreadProcessId(hwnd, &windowProcessId);

		ImePositionState state;
		state.hwnd = hwnd;
		activeState = &state;
		SetLastError(ERROR_SUCCESS);
		state.previousWndProc = reinterpret_cast<WNDPROC>(SetWindowLongPtrA(
			hwnd, GWLP_WNDPROC, reinterpret_cast<LONG_PTR>(&ProbeWndProc)));
		const DWORD subclassError = state.previousWndProc ? ERROR_SUCCESS : GetLastError();

		std::vector<EvidenceFile> evidence;
		bool evidenceComplete = state.previousWndProc != nullptr;
		for (std::size_t index = 0; index < Anchors.size(); ++index)
		{
			drawAnchorScreen(static_cast<int>(index));
			Sleep(40);
			const auto relative = std::filesystem::path("evidence") /
				(std::string(Anchors[index].id) + "_client.png");
			std::string backend;
			const bool saved = captureClientImage(hwnd, options.outputDirectory / relative, &backend);
			evidenceComplete = evidenceComplete && saved;
			if (saved)
				evidence.push_back({ normalizedRelative(relative), "synthetic anchor client evidence", backend });
		}
		state.activeAnchor.store(1, std::memory_order_release);
		drawAnchorScreen(1);

		int lastDrawnAnchor = 1;
		const ULONGLONG loopStarted = GetTickCount64();
		while (!state.closeRequested.load(std::memory_order_acquire) && IsWindow(hwnd))
		{
			if (options.autoCloseMilliseconds > 0 &&
				GetTickCount64() - loopStarted >= options.autoCloseMilliseconds)
			{
				recordLifecycle(state, "auto_close");
				state.closeRequested.store(true, std::memory_order_release);
				break;
			}
			const int currentAnchor = state.activeAnchor.load(std::memory_order_acquire);
			if (currentAnchor != lastDrawnAnchor)
			{
				drawAnchorScreen(currentAnchor);
				lastDrawnAnchor = currentAnchor;
			}
			Sleep(16);
		}

		if (state.previousWndProc && IsWindow(hwnd))
			SetWindowLongPtrA(hwnd, GWLP_WNDPROC,
				reinterpret_cast<LONG_PTR>(state.previousWndProc));
		activeState = nullptr;
		if (IsWindow(hwnd)) closegraph();

		std::vector<ImePositionEvent> events;
		{
			std::lock_guard<std::mutex> lock(state.eventsMutex);
			events = state.events;
		}
		std::array<bool, Anchors.size()> anchorSucceeded{};
		bool imeLifecycleObserved = false;
		bool imeLifecycleSucceeded = false;
		for (const auto& event : events)
		{
			if (event.positionTrigger == "ime_lifecycle") imeLifecycleObserved = true;
			if (event.positionTrigger == "ime_lifecycle" && event.contextAvailable &&
				event.compositionSucceeded && event.candidateSucceeded)
			{
				imeLifecycleSucceeded = true;
			}
			if (!event.apiAttempted || !event.compositionSucceeded || !event.candidateSucceeded)
				continue;
			for (std::size_t index = 0; index < Anchors.size(); ++index)
			{
				if (event.anchorId == Anchors[index].id) anchorSucceeded[index] = true;
			}
		}
		bool allAnchorsSucceeded = true;
		for (const bool value : anchorSucceeded) allAnchorsSucceeded = allAnchorsSucceeded && value;
		const bool completed = evidenceComplete && allAnchorsSucceeded && imeLifecycleSucceeded;

		std::ostringstream result;
		result << "{\"expectedImeName\":" << jsonString(options.expectedImeName)
			<< ",\"windowOwnerThreadId\":" << windowOwnerThread
			<< ",\"windowProcessId\":" << windowProcessId
			<< ",\"subclassInstalled\":";
		writeBoolean(result, state.previousWndProc != nullptr);
		result << ",\"subclassError\":" << subclassError
			<< ",\"allAnchorsApiSucceeded\":";
		writeBoolean(result, allAnchorsSucceeded);
		result << ",\"imeLifecycleObserved\":";
		writeBoolean(result, imeLifecycleObserved);
		result << ",\"imeLifecycleSucceeded\":";
		writeBoolean(result, imeLifecycleSucceeded);
		result << ",\"anchors\":[";
		for (std::size_t index = 0; index < Anchors.size(); ++index)
		{
			if (index) result << ',';
			writeAnchor(result, Anchors[index]);
		}
		result << "],\"imeEvents\":[";
		for (std::size_t index = 0; index < events.size(); ++index)
		{
			if (index) result << ',';
			writeEvent(result, events[index]);
		}
		result << "]}";

		std::string error;
		if (!writeRawResult(options, dpiConfiguration, environment, completed,
			evidence, result.str(), error))
		{
			std::cerr << error << '\n';
			return 4;
		}
		return completed ? 0 : 5;
	}
}