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

// Zen command line client utility
//

#include "consoleprogress.h"

#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/timer.h>
#include <zencore/windows.h>
#include <zenutil/consoletui.h>
#include <zenutil/progress.h>

#if !ZEN_PLATFORM_WINDOWS
#	include <csignal>
#endif

ZEN_THIRD_PARTY_INCLUDES_START
#include <gsl/gsl-lite.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

//////////////////////////////////////////////////////////////////////////

namespace zen {

// Global tracking for scroll region cleanup on abnormal termination (Ctrl+C etc.)
// Only one ProgressBar can own a scroll region at a time.
static std::atomic<class ConsoleProgressBar*> g_ActiveScrollRegionOwner{nullptr};
static std::atomic<uint32_t>				  g_ActiveScrollRegionRows{0};

static void
ResetScrollRegionRaw()
{
	// Signal-safe: emit raw escape sequences to restore terminal state.
	// These are async-signal-safe on POSIX (write()) and safe in console
	// ctrl handlers on Windows (WriteConsole is allowed).
	uint32_t Rows = g_ActiveScrollRegionRows.load(std::memory_order_acquire);
	if (Rows >= 3)
	{
		// Move to status line, erase it, reset scroll region, move cursor to end of content
		TuiMoveCursor(Rows, 1);
		TuiEraseLine();
		TuiResetScrollRegion();
		TuiMoveCursor(Rows - 1, 1);
	}
	else
	{
		TuiResetScrollRegion();
	}
	TuiShowCursor(true);
	TuiFlush();
}

#if ZEN_PLATFORM_WINDOWS
static BOOL WINAPI
ScrollRegionCtrlHandler(DWORD CtrlType)
{
	if (CtrlType == CTRL_C_EVENT || CtrlType == CTRL_BREAK_EVENT)
	{
		ResetScrollRegionRaw();
	}
	// Return FALSE so the default handler (process termination) still runs
	return FALSE;
}
#else
static struct sigaction s_PrevSigIntAction;
static struct sigaction s_PrevSigTermAction;

static void
ScrollRegionSignalHandler(int Signal)
{
	ResetScrollRegionRaw();

	// Re-raise with the previous handler
	struct sigaction* PrevAction = (Signal == SIGINT) ? &s_PrevSigIntAction : &s_PrevSigTermAction;
	sigaction(Signal, PrevAction, nullptr);
	raise(Signal);
}
#endif

static void
InstallScrollRegionCleanupHandler()
{
#if ZEN_PLATFORM_WINDOWS
	SetConsoleCtrlHandler(ScrollRegionCtrlHandler, TRUE);
#else
	struct sigaction Action = {};
	Action.sa_handler		= ScrollRegionSignalHandler;
	Action.sa_flags			= SA_RESETHAND;	 // one-shot
	sigemptyset(&Action.sa_mask);
	sigaction(SIGINT, &Action, &s_PrevSigIntAction);
	sigaction(SIGTERM, &Action, &s_PrevSigTermAction);
#endif
}

static void
RemoveScrollRegionCleanupHandler()
{
#if ZEN_PLATFORM_WINDOWS
	SetConsoleCtrlHandler(ScrollRegionCtrlHandler, FALSE);
#else
	sigaction(SIGINT, &s_PrevSigIntAction, nullptr);
	sigaction(SIGTERM, &s_PrevSigTermAction, nullptr);
#endif
}

#if ZEN_PLATFORM_WINDOWS
static HANDLE
GetConsoleHandle()
{
	static HANDLE hStdOut = GetStdHandle(STD_OUTPUT_HANDLE);
	return hStdOut;
}
#endif

static void
OutputToConsoleRaw(const char* String, size_t Length)
{
#if ZEN_PLATFORM_WINDOWS
	HANDLE hStdOut = GetConsoleHandle();
	if (TuiIsStdoutTty())
	{
		WriteConsoleA(hStdOut, String, (DWORD)Length, 0, 0);
	}
	else
	{
		::WriteFile(hStdOut, (LPCVOID)String, (DWORD)Length, 0, 0);
	}
#else
	fwrite(String, 1, Length, stdout);
#endif
}

static void
OutputToConsoleRaw(const std::string& String)
{
	OutputToConsoleRaw(String.c_str(), String.length());
}

static void
OutputToConsoleRaw(const StringBuilderBase& SB)
{
	OutputToConsoleRaw(SB.c_str(), SB.Size());
}

static uint32_t
GetUpdateDelayMS(ConsoleProgressMode InMode)
{
	switch (InMode)
	{
		case ConsoleProgressMode::Plain:
			return 5000;
		case ConsoleProgressMode::Pretty:
			return 200;
		case ConsoleProgressMode::Log:
			return 2000;
		case ConsoleProgressMode::Quiet:
			return 5000;
		default:
			ZEN_ASSERT(false);
			return 0;
	}
}

class ConsoleProgressBar : public ProgressBase::ProgressBar
{
public:
	explicit ConsoleProgressBar(ConsoleProgressMode InMode, std::string_view InSubTask);
	~ConsoleProgressBar();

	void UpdateState(const State& NewState, bool DoLinebreak) override;
	void ForceLinebreak() override;
	void Finish() override;

private:
	void SetupScrollRegion();
	void TeardownScrollRegion();
	void RenderStatusLine(std::string_view Line);

	const ConsoleProgressMode m_Mode;
	Stopwatch				  m_SW;
	uint64_t				  m_LastUpdateMS;
	uint64_t				  m_PausedMS;
	State					  m_State;
	const std::string		  m_SubTask;
	size_t					  m_LastOutputLength   = 0;
	bool					  m_ScrollRegionActive = false;
	uint32_t				  m_ScrollRegionRows   = 0;
};

ConsoleProgressBar::ConsoleProgressBar(ConsoleProgressMode InMode, std::string_view InSubTask)
: m_Mode((!TuiIsStdoutTty() && InMode == ConsoleProgressMode::Pretty) ? ConsoleProgressMode::Plain : InMode)
, m_LastUpdateMS((uint64_t)-1)
, m_PausedMS(0)
, m_SubTask(InSubTask)
{
	ZEN_ASSERT(InSubTask.find('\"') == std::string_view::npos);
	if (!m_SubTask.empty())
	{
		if (m_Mode == ConsoleProgressMode::Log)
		{
			std::string String = fmt::format("@progress push \"{}\"\n", m_SubTask);
			OutputToConsoleRaw(String);
		}
	}

	if (m_Mode == ConsoleProgressMode::Pretty)
	{
		SetupScrollRegion();
	}
}

ConsoleProgressBar::~ConsoleProgressBar()
{
	try
	{
		TeardownScrollRegion();
		ForceLinebreak();
		if (!m_SubTask.empty())
		{
			if (m_Mode == ConsoleProgressMode::Log)
			{
				const std::string String("@progress pop\n");
				OutputToConsoleRaw(String);
			}
		}
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("ConsoleProgressBar::~ConsoleProgressBar() failed with {}", Ex.what());
	}
}

void
ConsoleProgressBar::SetupScrollRegion()
{
	// Only one scroll region owner at a time; nested bars fall back to the inline \r path.
	if (g_ActiveScrollRegionOwner.load(std::memory_order_acquire) != nullptr)
	{
		return;
	}

	uint32_t Rows = TuiConsoleRows(0);
	if (Rows < 3)
	{
		return;
	}

	TuiEnableOutput();

	// Ensure cursor is not on the last row before we install the region.
	// Print a newline to push content up if needed, then set the region.
	OutputToConsoleRaw("\n", 1);
	TuiSetScrollRegion(1, Rows - 1);

	// Move cursor into the scroll region so normal output stays there
	TuiMoveCursor(Rows - 1, 1);

	m_ScrollRegionActive = true;
	m_ScrollRegionRows	 = Rows;

	g_ActiveScrollRegionRows.store(Rows, std::memory_order_release);
	g_ActiveScrollRegionOwner.store(this, std::memory_order_release);
	InstallScrollRegionCleanupHandler();
}

void
ConsoleProgressBar::TeardownScrollRegion()
{
	if (!m_ScrollRegionActive)
	{
		return;
	}
	m_ScrollRegionActive = false;

	RemoveScrollRegionCleanupHandler();
	g_ActiveScrollRegionOwner.store(nullptr, std::memory_order_release);
	g_ActiveScrollRegionRows.store(0, std::memory_order_release);

	// Emit all teardown escape sequences as a single atomic write
	ExtendableStringBuilder<128> Buf;
	Buf << fmt::format("\x1b[{};1H", m_ScrollRegionRows)	   // move to status line
		<< "\x1b[2K"										   // erase it
		<< "\x1b[r"											   // reset scroll region
		<< fmt::format("\x1b[{};1H", m_ScrollRegionRows - 1);  // move to end of content
	OutputToConsoleRaw(Buf);
	TuiFlush();
}

void
ConsoleProgressBar::RenderStatusLine(std::string_view Line)
{
	// Handle terminal resizes by re-querying row count
	uint32_t CurrentRows = TuiConsoleRows(0);
	if (CurrentRows >= 3 && CurrentRows != m_ScrollRegionRows)
	{
		// Terminal was resized - reinstall scroll region
		TuiSetScrollRegion(1, CurrentRows - 1);
		m_ScrollRegionRows = CurrentRows;
	}

	// Build the entire escape sequence as a single string so the console write
	// is atomic and log output from other threads cannot interleave.
	ExtendableStringBuilder<512> Buf;
	Buf << "\x1b"
		   "7"											  // ESC 7  - save cursor
		<< fmt::format("\x1b[{};1H", m_ScrollRegionRows)  // move to bottom row
		<< "\x1b[2K"									  // erase entire line
		<< Line											  // progress bar content
		<< "\x1b"
		   "8";	 // ESC 8  - restore cursor
	OutputToConsoleRaw(Buf);
}

void
ConsoleProgressBar::UpdateState(const State& NewState, bool DoLinebreak)
{
	ZEN_ASSERT(NewState.TotalCount >= NewState.RemainingCount);
	ZEN_ASSERT(NewState.Task.find('\"') == std::string::npos);
	if (DoLinebreak == false && m_State == NewState)
	{
		return;
	}

	uint64_t ElapsedTimeMS = NewState.OptionalElapsedTime == (uint64_t)-1 ? m_SW.GetElapsedTimeMs() : NewState.OptionalElapsedTime;
	if (m_LastUpdateMS != (uint64_t)-1)
	{
		if (!DoLinebreak && (NewState.Status == m_State.Status) && (NewState.Task == m_State.Task) &&
			((m_LastUpdateMS + 200) > ElapsedTimeMS))
		{
			return;
		}
		if (m_State.Status == State::EStatus::Paused)
		{
			uint64_t ElapsedSinceLast = ElapsedTimeMS - m_LastUpdateMS;
			m_PausedMS += ElapsedSinceLast;
		}
	}

	m_LastUpdateMS = ElapsedTimeMS;

	std::string Task = NewState.Task;
	switch (NewState.Status)
	{
		case State::EStatus::Aborted:
			Task = "Aborting";
			break;
		case State::EStatus::Paused:
			Task = "Paused";
			break;
		default:
			break;
	}
	if (NewState.Task.length() > Task.length())
	{
		Task += std::string(NewState.Task.length() - Task.length(), ' ');
	}

	const size_t PercentDone =
		NewState.TotalCount > 0u ? gsl::narrow<uint8_t>((100 * (NewState.TotalCount - NewState.RemainingCount)) / NewState.TotalCount) : 0u;

	uint64_t Completed	  = NewState.TotalCount - NewState.RemainingCount;
	uint64_t ETAElapsedMS = ElapsedTimeMS - m_PausedMS;
	uint64_t ETAMS = ((m_State.TotalCount == NewState.TotalCount) && (NewState.Status == State::EStatus::Running)) && (PercentDone > 5)
						 ? (ETAElapsedMS * NewState.RemainingCount) / Completed
						 : 0;
	const std::string ETAString = (ETAMS > 0) ? fmt::format(" ETA {}", NiceTimeSpanMs(ETAMS)) : "";

	if (m_Mode == ConsoleProgressMode::Plain)
	{
		const std::string Details = (!NewState.Details.empty()) ? fmt::format(": {}", NewState.Details) : "";
		const std::string Output  = fmt::format("{} {}% {}{}{}\n", Task, PercentDone, NiceTimeSpanMs(ElapsedTimeMS), ETAString, Details);
		OutputToConsoleRaw(Output);
		m_State = NewState;
	}
	else if (m_Mode == ConsoleProgressMode::Pretty)
	{
		size_t ProgressBarSize = 20;

		size_t ProgressBarCount = (ProgressBarSize * PercentDone) / 100;

		uint32_t ConsoleColumns = TuiConsoleColumns(1024);

		const std::string PercentString = fmt::format("{:#3}%", PercentDone);

		const std::string ProgressBarString =
			fmt::format(": |{}{}|", std::string(ProgressBarCount, '#'), std::string(ProgressBarSize - ProgressBarCount, ' '));

		const std::string ElapsedString = fmt::format(": {}", NiceTimeSpanMs(ElapsedTimeMS));

		const std::string DetailsString = (!NewState.Details.empty()) ? fmt::format(". {}", NewState.Details) : "";

		ExtendableStringBuilder<256> OutputBuilder;

		OutputBuilder << Task << " " << PercentString;
		if (OutputBuilder.Size() + 1 < ConsoleColumns)
		{
			size_t RemainingSpace = ConsoleColumns - (OutputBuilder.Size() + 1);
			bool   ElapsedFits	  = RemainingSpace >= ElapsedString.length();
			RemainingSpace -= ElapsedString.length();
			bool ETAFits = ElapsedFits && RemainingSpace >= ETAString.length();
			RemainingSpace -= ETAString.length();
			bool DetailsFits = ETAFits && RemainingSpace >= DetailsString.length();
			RemainingSpace -= DetailsString.length();
			bool ProgressBarFits = DetailsFits && RemainingSpace >= ProgressBarString.length();
			RemainingSpace -= ProgressBarString.length();

			if (ProgressBarFits)
			{
				OutputBuilder << ProgressBarString;
			}
			if (ElapsedFits)
			{
				OutputBuilder << ElapsedString;
			}
			if (ETAFits)
			{
				OutputBuilder << ETAString;
			}
			if (DetailsFits)
			{
				OutputBuilder << DetailsString;
			}
		}

		if (m_ScrollRegionActive)
		{
			// Render on the pinned bottom status line
			RenderStatusLine(OutputBuilder.ToView());
		}
		else
		{
			// Fallback: inline \r-based overwrite (terminal too small for scroll region)
			std::string_view	   Output = OutputBuilder.ToView();
			std::string::size_type EraseLength =
				m_LastOutputLength > (Output.length() + 1) ? (m_LastOutputLength - Output.length() - 1) : 0;
			ExtendableStringBuilder<256> LineToPrint;

			if (Output.length() + 1 + EraseLength >= ConsoleColumns)
			{
				if (m_LastOutputLength > 0)
				{
					LineToPrint << "\n";
				}
				LineToPrint << Output;
				DoLinebreak = true;
			}
			else
			{
				LineToPrint << "\r" << Output << std::string(EraseLength, ' ');
			}

			if (DoLinebreak)
			{
				LineToPrint << "\n";
			}

			OutputToConsoleRaw(LineToPrint);
			m_LastOutputLength = DoLinebreak ? 0 : (Output.length() + 1);  // +1 for \r prefix
		}

		m_State = NewState;
	}
	else if (m_Mode == ConsoleProgressMode::Log)
	{
		if (m_State.Task != NewState.Task ||
			m_State.Details != NewState.Details)  // TODO: Should we output just because details change? Will this spam the log collector?
		{
			std::string Details = (!NewState.Details.empty()) ? fmt::format(": {}", NewState.Details) : "";
			for (std::string::value_type& Char : Details)
			{
				if (Char == '"')
				{
					Char = '\'';
				}
			}
			const std::string Message =
				fmt::format("@progress \"{} {}{}{}\"\n", NewState.Task, NiceTimeSpanMs(ElapsedTimeMS), ETAString, Details);
			OutputToConsoleRaw(Message);
		}

		const size_t OldPercentDone =
			m_State.TotalCount > 0u ? gsl::narrow<uint8_t>((100 * (m_State.TotalCount - m_State.RemainingCount)) / m_State.TotalCount) : 0u;

		if (OldPercentDone != PercentDone)
		{
			const std::string Progress = fmt::format("@progress {}%\n", PercentDone);
			OutputToConsoleRaw(Progress);
		}
		m_State = NewState;
	}
}

void
ConsoleProgressBar::ForceLinebreak()
{
	if (m_LastOutputLength > 0)
	{
		State NewState = m_State;
		UpdateState(NewState, /*DoLinebreak*/ true);
	}
}

void
ConsoleProgressBar::Finish()
{
	TeardownScrollRegion();

	if (m_LastOutputLength > 0 || m_State.RemainingCount > 0)
	{
		State NewState			= m_State;
		NewState.RemainingCount = 0;
		NewState.Details		= "";
		UpdateState(NewState, /*DoLinebreak*/ true);
	}
	m_State			   = State{};
	m_LastOutputLength = 0;
	m_SW.Reset();
}

class ConsoleProgress : public ProgressBase
{
public:
	ConsoleProgress(ConsoleProgressMode InMode) : m_Mode(InMode) {}

	virtual void SetLogOperationName(std::string_view Name) override
	{
		ZEN_ASSERT(Name.find('\"') == std::string_view::npos);
		if (m_Mode == ConsoleProgressMode::Log)
		{
			std::string String = fmt::format("@progress \"{}\"\n", Name);
			OutputToConsoleRaw(String);
		}
	}

	virtual void SetLogOperationProgress(uint32_t StepIndex, uint32_t StepCount) override
	{
		if (m_Mode == ConsoleProgressMode::Log)
		{
			const size_t PercentDone = StepCount > 0u ? gsl::narrow<uint8_t>((100 * StepIndex) / StepCount) : 0u;
			std::string	 String		 = fmt::format("@progress {}%\n", PercentDone);
			OutputToConsoleRaw(String);
		}
	}
	virtual void PushLogOperation(std::string_view Name) override
	{
		ZEN_ASSERT(Name.find('\"') == std::string_view::npos);
		if (m_Mode == ConsoleProgressMode::Log)
		{
			std::string String = fmt::format("@progress push \"{}\"\n", Name);
			OutputToConsoleRaw(String);
		}
	}

	virtual void PopLogOperation() override
	{
		if (m_Mode == ConsoleProgressMode::Log)
		{
			const std::string String("@progress pop\n");
			OutputToConsoleRaw(String);
		}
	}

	virtual uint32_t GetProgressUpdateDelayMS() const override { return GetUpdateDelayMS(m_Mode); }

	virtual std::unique_ptr<ProgressBase::ProgressBar> CreateProgressBar(std::string_view InSubTask) override
	{
		return std::make_unique<ConsoleProgressBar>(m_Mode, InSubTask);
	}

private:
	ConsoleProgressMode m_Mode;
};

ProgressBase*
CreateConsoleProgress(ConsoleProgressMode InMode)
{
	return new ConsoleProgress(InMode);
}

}  // namespace zen