/***************************************************************
 *
 * Copyright (C) 1990-2007, Condor Team, Computer Sciences Department,
 * University of Wisconsin-Madison, WI.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you
 * may not use this file except in compliance with the License.  You may
 * obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 ***************************************************************/


#include "condor_common.h"

#include <set>
#include "condor_debug.h"
#include "condor_fix_assert.h"
#include "condor_io.h"
#include "condor_classad.h"
#include "condor_attributes.h"
#include "condor_daemon_core.h"
#include "condor_config.h"
#include "condor_query.h"
#include "rooster.h"

Rooster::Rooster():
	m_polling_interval(-1),
	m_polling_timer(-1),
	m_max_unhibernate(0)
{
}

Rooster::~Rooster()
{
	stop();
}

void Rooster::init()
{
	config();
}

void Rooster::config()
{
	int old_polling_interval = m_polling_interval;
	m_polling_interval = param_integer("ROOSTER_INTERVAL",300);
	if( m_polling_interval < 0 ) {
		dprintf(D_ALWAYS,
				"ROOSTER_INTERVAL is less than 0, so no unhibernate checks "
				"will be made.\n");
		if( m_polling_timer != -1 ) {
			daemonCore->Cancel_Timer(m_polling_timer);
			m_polling_timer = -1;
		}
	}
	else if( m_polling_timer >= 0 ) {
		if( old_polling_interval != m_polling_interval ) {
			daemonCore->Reset_Timer(
				m_polling_timer,
				m_polling_interval,
				m_polling_interval);
		}
	}
	else {
		m_polling_timer = daemonCore->Register_Timer(
			m_polling_interval,
			m_polling_interval,
			(TimerHandlercpp)&Rooster::poll,
			"Rooster::poll",
			this );
	}
	if( old_polling_interval != m_polling_interval && m_polling_interval > 0 )
	{
		dprintf(D_ALWAYS,
				"Will perform unhibernate checks every ROOSTER_INTERVAL=%d "
				"seconds.\n", m_polling_interval);
	}

	ASSERT( param(m_unhibernate_constraint,"ROOSTER_UNHIBERNATE") );


	ASSERT( param(m_wakeup_cmd,"ROOSTER_WAKEUP_CMD") );

	m_wakeup_args.Clear();
	std::string error_msg;
	if( !m_wakeup_args.AppendArgsV2Quoted(m_wakeup_cmd.c_str(), error_msg) ) {
		EXCEPT("Invalid wakeup command %s: %s",
			   m_wakeup_cmd.c_str(), error_msg.c_str());
	}

	std::string rank;
	param(rank,"ROOSTER_UNHIBERNATE_RANK");
	if( rank.empty() ) {
		m_rank_ad.Delete(ATTR_RANK);
	}
	else {
		if( !m_rank_ad.AssignExpr(ATTR_RANK,rank.c_str()) ) {
			EXCEPT("Invalid expression for ROOSTER_UNHIBERNATE_RANK: %s",
				   rank.c_str());
		}
	}

	m_max_unhibernate = param_integer("ROOSTER_MAX_UNHIBERNATE",0,0);
}

void Rooster::stop()
{
	if( m_polling_timer != -1 ) {
		daemonCore->Cancel_Timer(m_polling_timer);
		m_polling_timer = -1;
	}
}

void Rooster::poll( int /* timerID */ )
{
	dprintf(D_FULLDEBUG,"Cock-a-doodle-doo! (Time to look for machines to wake up.)\n");

	CondorQuery unhibernateQuery(STARTD_AD);
	ExprTree *requirements = nullptr;

	if( ParseClassAdRvalExpr( m_unhibernate_constraint.c_str(), requirements )!=0 || requirements==nullptr )
	{
		EXCEPT("Invalid expression for ROOSTER_UNHIBERNATE: %s",
			   m_unhibernate_constraint.c_str());
	}

	unhibernateQuery.addANDConstraint(m_unhibernate_constraint.c_str());

	CollectorList* collects = daemonCore->getCollectorList();
	ASSERT( collects );

	QueryResult result;
	std::vector<std::unique_ptr<ClassAd>> startdAds;
	result = collects->query(unhibernateQuery,startdAds);
	if( result != Q_OK ) {
		dprintf(D_ALWAYS,
				"Couldn't fetch startd ads using constraint "
				"ROOSTER_UNHIBERNATE=%s: %s\n",
				m_unhibernate_constraint.c_str(), getStrQueryResult(result));
		delete requirements;
		return;
	}

	dprintf(D_FULLDEBUG,"Got %zu startd ads matching ROOSTER_UNHIBERNATE=%s\n",
			startdAds.size(), m_unhibernate_constraint.c_str());

	// Assign HibernationRank once to each startd ad, just in case the classad expression
	// is not stable. (e.g. it depends on time or random numbers)

	static const char *ATTR_HIBERNATION_RANK = "_CondorHibernationRank";

	for (auto& startdAd: startdAds) {
		double hibernation_rank = 0.0;
		EvalFloat(ATTR_RANK,&m_rank_ad,startdAd.get(),hibernation_rank);
		startdAd->Assign(ATTR_HIBERNATION_RANK,hibernation_rank);
	}

	// Sort startd ads by HibernationRank (higher rank first)
	std::ranges::sort(startdAds,
		[](const std::unique_ptr<ClassAd>& ad1, const std::unique_ptr<ClassAd>& ad2) {
			double rank1 = 0.0;
			double rank2 = 0.0;
			ad1->LookupFloat(ATTR_HIBERNATION_RANK, rank1);
			ad2->LookupFloat(ATTR_HIBERNATION_RANK, rank2);

			if (fabs(rank1 - rank2) < DBL_EPSILON) {
				return ad1.get() > ad2.get(); // pointer comparison to have a consistent order
			}
			return rank1 > rank2; // higher rank first
		}
	);

	int num_woken = 0;
	std::set<std::string> machines_done;
	for (const auto& startdAdPtr : startdAds) {
		ClassAd *startd_ad = startdAdPtr.get();
		std::string machine;
		std::string name;
		startd_ad->LookupString(ATTR_MACHINE,machine);
		startd_ad->LookupString(ATTR_NAME,name);

		if( machines_done.count(machine)!=0 ) {
			dprintf(D_FULLDEBUG,
					"Skipping %s: already attempted to wake up %s in this cycle.\n",
					name.c_str(),machine.c_str());
			continue;
		}

			// in case the unhibernate expression is time-sensitive,
			// re-evaluate it now to make sure it still passes
		if( !EvalExprBool(startd_ad,requirements) ) {
			dprintf(D_ALWAYS,
					"Skipping %s: ROOSTER_UNHIBERNATE is no longer true.\n",
					name.c_str());
			continue;
		}

		if( wakeUp(startd_ad) ) {
			machines_done.emplace(machine);

			if( ++num_woken >= m_max_unhibernate && m_max_unhibernate > 0 ) {
				dprintf(D_ALWAYS,
						"Reached ROOSTER_MAX_UNHIBERNATE=%d in this cycle.\n",
						m_max_unhibernate);
				break;
			}
		}
	}

	delete requirements;
	requirements = NULL;

	if (!startdAds.empty() ) {
		dprintf(D_FULLDEBUG,"Done sending wakeup calls.\n");
	}
}

bool
Rooster::wakeUp(ClassAd *startd_ad)
{
	ASSERT( startd_ad );

	std::string name;
	startd_ad->LookupString(ATTR_NAME,name);

	dprintf(D_ALWAYS,"Sending wakeup call to %s.\n",name.c_str());

	int stdin_pipe_fds[2];
	stdin_pipe_fds[0] = -1; // child's side
	stdin_pipe_fds[1] = -1; // my side
	if( !daemonCore->Create_Pipe(stdin_pipe_fds) ) {
		dprintf(D_ALWAYS,"Rooster::wakeUp: failed to create stdin pipe.\n");
		return false;
	}

	int stdout_pipe_fds[2];
	stdout_pipe_fds[0] = -1; // my side
	stdout_pipe_fds[1] = -1; // child's side
	if( !daemonCore->Create_Pipe(stdout_pipe_fds) ) {
		dprintf(D_ALWAYS,"Rooster::wakeUp: failed to create stdout pipe.\n");
		daemonCore->Close_Pipe(stdin_pipe_fds[0]);
		daemonCore->Close_Pipe(stdin_pipe_fds[1]);
		return false;
	}

	int std_fds[3];
	std_fds[0] = stdin_pipe_fds[0];
	std_fds[1] = stdout_pipe_fds[1];
	std_fds[2] = stdout_pipe_fds[1];

	OptionalCreateProcessArgs cpArgs;
	int pid = daemonCore->CreateProcessNew(
		m_wakeup_args.GetArg(0),
		m_wakeup_args,
		cpArgs.priv(PRIV_CONDOR_FINAL)
			.reaperID(0)
			.wantCommandPort(FALSE)
			.wantUDPCommandPort(FALSE)
			.std(std_fds));

	daemonCore->Close_Pipe(stdin_pipe_fds[0]);
	daemonCore->Close_Pipe(stdout_pipe_fds[1]);

	if( pid == -1 ) {
		dprintf(D_ALWAYS,"Failed to run %s: %s\n",
				m_wakeup_cmd.c_str(), strerror(errno));
		daemonCore->Close_Pipe(stdin_pipe_fds[1]);
		daemonCore->Close_Pipe(stdout_pipe_fds[0]);
		return false;
	}

	std::string stdin_str;
	sPrintAd(stdin_str, *startd_ad);

		// Beware: the following code assumes that we will not
		// deadlock by filling up the stdin pipe while the tool blocks
		// filling up the stdout pipe.  The tool must consume all
		// input before generating more than a pipe buffer full of output.

	int n = daemonCore->Write_Pipe(stdin_pipe_fds[1],stdin_str.c_str(),stdin_str.length());
	if( n < 0 || (unsigned)n != stdin_str.length() ) {
		dprintf(D_ALWAYS,"Rooster::wakeUp: failed to write to %s: %s\n",
				m_wakeup_cmd.c_str(), strerror(errno));
		daemonCore->Close_Pipe(stdin_pipe_fds[1]);
		daemonCore->Close_Pipe(stdout_pipe_fds[0]);
		return false;
	}

		// done writing to tool
	daemonCore->Close_Pipe(stdin_pipe_fds[1]);

	std::string stdout_str;
	while( true ) {
		char pipe_buf[1024];
		n = daemonCore->Read_Pipe(stdout_pipe_fds[0],pipe_buf,1023);
		if( n <= 0 ) {
			break;
		}
		ASSERT( n < 1024 );
		pipe_buf[n] = '\0';
		stdout_str += pipe_buf;
	}

		// done reading from tool
	daemonCore->Close_Pipe(stdout_pipe_fds[0]);

	if( stdout_str.length() ) {
			// log debugging output from the tool
		dprintf(D_ALWAYS|D_NOHEADER,"%s",stdout_str.c_str());
	}

		// Would be nice to get final exit status of tool, but
		// daemonCore() doesn't provide a waitpid() equivalent.
		// Why didn't I just use my_popen()?  Because it doesn't
		// allow us to write to the tool _and_ log debugging output.
	return true;
}
