//------------------------------------------------------------------------------
// <copyright file="LqrSMACProcessTree.cpp" company="LiquidRobotics">
//
// Copyright (c) 2014 Liquid Robotics, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//
// This software is associated with a release.  Liquid Robotics makes no
// representation that the design, interfaces or component distribution will
// remain unchanged in future releases.
//
// </copyright>
//
// <summary>
// </summary>
//
// <owner></owner>
//------------------------------------------------------------------------------

#include "LqrSMACProcessTree.h"
#include <fstream>
#include <iostream>
#include <string>
#include <boost/algorithm/string.hpp>
#include <boost/filesystem.hpp>
#include <boost/foreach.hpp>
#include <boost/lexical_cast.hpp>

#ifdef NEW_FILESYSTEM
#undef NEW_FILESYSTEM
#endif
#ifdef BOOST_FILESYSTEM_VERSION
#if BOOST_FILESYSTEM_VERSION == 3
#define NEW_FILESYSTEM
#endif
#endif

using namespace std;
using namespace boost;
using namespace boost::filesystem;

namespace LiquidRobotics {
namespace SMACFramework {

ProcessTree::ProcessTree(bool autoLoad, bool includeCmdLines)
{
	if (autoLoad)
		Refresh(includeCmdLines);
}

ProcessTree::~ProcessTree()
{
}

bool ProcessTree::HaveNode(pid_t pid) const
{
	NodeMap::const_iterator it = nodeMap.find(pid);

	return it != nodeMap.end();
}

const ProcessTree::Node& ProcessTree::GetNode(pid_t pid) const
{
	return *nodeMap.find(pid)->second;
}

pid_t ProcessTree::GetParentOfChild(pid_t childPID) const
{
	return GetNode(childPID).parentPID;
}

void ProcessTree::Refresh(bool includeCmdLines)
{
	nodeMap.clear();

	const string procDir = "/proc";

	if (!exists(procDir))
		return;

	directory_iterator endIter;

	for (directory_iterator iter(procDir); iter != endIter; ++iter)
		if (is_directory(iter->status()))
		{
			Node::Ptr node;
			node.reset(new Node);

#ifdef NEW_FILESYSTEM
			const string dirFileName = iter->path().filename().string();
#else
			const string dirFileName = iter->filename();
#endif

			try
			{
				node->pid = lexical_cast<pid_t>(dirFileName);
			}
			catch(bad_lexical_cast& ex)
			{
				continue;
			}

			try
			{
				path newPath(iter->path() / "stat");
				string fileName = newPath.string();
				ifstream statFile(fileName.c_str());
				string ignored;

				statFile >> node->pid >> node->procName >> ignored >> node->parentPID;
			}
			catch(...)
			{
				continue;
			}

			string& procName = node->procName;
			const string::size_type len0 = procName.length();

			if ((len0 >= 2) && (procName[0] == '(') && (procName[len0-1] == ')'))
				procName = procName.substr(1, len0-2);

			if (includeCmdLines)
				try
				{
					path newPath(iter->path() / "cmdline");
					string fileName(newPath.string());
					ifstream cmdLineFile(fileName.c_str());
					stringstream outstm;

					while(cmdLineFile.good())
					{
						char c = cmdLineFile.get();

						if (cmdLineFile.good())
							outstm << c;
					}

					node->commandLine = outstm.str();
				}
				catch(...)
				{
				}

			nodeMap[node->pid] = node;
		}

	for (NodeMap::iterator it = nodeMap.begin(); it != nodeMap.end(); ++it)
	{
		NodeMap::iterator inner = nodeMap.find(it->second->parentPID);

		if (inner != nodeMap.end())
			inner->second->children.insert(it->first);
	}
}

void ProcessTree::GetChildrenOfPID(PIDSet& resultSet, pid_t pid, bool recursive) const
{
	resultSet.clear();
	AddChildrenOfPID(resultSet, pid, recursive);
}

void ProcessTree::AddChildrenOfPID(PIDSet& resultSet, pid_t pid, bool recursive) const
{
	const Node& node = GetNode(pid);

	for (PIDSet::const_iterator it = node.children.begin(); it != node.children.end(); ++it)
		if (resultSet.find(*it) == resultSet.end())
		{
			resultSet.insert(*it);

			if (recursive)
				AddChildrenOfPID(resultSet, *it, true);
		}
}

void ProcessTree::Stream(ostream& out) const
{
	out << "nodeMap contains " << nodeMap.size() << " item(s)";

	if (nodeMap.size())
	{
		unsigned idx = 0;

		out << ":";

		for (NodeMap::const_iterator it = nodeMap.begin(); it != nodeMap.end(); ++it)
			out << endl << (++idx) << ". { key=" << it->first << ", value={ " << *it->second << " }";
	}
}

void ProcessTree::GetPidListBelow(pid_t pid, PIDList& result, bool reversed) const
{
	PIDSet tracking;

	result.clear();
	GetNode(pid);		// throws if no such pid is in the hierarchy

	if (reversed)
		AddPidListBelowReversed(pid, result, tracking);
	else
		AddPidListBelow(pid, result, tracking);
}

void ProcessTree::AddPidListBelow(pid_t pid, PIDList& result, PIDSet& tracking) const
{
	const Node& node = GetNode(pid);

	for (PIDSet::const_iterator it = node.children.begin(); it != node.children.end(); ++it)
		if ((tracking.find(*it) == tracking.end()) && HaveNode(*it))
		{
			tracking.insert(*it);
			result.push_back(*it);
			AddPidListBelow(*it, result, tracking);
		}
}

void ProcessTree::AddPidListBelowReversed(pid_t pid, PIDList& result, PIDSet& tracking) const
{
	const Node& node = GetNode(pid);

	for (PIDSet::const_iterator it = node.children.begin(); it != node.children.end(); ++it)
		if ((tracking.find(*it) == tracking.end()) && HaveNode(*it))
		{
			tracking.insert(*it);
			AddPidListBelowReversed(*it, result, tracking);
			result.push_back(*it);
		}
}

ostream& operator<<(ostream& out, const ProcessTree::Node& node)
{
	out << "pid=" << node.pid << " parentPID=" << node.parentPID << " procName='" << node.procName << "' commandLine='" << node.commandLine << "' children={";

	for (ProcessTree::PIDSet::const_iterator it = node.children.begin(); it != node.children.end(); ++it)
		out << " " << *it;

	out << " }";
	return out;
}

ostream& operator<<(ostream& out, const ProcessTree& tree)
{
	tree.Stream(out);
	return out;
}

}
}
