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

// Zen command line client utility
//

#include "progressbar.h"

#include <zencore/logging.h>
#include <zencore/windows.h>
#include <zenremotestore/operationlogoutput.h>
#include <zenutil/consoletui.h>

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

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

namespace zen {

#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());
}

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

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

void
ProgressBar::SetLogOperationProgress(Mode InMode, uint32_t StepIndex, uint32_t StepCount)
{
	if (InMode == Mode::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);
	}
}

void
ProgressBar::PushLogOperation(Mode InMode, std::string_view Name)
{
	if (InMode == Mode::Log)
	{
		std::string String = fmt::format("@progress push \"{}\"\n", Name);
		OutputToConsoleRaw(String);
	}
}

void
ProgressBar::PopLogOperation(Mode InMode)
{
	if (InMode == Mode::Log)
	{
		const std::string String("@progress pop\n");
		OutputToConsoleRaw(String);
	}
}

ProgressBar::ProgressBar(Mode InMode, std::string_view InSubTask)
: m_Mode((!TuiIsStdoutTty() && InMode == Mode::Pretty) ? Mode::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())
	{
		PushLogOperation(InMode, m_SubTask);
	}
}

ProgressBar::~ProgressBar()
{
	try
	{
		ForceLinebreak();
		if (!m_SubTask.empty())
		{
			PopLogOperation(m_Mode);
		}
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("ProgressBar::~ProgressBar() failed with {}", Ex.what());
	}
}

void
ProgressBar::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 == Mode::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 == Mode::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 << "\r" << 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;
			}
		}

		std::string_view	   Output	   = OutputBuilder.ToView();
		std::string::size_type EraseLength = m_LastOutputLength > Output.length() ? (m_LastOutputLength - Output.length()) : 0;

		ExtendableStringBuilder<256> LineToPrint;

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

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

		OutputToConsoleRaw(LineToPrint);

		m_LastOutputLength = DoLinebreak ? 0 : Output.length();
		m_State			   = NewState;
	}
	else if (m_Mode == Mode::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
ProgressBar::ForceLinebreak()
{
	if (m_LastOutputLength > 0)
	{
		State NewState = m_State;
		UpdateState(NewState, /*DoLinebreak*/ true);
	}
}

void
ProgressBar::Finish()
{
	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();
}

bool
ProgressBar::IsSameTask(std::string_view Task) const
{
	return Task == m_State.Task;
}

bool
ProgressBar::HasActiveTask() const
{
	return !m_State.Task.empty();
}

class ConsoleOpLogProgressBar : public OperationLogOutput::ProgressBar
{
public:
	ConsoleOpLogProgressBar(zen::ProgressBar::Mode InMode, std::string_view InSubTask) : m_Inner(InMode, InSubTask) {}

	virtual void UpdateState(const State& NewState, bool DoLinebreak)
	{
		zen::ProgressBar::State State = {.Task			 = NewState.Task,
										 .Details		 = NewState.Details,
										 .TotalCount	 = NewState.TotalCount,
										 .RemainingCount = NewState.RemainingCount,
										 .Status		 = ConvertStatus(NewState.Status)};
		m_Inner.UpdateState(State, DoLinebreak);
	}
	virtual void Finish() { m_Inner.Finish(); }

private:
	zen::ProgressBar::State::EStatus ConvertStatus(State::EStatus Status)
	{
		switch (Status)
		{
			case State::EStatus::Running:
				return zen::ProgressBar::State::EStatus::Running;
			case State::EStatus::Aborted:
				return zen::ProgressBar::State::EStatus::Aborted;
			case State::EStatus::Paused:
				return zen::ProgressBar::State::EStatus::Paused;
			default:
				return (zen::ProgressBar::State::EStatus)Status;
		}
	}
	zen::ProgressBar m_Inner;
};

class ConsoleOpLogOutput : public OperationLogOutput
{
public:
	ConsoleOpLogOutput(zen::ProgressBar::Mode InMode) : m_Mode(InMode) {}
	virtual void EmitLogMessage(const logging::LogPoint& Point, fmt::format_args Args) override
	{
		logging::EmitConsoleLogMessage(Point, Args);
	}

	virtual void SetLogOperationName(std::string_view Name) override { zen::ProgressBar::SetLogOperationName(m_Mode, Name); }
	virtual void SetLogOperationProgress(uint32_t StepIndex, uint32_t StepCount) override
	{
		zen::ProgressBar::SetLogOperationProgress(m_Mode, StepIndex, StepCount);
	}
	virtual uint32_t GetProgressUpdateDelayMS() override { return GetUpdateDelayMS(m_Mode); }

	virtual ProgressBar* CreateProgressBar(std::string_view InSubTask) override { return new ConsoleOpLogProgressBar(m_Mode, InSubTask); }

private:
	zen::ProgressBar::Mode m_Mode;
};

OperationLogOutput*
CreateConsoleLogOutput(ProgressBar::Mode InMode)
{
	return new ConsoleOpLogOutput(InMode);
}

}  // namespace zen