#include "Workflow.hpp"

#include <tsl/robin_set.h>
#include <algorithm>
#include <cassert>
#include <queue>
#include <utility>

WorkflowConnection::WorkflowConnection()
	: MultiConnections{}
	, SingleConnection{ WorkflowNode::kInvalidId, -1 }
	, ConnectionDirection{ OneToMany } {
}

bool WorkflowConnection::IsValid() const {
	return SingleConnection.Node != WorkflowNode::kInvalidId;
}

std::span<WorkflowConnection::ConnectionPoint> WorkflowConnection::GetSourcePoints() {
	switch (ConnectionDirection) {
		case ManyToOne: return MultiConnections;
		case OneToMany: return { &SingleConnection, 1 };
	}
}

std::span<const WorkflowConnection::ConnectionPoint> WorkflowConnection::GetSourcePoints() const {
	switch (ConnectionDirection) {
		case ManyToOne: return MultiConnections;
		case OneToMany: return { &SingleConnection, 1 };
	}
}

std::span<WorkflowConnection::ConnectionPoint> WorkflowConnection::GetDestinationPoints() {
	switch (ConnectionDirection) {
		case ManyToOne: return { &SingleConnection, 1 };
		case OneToMany: return MultiConnections;
	}
}

std::span<const WorkflowConnection::ConnectionPoint> WorkflowConnection::GetDestinationPoints() const {
	switch (ConnectionDirection) {
		case ManyToOne: return { &SingleConnection, 1 };
		case OneToMany: return MultiConnections;
	}
}

bool WorkflowNode::InputPin::IsConstantConnection() const {
	return ConnectionToConst && IsConnected();
}

bool WorkflowNode::InputPin::IsConnected() const {
	return Connection != WorkflowConnection::kInvalidId;
}

BaseValue::Kind WorkflowNode::InputPin::GetMatchingType() const {
	return MatchingType;
}

WorkflowConnection::Direction WorkflowNode::InputPin::GetSupportedDirection() const {
	return AllowsMultipleConnections ? WorkflowConnection::ManyToOne : WorkflowConnection::OneToMany;
}

bool WorkflowNode::OutputPin::IsConnected() const {
	return Connection != WorkflowConnection::kInvalidId;
}

BaseValue::Kind WorkflowNode::OutputPin::GetMatchingType() const {
	return MatchingType;
}

WorkflowConnection::Direction WorkflowNode::OutputPin::GetSupportedDirection() const {
	return AllowsMultipleConnections ? WorkflowConnection::OneToMany : WorkflowConnection::ManyToOne;
}

WorkflowNode::WorkflowNode(Type type, Kind kind)
	: mType{ type }
	, mKind{ kind }
	, mDepth{ -1 } {
}

Vec2i WorkflowNode::GetPosition() const {
	return mPosition;
}

void WorkflowNode::SetPosition(const Vec2i& position) {
	mPosition = position;
}

size_t WorkflowNode::GetId() const {
	return mId;
}

WorkflowNode::Type WorkflowNode::GetType() const {
	return mType;
}

WorkflowNode::Kind WorkflowNode::GetKind() const {
	return mKind;
}

int WorkflowNode::GetDepth() const {
	return mDepth;
}

void WorkflowNode::ConnectInput(int nodeId, WorkflowNode& output, int outputNodeId) {
	mWorkflow->Connect(*this, nodeId, output, outputNodeId);
}

void WorkflowNode::DisconnectInput(int nodeId) {
	mWorkflow->DisconnectByDestination(*this, nodeId);
}

bool WorkflowNode::IsInputConnected(int nodeId) const {
	return mInputs[nodeId].IsConnected();
}

void WorkflowNode::ConnectOutput(int nodeId, WorkflowNode& input, int inputNodeId) {
	mWorkflow->Connect(input, inputNodeId, *this, nodeId);
}

void WorkflowNode::DisconnectOutput(int nodeId) {
	mWorkflow->DisconnectBySource(*this, nodeId);
}

bool WorkflowNode::IsOutputConnected(int nodeId) const {
	return mOutputs[nodeId].IsConnected();
}

WorkflowNode::InputPin& WorkflowNode::InsertInputPin(int atIdx) {
	assert(atIdx >= 0 && atIdx < mInputs.size());

	mInputs.push_back(InputPin{});
	for (int i = (int)mInputs.size() - 1, end = atIdx + 1; i >= end; --i) {
		SwapInputPin(i, i + 1);
	}

	return mInputs[atIdx];
}

void WorkflowNode::RemoveInputPin(int pin) {
	DisconnectInput(pin);
	for (int i = 0, end = (int)mInputs.size() - 1; i < end; ++i) {
		SwapInputPin(i, i + 1);
	}
	mInputs.resize(mInputs.size() - 1);
}

void WorkflowNode::SwapInputPin(int a, int b) {
	auto& pinA = mInputs[a];
	auto& pinB = mInputs[b];

	if (mWorkflow) {
		// Delay assignment to ConnectionPoint::Pin so that the pinB loop (if it happens to run and looks in the same Connection) doesn't match the point updated in the pinA loop
		using Pt = WorkflowConnection::ConnectionPoint;
		Pt* pointA = nullptr;
		Pt* pointB = nullptr;

		if (pinA.IsConnected() && !pinA.IsConstantConnection()) {
			auto pts = mWorkflow->GetConnectionById(pinA.Connection)->GetDestinationPoints();
			auto it = std::find(pts.begin(), pts.end(), Pt{ mId, a });
			pointA = it == pts.end() ? nullptr : &*it;
		}
		if (pinB.IsConnected() && !pinB.IsConstantConnection()) {
			auto pts = mWorkflow->GetConnectionById(pinB.Connection)->GetDestinationPoints();
			auto it = std::find(pts.begin(), pts.end(), Pt{ mId, b });
			pointB = it == pts.end() ? nullptr : &*it;
		}

		if (pointA) pointA->Pin = b;
		if (pointB) pointB->Pin = a;
	}

	std::swap(pinA, pinB);
}

WorkflowNode::OutputPin& WorkflowNode::InsertOutputPin(int atIdx) {
	assert(atIdx >= 0 && atIdx < mOutputs.size());

	mOutputs.push_back(OutputPin{});
	for (int i = (int)mOutputs.size() - 1, end = atIdx + 1; i >= end; --i) {
		SwapOutputPin(i, i + 1);
	}

	return mOutputs[atIdx];
}

void WorkflowNode::RemoveOutputPin(int pin) {
	DisconnectOutput(pin);
	for (int i = 0, end = (int)mOutputs.size() - 1; i < end; ++i) {
		SwapInputPin(i, i + 1);
	}
	mOutputs.resize(mOutputs.size() - 1);
}

void WorkflowNode::SwapOutputPin(int a, int b) {
	auto& pinA = mOutputs[a];
	auto& pinB = mOutputs[b];

	if (mWorkflow) {
		using Pt = WorkflowConnection::ConnectionPoint;
		Pt* pointA = nullptr;
		Pt* pointB = nullptr;

		if (pinA.IsConnected()) {
			auto pts = mWorkflow->GetConnectionById(pinA.Connection)->GetSourcePoints();
			auto it = std::find(pts.begin(), pts.end(), Pt{ mId, a });
			pointA = it == pts.end() ? nullptr : &*it;
		}
		if (pinB.IsConnected()) {
			auto pts = mWorkflow->GetConnectionById(pinB.Connection)->GetSourcePoints();
			auto it = std::find(pts.begin(), pts.end(), Pt{ mId, b });
			pointB = it == pts.end() ? nullptr : &*it;
		}

		if (pointA) pointA->Pin = b;
		if (pointB) pointB->Pin = a;
	}

	std::swap(pinA, pinB);
}

WorkflowConnection* Workflow::GetConnectionById(size_t id) {
	return &mConnections[id];
}

WorkflowNode* Workflow::GetStepById(size_t id) {
	return mNodes[id].get();
}

BaseValue* Workflow::GetConstantById(size_t id) {
	return mConstants[id].get();
}

void Workflow::AddStep(std::unique_ptr<WorkflowNode> step) {
	auto [storage, id] = AllocWorkflowStep();
	storage = std::move(step);
	storage->OnAttach(*this, id);
	storage->mWorkflow = this;
	storage->mId = id;
}

void Workflow::RemoveStep(size_t id) {
	auto& step = mNodes[id];
	if (step == nullptr) return;

	step->OnDetach();
	step->mWorkflow = nullptr;
	step->mId = WorkflowNode::kInvalidId;
}

void Workflow::RemoveConnection(size_t id) {
	auto& conn = mConnections[id];
	if (!conn.IsValid()) return;

	for (auto& point : conn.GetSourcePoints()) {
		auto& node = *mNodes[point.Node];
		auto& pin = node.mOutputs[point.Pin];
		pin.Connection = WorkflowNode::kInvalidId;
	}
	for (auto& point : conn.GetDestinationPoints()) {
		auto& node = *mNodes[point.Node];
		auto& pin = node.mInputs[point.Pin];
		pin.Connection = WorkflowNode::kInvalidId;
	}

	conn = {};
	mDepthsDirty = true;
}

bool Workflow::Connect(WorkflowNode& sourceNode, int sourcePin, WorkflowNode& destinationNode, int destinationPin) {
	auto& src = sourceNode.mOutputs[sourcePin];
	auto& dst = destinationNode.mInputs[destinationPin];

	// Equivalent to `if (o.GetSupportedDirection() != i.GetSupportedDirection()) return false;`
	// Cannot connect two pins of different type
	if (src.AllowsMultipleConnections == dst.AllowsMultipleConnections) return false;
	// Would be same as `dst.GetSupportedDirection()` because we validated that they are the same above
	auto direction = src.GetSupportedDirection();

	// TODO report error to user?
	if (src.GetMatchingType() != dst.GetMatchingType()) {
		return false;
	}

	using Pt = WorkflowConnection::ConnectionPoint;
	Pt* srcConnPt;
	Pt* dstConnPt;
	size_t connId;

	switch (direction) {
		case WorkflowConnection::ManyToOne: {
			if (dst.IsConnected()) return false;

			WorkflowConnection* conn;
			if (src.IsConnected()) {
				conn = &mConnections[src.Connection];
				connId = src.Connection;
			} else {
				auto p = AllocWorkflowConnection();
				conn = &p.first;
				connId = p.second;
			}

			srcConnPt = &conn->MultiConnections.emplace_back();
			dstConnPt = &conn->SingleConnection;
		} break;

		case WorkflowConnection::OneToMany: {
			if (src.IsConnected()) return false;

			WorkflowConnection* conn;
			if (dst.IsConnected()) {
				conn = &mConnections[src.Connection];
				connId = src.Connection;
			} else {
				auto p = AllocWorkflowConnection();
				conn = &p.first;
				connId = p.second;
			}

			srcConnPt = &conn->SingleConnection;
			dstConnPt = &conn->MultiConnections.emplace_back();
		} break;
	}

	srcConnPt->Node = sourceNode.GetId();
	srcConnPt->Pin = sourcePin;
	dstConnPt->Node = destinationNode.GetId();
	dstConnPt->Pin = destinationPin;

	src.Connection = connId;
	dst.Connection = connId;

	mDepthsDirty = true;
	return true;
}

// TODO cleanup these two implementation

bool Workflow::DisconnectBySource(WorkflowNode& sourceNode, int sourcePin) {
	auto& sp = sourceNode.mOutputs[sourcePin];
	if (!sp.IsConnected()) return false;

	auto& conn = mConnections[sp.Connection];

	using Pt = WorkflowConnection::ConnectionPoint;
	switch (sp.GetSupportedDirection()) {
		case WorkflowConnection::ManyToOne: {
			// Recessive pin, remove ConnectionPoint associated with this pin only

			auto& vec = conn.MultiConnections;
			vec.erase(std::remove(vec.begin(), vec.end(), Pt{ sourceNode.GetId(), sourcePin }), vec.end());

			sp.Connection = WorkflowNode::kInvalidId;
		} break;

		case WorkflowConnection::OneToMany: {
			// Dominate pin, removes whole connection

			for (auto& pt : conn.MultiConnections) {
				auto& node = *mNodes[pt.Node];
				node.mInputs[pt.Pin].Connection = WorkflowNode::kInvalidId;
			}
			sp.Connection = WorkflowNode::kInvalidId;

			conn = {};
		} break;
	}

	mDepthsDirty = true;
	return true;
}

bool Workflow::DisconnectByDestination(WorkflowNode& destinationNode, int destinationPin) {
	auto& dp = destinationNode.mInputs[destinationPin];
	if (!dp.IsConnected()) return false;
	if (dp.IsConstantConnection()) {
		dp.ConnectionToConst = false;
		dp.Connection = WorkflowNode::kInvalidId;
		return true;
	}

	auto& conn = mConnections[dp.Connection];

	using Pt = WorkflowConnection::ConnectionPoint;
	switch (dp.GetSupportedDirection()) {
		case WorkflowConnection::ManyToOne: {
			// Dominate pin, removes whole connection

			for (auto& pt : conn.MultiConnections) {
				auto& node = *mNodes[pt.Node];
				node.mOutputs[pt.Pin].Connection = WorkflowNode::kInvalidId;
			}
			dp.Connection = WorkflowNode::kInvalidId;

			conn = {};
		} break;

		case WorkflowConnection::OneToMany: {
			// Recessive pin, remove ConnectionPoint associated with this pin only

			auto& vec = conn.MultiConnections;
			vec.erase(std::remove(vec.begin(), vec.end(), Pt{ destinationNode.GetId(), destinationPin }), vec.end());

			dp.Connection = WorkflowNode::kInvalidId;
		} break;
	}

	mDepthsDirty = true;
	return true;
}

const std::vector<std::vector<size_t>>& Workflow::GetDepthGroups() const {
	return mDepthGroups;
}

bool Workflow::DoesDepthNeedsUpdate() const {
	return mDepthsDirty;
}

Workflow::GraphUpdateResult Workflow::UpdateGraph(bool getInfo) {
	if (!mDepthsDirty) {
		return GraphUpdate_NoWorkToDo{};
	}

	// Terminology:
	// - Dependency = nodes its input pins are connected to
	// - Dependents = nodes its output pins are connected to

	struct WorkingNode {
		// The max depth out of all dependency nodes, maintained during the traversal and committed as the actual depth
		// when all dependencies of this node has been resolved. Add 1 to get the depth that will be assigned to the node.
		int MaximumDepth = 0;
		int FulfilledInputCount = 0;
	};

	std::vector<WorkingNode> workingNodes;
	std::queue<size_t> q;

	// Check if all dependencies of this node is satisfied
	auto CheckNodeDependencies = [&](WorkflowNode& node) -> bool {
		for (auto& pin : node.mInputs) {
			if (!pin.IsConnected()) {
				return false;
			}
		}
		return true;
	};

	workingNodes.reserve(mNodes.size());
	{
		std::vector<size_t> unsatisfiedNodes;
		for (size_t i = 0; i < mNodes.size(); ++i) {
			auto& node = mNodes[i];
			workingNodes.push_back(WorkingNode{});

			if (!node) continue;

			if (!CheckNodeDependencies(*node)) {
				if (getInfo) unsatisfiedNodes.push_back(i);
			}

			node->mDepth = -1;

			// Start traversing with the input nodes
			if (node->GetType() == WorkflowNode::InputType) {
				q.push(i);
			}
		}

		if (!unsatisfiedNodes.empty()) {
			return GraphUpdate_UnsatisfiedDependencies{ std::move(unsatisfiedNodes) };
		}
	}

	auto HasCyclicReference = [&](WorkflowNode& node) -> bool {
		// TODO
		return false;
	};

	auto ProcessNode = [&](WorkflowNode& node) -> void {
		for (auto& pin : node.mOutputs) {
			if (!pin.IsConnected()) continue;
			auto& conn = mConnections[pin.Connection];

			for (auto& point : conn.GetDestinationPoints()) {
				auto& wn = workingNodes[point.Node];
				auto& n = *mNodes[point.Node].get();

				if (HasCyclicReference(n)) {
					// TODO
					break;
				}

				wn.FulfilledInputCount++;
				wn.MaximumDepth = std::max(node.mDepth, wn.MaximumDepth);

				// Node's dependency is fulfilled, we can process its dependents next
				// We use >= here because for a many-to-one pin, the dependency is an "or" relation ship, i.e. any of the nodes firing before this will fulfill the requirement
				if (n.mInputs.size() >= wn.FulfilledInputCount) {
					n.mDepth = wn.MaximumDepth + 1;
				}
			}
		}
	};

	int processedNodes = 0;
	while (!q.empty()) {
		auto& wn = workingNodes[q.front()];
		auto& n = *mNodes[q.front()];
		q.pop();
		processedNodes++;

		ProcessNode(n);
	}

	if (processedNodes < mNodes.size()) {
		// There is unreachable nodes
		if (!getInfo) {
			return GraphUpdate_UnreachableNodes{};
		}

		std::vector<size_t> unreachableNodes;
		for (size_t i = 0; i < mNodes.size(); ++i) {
			auto& wn = workingNodes[i];
			auto& n = *mNodes[i];

			// This is a reachable node
			if (n.mDepth != -1) continue;

			unreachableNodes.push_back(i);
		}

		return GraphUpdate_UnreachableNodes{ std::move(unreachableNodes) };
	}

	return GraphUpdate_Success{};
}

std::pair<WorkflowConnection&, size_t> Workflow::AllocWorkflowConnection() {
	for (size_t idx = 0; idx < mConnections.size(); ++idx) {
		auto& elm = mConnections[idx];
		if (!elm.IsValid()) {
			return { elm, idx };
		}
	}

	auto id = mConnections.size();
	return { mConnections.emplace_back(WorkflowConnection{}), id };
}

std::pair<std::unique_ptr<WorkflowNode>&, size_t> Workflow::AllocWorkflowStep() {
	for (size_t idx = 0; idx < mNodes.size(); ++idx) {
		auto& elm = mNodes[idx];
		if (elm == nullptr) {
			return { elm, idx };
		}
	}

	auto id = mNodes.size();
	return { mNodes.emplace_back(std::unique_ptr<WorkflowNode>()), id };
}