// Copyright Epic Games, Inc. All Rights Reserved.

#include <zenutil/consoletui.h>

#include <zencore/string.h>
#include <zencore/zencore.h>

#include <fmt/format.h>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
#else
#	include <poll.h>
#	include <sys/ioctl.h>
#	include <termios.h>
#	include <unistd.h>
#endif

#include <algorithm>
#include <cstdio>

#if !ZEN_PLATFORM_WINDOWS
#	include <cerrno>
#	include <csignal>
#endif

namespace zen {

//////////////////////////////////////////////////////////////////////////
// Platform-specific terminal helpers

#if ZEN_PLATFORM_WINDOWS

static bool
CheckIsInteractiveTerminal()
{
	DWORD dwMode = 0;
	return GetConsoleMode(GetStdHandle(STD_INPUT_HANDLE), &dwMode) && GetConsoleMode(GetStdHandle(STD_OUTPUT_HANDLE), &dwMode);
}

static void
EnableVirtualTerminal()
{
	HANDLE hStdOut = GetStdHandle(STD_OUTPUT_HANDLE);
	DWORD  dwMode  = 0;
	if (GetConsoleMode(hStdOut, &dwMode))
	{
		SetConsoleMode(hStdOut, dwMode | ENABLE_VIRTUAL_TERMINAL_PROCESSING);
	}
}

// RAII guard: sets the console output code page for the lifetime of the object and
// restores the original on destruction. Required for UTF-8 glyphs to render correctly
// via printf/fflush since the default console code page is not UTF-8.
class ConsoleCodePageGuard
{
public:
	explicit ConsoleCodePageGuard(UINT NewCP) : m_OldCP(GetConsoleOutputCP()) { SetConsoleOutputCP(NewCP); }
	~ConsoleCodePageGuard() { SetConsoleOutputCP(m_OldCP); }

private:
	UINT m_OldCP;
};

// RAII guard: ORs additional flags into the console input mode and restores the
// original mode on destruction. Used to enable ENABLE_WINDOW_INPUT so that
// WINDOW_BUFFER_SIZE_EVENT records are delivered to ReadConsoleInputA.
class ConsoleInputModeGuard
{
public:
	ConsoleInputModeGuard(HANDLE Handle, DWORD AddFlags) : m_Handle(Handle)
	{
		if (GetConsoleMode(Handle, &m_OldMode) && SetConsoleMode(Handle, m_OldMode | AddFlags))
		{
			m_Valid = true;
		}
	}

	~ConsoleInputModeGuard()
	{
		if (m_Valid)
		{
			SetConsoleMode(m_Handle, m_OldMode);
		}
	}

private:
	HANDLE m_Handle;
	DWORD  m_OldMode = 0;
	bool   m_Valid	 = false;
};

enum class ConsoleKey
{
	Unknown,
	ArrowUp,
	ArrowDown,
	PageUp,
	PageDown,
	Enter,
	Escape,
	Resize,
};

static ConsoleKey
ReadKey()
{
	HANDLE		 hStdin = GetStdHandle(STD_INPUT_HANDLE);
	INPUT_RECORD Record{};
	DWORD		 dwRead = 0;
	while (true)
	{
		if (!ReadConsoleInputA(hStdin, &Record, 1, &dwRead))
		{
			return ConsoleKey::Escape;	// treat read error as cancel
		}
		if (Record.EventType == WINDOW_BUFFER_SIZE_EVENT)
		{
			return ConsoleKey::Resize;
		}
		if (Record.EventType == KEY_EVENT && Record.Event.KeyEvent.bKeyDown)
		{
			switch (Record.Event.KeyEvent.wVirtualKeyCode)
			{
				case VK_UP:
					return ConsoleKey::ArrowUp;
				case VK_DOWN:
					return ConsoleKey::ArrowDown;
				case VK_PRIOR:
					return ConsoleKey::PageUp;
				case VK_NEXT:
					return ConsoleKey::PageDown;
				case VK_RETURN:
					return ConsoleKey::Enter;
				case VK_ESCAPE:
					return ConsoleKey::Escape;
				default:
					break;
			}
		}
	}
}

#else  // POSIX

static bool
CheckIsInteractiveTerminal()
{
	return isatty(STDIN_FILENO) && isatty(STDOUT_FILENO);
}

static void
EnableVirtualTerminal()
{
	// ANSI escape codes are native on POSIX terminals; nothing to do
}

// RAII guard: switches the terminal to raw/unbuffered input mode and restores
// the original attributes on destruction.
class RawModeGuard
{
public:
	RawModeGuard()
	{
		if (tcgetattr(STDIN_FILENO, &m_OldAttrs) != 0)
		{
			return;
		}

		struct termios Raw = m_OldAttrs;
		Raw.c_iflag &= ~static_cast<tcflag_t>(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
		Raw.c_cflag |= CS8;
		Raw.c_lflag &= ~static_cast<tcflag_t>(ECHO | ICANON | IEXTEN | ISIG);
		Raw.c_cc[VMIN]	= 1;
		Raw.c_cc[VTIME] = 0;
		if (tcsetattr(STDIN_FILENO, TCSANOW, &Raw) == 0)
		{
			m_Valid = true;
		}
	}

	~RawModeGuard()
	{
		if (m_Valid)
		{
			tcsetattr(STDIN_FILENO, TCSANOW, &m_OldAttrs);
		}
	}

	bool IsValid() const { return m_Valid; }

private:
	struct termios m_OldAttrs = {};
	bool		   m_Valid	  = false;
};

static int
ReadByteWithTimeout(int TimeoutMs)
{
	struct pollfd Pfd
	{
		STDIN_FILENO, POLLIN, 0
	};
	if (poll(&Pfd, 1, TimeoutMs) > 0 && (Pfd.revents & POLLIN))
	{
		unsigned char c = 0;
		if (read(STDIN_FILENO, &c, 1) == 1)
		{
			return static_cast<int>(c);
		}
	}
	return -1;
}

// State for fullscreen live mode (alternate screen + raw input)
static struct termios s_SavedAttrs = {};
static bool			  s_InLiveMode = false;

// SIGWINCH delivery: the handler sets a flag, ReadKey treats a read() that
// returns EINTR with the flag set as a Resize event.
static volatile sig_atomic_t s_ResizePending = 0;

static void
SigwinchHandler(int)
{
	s_ResizePending = 1;
}

// RAII guard: installs a SIGWINCH handler without SA_RESTART so that the
// blocking read() in ReadKey returns EINTR when the terminal is resized.
class SigwinchGuard
{
public:
	SigwinchGuard()
	{
		struct sigaction Action = {};
		Action.sa_handler		= SigwinchHandler;
		Action.sa_flags			= 0;  // intentionally NOT SA_RESTART
		sigemptyset(&Action.sa_mask);
		if (sigaction(SIGWINCH, &Action, &m_OldAction) == 0)
		{
			m_Valid = true;
		}
	}

	~SigwinchGuard()
	{
		if (m_Valid)
		{
			sigaction(SIGWINCH, &m_OldAction, nullptr);
		}
	}

private:
	struct sigaction m_OldAction = {};
	bool			 m_Valid	 = false;
};

enum class ConsoleKey
{
	Unknown,
	ArrowUp,
	ArrowDown,
	PageUp,
	PageDown,
	Enter,
	Escape,
	Resize,
};

static ConsoleKey
ReadKey()
{
	unsigned char c = 0;
	ssize_t		  n = read(STDIN_FILENO, &c, 1);
	if (n != 1)
	{
		if (n < 0 && errno == EINTR && s_ResizePending)
		{
			s_ResizePending = 0;
			return ConsoleKey::Resize;
		}
		return ConsoleKey::Escape;	// treat read error as cancel
	}

	if (c == 27)  // ESC byte or start of an escape sequence
	{
		int Next = ReadByteWithTimeout(50);
		if (Next == '[')
		{
			int Final = ReadByteWithTimeout(50);
			if (Final == 'A')
			{
				return ConsoleKey::ArrowUp;
			}
			if (Final == 'B')
			{
				return ConsoleKey::ArrowDown;
			}
			// PageUp / PageDown arrive as ESC[5~ and ESC[6~ respectively.
			if (Final == '5' || Final == '6')
			{
				const ConsoleKey Mapped = (Final == '5') ? ConsoleKey::PageUp : ConsoleKey::PageDown;
				int				 Tilde	= ReadByteWithTimeout(50);
				if (Tilde == '~')
				{
					return Mapped;
				}
				return ConsoleKey::Unknown;
			}
		}
		return ConsoleKey::Escape;
	}

	if (c == '\r' || c == '\n')
	{
		return ConsoleKey::Enter;
	}

	return ConsoleKey::Unknown;
}

#endif	// ZEN_PLATFORM_WINDOWS / POSIX

//////////////////////////////////////////////////////////////////////////
// Public API

uint32_t
TuiConsoleColumns(uint32_t Default)
{
#if ZEN_PLATFORM_WINDOWS
	CONSOLE_SCREEN_BUFFER_INFO Csbi = {};
	if (GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &Csbi))
	{
		// Use visible window width, not buffer width — legacy cmd.exe configs may have
		// a buffer wider than the window, which would cause wrapping for callers that
		// size their output to this value.
		return static_cast<uint32_t>(Csbi.srWindow.Right - Csbi.srWindow.Left + 1);
	}
#else
	struct winsize Ws = {};
	if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &Ws) == 0 && Ws.ws_col > 0)
	{
		return static_cast<uint32_t>(Ws.ws_col);
	}
#endif
	return Default;
}

void
TuiEnableOutput()
{
	EnableVirtualTerminal();
#if ZEN_PLATFORM_WINDOWS
	SetConsoleOutputCP(CP_UTF8);
#endif
}

bool
TuiIsStdoutTty()
{
#if ZEN_PLATFORM_WINDOWS
	static bool Cached = [] {
		DWORD dwMode = 0;
		return GetConsoleMode(GetStdHandle(STD_OUTPUT_HANDLE), &dwMode) != 0;
	}();
	return Cached;
#else
	static bool Cached = isatty(STDOUT_FILENO) != 0;
	return Cached;
#endif
}

bool
IsTuiAvailable()
{
	static bool Cached = CheckIsInteractiveTerminal();
	return Cached;
}

int
TuiPickOne(std::string_view Title, std::span<const std::string> Items)
{
	EnableVirtualTerminal();

#if ZEN_PLATFORM_WINDOWS
	ConsoleCodePageGuard CodePageGuard(CP_UTF8);
	// Enable WINDOW_BUFFER_SIZE_EVENT delivery so ReadKey can observe terminal
	// resizes; default mode does not include ENABLE_WINDOW_INPUT.
	ConsoleInputModeGuard InputModeGuard(GetStdHandle(STD_INPUT_HANDLE), ENABLE_WINDOW_INPUT);
#else
	RawModeGuard RawMode;
	if (!RawMode.IsValid())
	{
		return -1;
	}
	SigwinchGuard  ResizeGuard;
#endif

	const int Count = static_cast<int>(Items.size());
	if (Count == 0)
	{
		return -1;
	}

	int SelectedIndex = 0;

	// Frame layout (in rows): top blank, title, blank, ViewportRows of items,
	// blank, hint — each terminated with \n. We must leave at least one row
	// of slack between the frame and the bottom of the terminal: otherwise the
	// initial render scrolls just enough to push the anchor row above the
	// visible window, and the subsequent `\033[<N>A` cursor-up clamps at row 1,
	// causing the frame to drift upward by one row on every keypress.
	constexpr int  Chrome	 = 5;
	constexpr int  SlackRows = 1;
	const uint32_t TermRows	 = TuiConsoleRows(0);
	int			   ViewportRows;
	if (TermRows == 0)
	{
		ViewportRows = Count;
	}
	else
	{
		ViewportRows = std::max(1, static_cast<int>(TermRows) - Chrome - SlackRows);
		if (ViewportRows > Count)
		{
			ViewportRows = Count;
		}
	}
	int FrameRows	 = ViewportRows + Chrome;
	int ScrollOffset = 0;

	// Display budget per item line: terminal width minus the 3-column indicator.
	// We treat byte length as an upper bound on display columns — for ASCII this
	// is exact, and for multi-byte UTF-8 it slightly under-fills, which is fine
	// (we'd rather under-fill than wrap). Refreshed whenever the terminal is
	// resized so labels never re-introduce the cursor-math drift caused by wrap.
	constexpr int kIndicatorCols = 3;
	int			  LabelBudget	 = std::max(0, static_cast<int>(TuiConsoleColumns(120)) - kIndicatorCols);

	// Walk back from a candidate UTF-8 byte cut point to a codepoint boundary so
	// truncation never splits a multi-byte sequence.
	auto Utf8BoundaryBefore = [](std::string_view S, int Pos) {
		while (Pos > 0 && (static_cast<unsigned char>(S[Pos]) & 0xC0) == 0x80)
		{
			--Pos;
		}
		return Pos;
	};

	bool FirstRender = true;

	// Build each frame into a single buffer and emit it with one TuiWrite call.
	// On Windows, individual printf calls each round-trip through the console's
	// VT parser, so batching the whole frame into one write is dramatically
	// faster and removes the visible per-line redraw the user observed.
	auto RenderAll = [&] {
		// Keep the selection inside the viewport.
		if (SelectedIndex < ScrollOffset)
		{
			ScrollOffset = SelectedIndex;
		}
		else if (SelectedIndex >= ScrollOffset + ViewportRows)
		{
			ScrollOffset = SelectedIndex - ViewportRows + 1;
		}

		ExtendableStringBuilder<4096> Frame;

		if (FirstRender)
		{
			// Hide the cursor on the very first frame.
			Frame.Append("\033[?25l");
		}
		else
		{
			// Move back to the start of the previous frame.
			fmt::format_to(StringBuilderAppender(Frame), "\033[{}A", FrameRows);
		}
		FirstRender = false;

		// Top blank.
		Frame.Append("\r\033[K\n");
		// Title.
		Frame.Append("\r\033[K");
		Frame.Append(Title);
		Frame.Append("\n");
		// Blank between title and items.
		Frame.Append("\r\033[K\n");

		// Items: emit exactly ViewportRows lines so the frame height matches FrameRows.
		for (int Row = 0; Row < ViewportRows; ++Row)
		{
			const int Idx = ScrollOffset + Row;
			Frame.Append("\r\033[K");
			if (Idx < Count)
			{
				const bool IsSelected = (Idx == SelectedIndex);
				if (IsSelected)
				{
					Frame.Append("\033[1;7m");
				}
				// \xe2\x96\xb6 = U+25B6 BLACK RIGHT-POINTING TRIANGLE
				Frame.Append(IsSelected ? " \xe2\x96\xb6 " : "   ");

				const std::string_view Label = Items[Idx];
				if (static_cast<int>(Label.size()) <= LabelBudget)
				{
					Frame.Append(Label);
				}
				else if (LabelBudget <= 3)
				{
					// Not enough room for a meaningful ellipsis; emit whatever fits.
					Frame.Append(std::string_view("...", static_cast<size_t>(LabelBudget)));
				}
				else
				{
					const int Cut = Utf8BoundaryBefore(Label, LabelBudget - 3);
					Frame.Append(Label.substr(0, static_cast<size_t>(Cut)));
					Frame.Append("...");
				}
				if (IsSelected)
				{
					Frame.Append("\033[0m");
				}
			}
			Frame.Append("\n");
		}

		// Blank between items and hint.
		Frame.Append("\r\033[K\n");

		// Hint footer.
		// \xe2\x86\x91 = U+2191 ^    \xe2\x86\x93 = U+2193 v
		Frame.Append(
			"\r\033[K  \033[2m\xe2\x86\x91/\xe2\x86\x93\033[0m navigate  "
			"\033[2mPgUp/PgDn\033[0m page  "
			"\033[2mEnter\033[0m confirm  "
			"\033[2mEsc\033[0m cancel");
		if (Count > ViewportRows)
		{
			fmt::format_to(StringBuilderAppender(Frame), "  \033[2m[{}/{}]\033[0m", SelectedIndex + 1, Count);
		}
		Frame.Append("\n");

		TuiWrite(Frame);
		TuiFlush();
	};

	RenderAll();

	int	 Result = -1;
	bool Done	= false;
	while (!Done)
	{
		ConsoleKey Key = ReadKey();
		switch (Key)
		{
			case ConsoleKey::ArrowUp:
				SelectedIndex = (SelectedIndex - 1 + Count) % Count;
				RenderAll();
				break;

			case ConsoleKey::ArrowDown:
				SelectedIndex = (SelectedIndex + 1) % Count;
				RenderAll();
				break;

			case ConsoleKey::PageUp:
				SelectedIndex = std::max(0, SelectedIndex - ViewportRows);
				RenderAll();
				break;

			case ConsoleKey::PageDown:
				SelectedIndex = std::min(Count - 1, SelectedIndex + ViewportRows);
				RenderAll();
				break;

			case ConsoleKey::Resize:
				{
					// Recompute viewport and label width for the new terminal size,
					// then clear the visible screen and redraw from row 1: the
					// previous frame's position relative to the visible viewport is
					// no longer trackable through cursor-up math after a resize.
					const uint32_t NewRows = TuiConsoleRows(0);
					if (NewRows > 0)
					{
						ViewportRows = std::max(1, static_cast<int>(NewRows) - Chrome - SlackRows);
						if (ViewportRows > Count)
						{
							ViewportRows = Count;
						}
					}
					FrameRows	= ViewportRows + Chrome;
					LabelBudget = std::max(0, static_cast<int>(TuiConsoleColumns(120)) - kIndicatorCols);
					TuiWrite("\033[H\033[2J");
					FirstRender = true;
					RenderAll();
					break;
				}

			case ConsoleKey::Enter:
				Result = SelectedIndex;
				Done   = true;
				break;

			case ConsoleKey::Escape:
				Done = true;
				break;

			default:
				break;
		}
	}

	// Restore cursor and add a blank line for visual separation.
	TuiWrite("\033[?25h\n");
	TuiFlush();

	return Result;
}

void
TuiEnterAlternateScreen()
{
	EnableVirtualTerminal();
#if ZEN_PLATFORM_WINDOWS
	SetConsoleOutputCP(CP_UTF8);
#endif

	// Enter alternate screen buffer + hide cursor in one write.
	TuiWrite("\033[?1049h\033[?25l");
	TuiFlush();

#if !ZEN_PLATFORM_WINDOWS
	if (tcgetattr(STDIN_FILENO, &s_SavedAttrs) == 0)
	{
		struct termios Raw = s_SavedAttrs;
		Raw.c_iflag &= ~static_cast<tcflag_t>(BRKINT | ICRNL | INPCK | ISTRIP | IXON);
		Raw.c_cflag |= CS8;
		Raw.c_lflag &= ~static_cast<tcflag_t>(ECHO | ICANON | IEXTEN | ISIG);
		Raw.c_cc[VMIN]	= 1;
		Raw.c_cc[VTIME] = 0;
		if (tcsetattr(STDIN_FILENO, TCSANOW, &Raw) == 0)
		{
			s_InLiveMode = true;
		}
	}
#endif
}

void
TuiExitAlternateScreen()
{
	// Show cursor + exit alternate screen buffer in one write.
	TuiWrite("\033[?25h\033[?1049l");
	TuiFlush();

#if !ZEN_PLATFORM_WINDOWS
	if (s_InLiveMode)
	{
		tcsetattr(STDIN_FILENO, TCSANOW, &s_SavedAttrs);
		s_InLiveMode = false;
	}
#endif
}

void
TuiCursorHome()
{
	TuiWrite("\033[H");
}

uint32_t
TuiConsoleRows(uint32_t Default)
{
#if ZEN_PLATFORM_WINDOWS
	CONSOLE_SCREEN_BUFFER_INFO Csbi = {};
	if (GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &Csbi))
	{
		return static_cast<uint32_t>(Csbi.srWindow.Bottom - Csbi.srWindow.Top + 1);
	}
#else
	struct winsize Ws = {};
	if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &Ws) == 0 && Ws.ws_row > 0)
	{
		return static_cast<uint32_t>(Ws.ws_row);
	}
#endif
	return Default;
}

bool
TuiPollQuit()
{
#if ZEN_PLATFORM_WINDOWS
	HANDLE hStdin  = GetStdHandle(STD_INPUT_HANDLE);
	DWORD  dwCount = 0;
	if (!GetNumberOfConsoleInputEvents(hStdin, &dwCount) || dwCount == 0)
	{
		return false;
	}
	INPUT_RECORD Record{};
	DWORD		 dwRead = 0;
	while (PeekConsoleInputA(hStdin, &Record, 1, &dwRead) && dwRead > 0)
	{
		ReadConsoleInputA(hStdin, &Record, 1, &dwRead);
		if (Record.EventType == KEY_EVENT && Record.Event.KeyEvent.bKeyDown)
		{
			WORD vk = Record.Event.KeyEvent.wVirtualKeyCode;
			char ch = Record.Event.KeyEvent.uChar.AsciiChar;
			if (vk == VK_ESCAPE || ch == 'q' || ch == 'Q')
			{
				return true;
			}
		}
	}
	return false;
#else
	// Non-blocking read: character 3 = Ctrl+C, 27 = Esc, 'q'/'Q' = quit
	int b = ReadByteWithTimeout(0);
	return (b == 3 || b == 27 || b == 'q' || b == 'Q');
#endif
}

void
TuiSetScrollRegion(uint32_t Top, uint32_t Bottom)
{
	char	  Buf[32];
	const int Len = std::snprintf(Buf, sizeof(Buf), "\033[%u;%ur", Top, Bottom);
	if (Len > 0)
	{
		TuiWrite(std::string_view(Buf, static_cast<size_t>(Len)));
	}
}

void
TuiResetScrollRegion()
{
	TuiWrite("\033[r");
}

void
TuiMoveCursor(uint32_t Row, uint32_t Col)
{
	char	  Buf[32];
	const int Len = std::snprintf(Buf, sizeof(Buf), "\033[%u;%uH", Row, Col);
	if (Len > 0)
	{
		TuiWrite(std::string_view(Buf, static_cast<size_t>(Len)));
	}
}

void
TuiSaveCursor()
{
	TuiWrite(
		"\033"
		"7");
}

void
TuiRestoreCursor()
{
	TuiWrite(
		"\033"
		"8");
}

void
TuiEraseLine()
{
	TuiWrite("\033[2K");
}

void
TuiWrite(std::string_view Text)
{
#if ZEN_PLATFORM_WINDOWS
	// On Windows, stdout to a console is line-buffered (or worse, effectively
	// per-character through the CRT's VT translation), so fwrite of a multi-line
	// frame turns into many individual WriteFile calls. Each one round-trips
	// through the conhost VT parser, which is visible as flicker / lag during
	// interactive redraws. When stdout is a real console, drain the CRT buffer
	// and emit the entire frame with one WriteConsoleW call.
	HANDLE hStdOut = GetStdHandle(STD_OUTPUT_HANDLE);
	DWORD  dwMode  = 0;
	if (hStdOut != INVALID_HANDLE_VALUE && GetConsoleMode(hStdOut, &dwMode))
	{
		fflush(stdout);

		ExtendableWideStringBuilder<4096> Wide;
		Utf8ToWide(Text, Wide);

		DWORD Written = 0;
		WriteConsoleW(hStdOut, Wide.Data(), static_cast<DWORD>(Wide.Size()), &Written, nullptr);
		return;
	}
#endif
	fwrite(Text.data(), 1, Text.size(), stdout);
}

void
TuiFlush()
{
	// Always fflush, even when TuiWrite bypassed the CRT — callers may have
	// emitted printf / fwrite output earlier in the sequence and rely on this
	// to drain the CRT buffer.
	fflush(stdout);
}

void
TuiShowCursor(bool Show)
{
	TuiWrite(Show ? "\033[?25h" : "\033[?25l");
}

}  // namespace zen